# Setup: dependencies, the shared helpers, non-interactive plotting.
import importlib.util
import subprocess
import sys
import warnings
from pathlib import Path

# 1. Dependencies are pinned in notebooks/requirements.txt.  Nothing is installed when the
#    pinned stack is already present (local runs, CI); a fresh Colab or Binder kernel installs
#    it once.  On Colab, run from a clone of the repository so that notebooks/_shared/ is
#    available (repository URL: TODO(confirm), spec section 13 item 3).
_needed = ("mne", "scipy", "matplotlib", "pandas", "pooch")
_missing = [p for p in _needed if importlib.util.find_spec(p) is None]
if _missing:
    _req = next((d / "requirements.txt" for d in (Path.cwd(), *Path.cwd().parents)
                 if (d / "requirements.txt").exists()), None)
    _cmd = [sys.executable, "-m", "pip", "install", "-q"]
    _cmd += ["-r", str(_req)] if _req else ["mne==1.10.2", "pooch>=1.8"]
    subprocess.check_call(_cmd)

# 2. Shared helpers, located relative to the working directory -- notebooks/<level>/ or
#    notebooks/ -- never through an absolute path.
_shared = next((d / "_shared" for d in (Path.cwd(), *Path.cwd().parents)
                if (d / "_shared" / "helpers_l3.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L3/ (or notebooks/) so that _shared/helpers_l3.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l3 as L3

# 3. Plotting: Jupyter's default inline backend renders static PNGs through Agg (no windows,
#    nothing blocks); outside Jupyter the helpers select Agg.  Every MNE figure is requested
#    with show=False and each figure cell ends with plt.show().
import matplotlib.pyplot as plt
import numpy as np
import mne

mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers_l3 imported from notebooks/_shared")
print(f"ERP CORE cache: {L3.erpcore_root().name}/ (resolved relative to the working directory, "
      "or $EEG_COURSE_ERPCORE); only the subjects this notebook names are fetched")

# The one Level-3 pipeline, printed rather than described.  Every Level-3 notebook and the C3
# capstone call the same helpers_l3.load_p3_epochs, so their numbers are comparable.
for key, value in L3.P3_PIPELINE.items():
    print(f"{key:15s} : {value}")
print()
print(f"a-priori measurement window : {L3.P3_WINDOW[0] * 1000:.0f}-{L3.P3_WINDOW[1] * 1000:.0f} ms "
      f"at {L3.P3_CHANNEL}, fixed in helpers_l3.P3_WINDOW")

SUBJECTS = list(L3.SUBSET_DEFAULT)
W = L3.P3_WINDOW
CH = L3.P3_CHANNEL

store = {}
for s in SUBJECTS:
    ep, nfo = L3.load_p3_epochs(s, verbose=False)
    store[s] = {"target": L3.condition_epochs(ep, "target").get_data(picks="eeg") * 1e6,
                "standard": L3.condition_epochs(ep, "standard").get_data(picks="eeg") * 1e6, "info": nfo}
times = ep.times
eeg = store[SUBJECTS[0]]["info"]["eeg_channels"]
i_ch = eeg.index(CH)
diff = {s: store[s]["target"].mean(0)[i_ch] - store[s]["standard"].mean(0)[i_ch] for s in SUBJECTS}
n_target = {s: store[s]["target"].shape[0] for s in SUBJECTS}

print(f"{len(SUBJECTS)} subjects loaded; target trials per subject: "
      + ", ".join(f"{L3.erpcore_subject_id(s)} {n_target[s]}" for s in SUBJECTS))
LOW = min(SUBJECTS, key=lambda s: n_target[s])
print(f"lowest-trial subject: {L3.erpcore_subject_id(LOW)} with {n_target[LOW]} target trials "
      f"(chosen by trial count alone, not by its amplitude)")

d_low = diff[LOW]
fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
ax = axes[0]
L3.plot_erp({f"{L3.erpcore_subject_id(LOW)} difference wave (n = {n_target[LOW]} target trials)": d_low},
            times, window=W, ax=ax, title=f"Five measurements of one difference wave at {CH} (uV)")
m = (times >= W[0]) & (times <= W[1])
pk = L3.peak_amplitude(d_low, None, W, times=times)
pk_t = L3.peak_latency(d_low, None, W, times=times)
mn = L3.mean_amplitude(d_low, None, W, times=times)
fal = L3.fractional_area_latency(d_low, None, W, times=times)
ax.plot(pk_t * 1000, pk, "v", color="tab:red", ms=9, label=f"peak {pk:+.2f} uV at {pk_t * 1000:.0f} ms")
ax.hlines(mn, W[0] * 1000, W[1] * 1000, color="tab:green", lw=2, label=f"mean {mn:+.2f} uV")
ax.axvline(fal * 1000, color="tab:purple", lw=1.5, ls="--", label=f"50% area latency {fal * 1000:.0f} ms")
ax.fill_between(times[m] * 1000, 0, np.clip(d_low[m], 0, None), color="tab:orange", alpha=0.25,
                label=f"positive area {L3.area_amplitude(d_low, None, W, times=times):.3f} uV*s")
ax.legend(fontsize=8)

for s in SUBJECTS:
    axes[1].plot(times * 1000, diff[s], lw=0.9, alpha=0.7)
    p_t = L3.peak_latency(diff[s], None, W, times=times)
    axes[1].plot(p_t * 1000, L3.peak_amplitude(diff[s], None, W, times=times), "v", ms=5, color="tab:red")
axes[1].axvspan(W[0] * 1000, W[1] * 1000, color="tab:orange", alpha=0.18, lw=0)
axes[1].axhline(0, color="gray", lw=0.6)
axes[1].set(xlabel="Time from stimulus (ms)", ylabel="Amplitude (uV)",
            title=f"Each subject's difference wave and its peak ({CH}, uV) -- the peaks scatter in both axes")
axes[1].grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"Per subject, {CH}, target minus standard, window {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms")
header = f"{'subject':9s} {'nT':>4s} " + " ".join(f"{k:>14s}" for k in L3.MEASURES)
print(header)
table = {k: [] for k in L3.MEASURES}
for s in SUBJECTS:
    vals = L3.measure_all(diff[s], None, W, times=times)
    for k, v in vals.items():
        table[k].append(v)
    print(f"{L3.erpcore_subject_id(s):9s} {n_target[s]:4d} "
          + " ".join(f"{vals[k]:14.4f}" for k in L3.MEASURES))
table = {k: np.array(v) for k, v in table.items()}
print()
print(f"{'method':30s} {'unit':7s} {'mean':>10s} {'SD':>10s} {'SEM':>9s} {'CV':>8s}")
for k, (fn, unit) in L3.MEASURES.items():
    v = table[k]
    print(f"{k:30s} {unit:7s} {v.mean():10.4f} {v.std(ddof=1):10.4f} "
          f"{v.std(ddof=1) / np.sqrt(len(v)):9.4f} {abs(v.std(ddof=1) / v.mean()):8.3f}")

rng = np.random.default_rng(20260918)
N_REP = 60
sizes = [5, 8, 12, 16, 24, 32]
sizes = [n for n in sizes if n <= min(n_target.values())] + [min(n_target.values())]
sizes = sorted(set(sizes))

curve = {"peak amplitude": np.zeros((len(SUBJECTS), len(sizes))),
         "mean amplitude": np.zeros((len(SUBJECTS), len(sizes)))}
for k, s in enumerate(SUBJECTS):
    tgt = store[s]["target"][:, i_ch, :]
    std_avg = store[s]["standard"].mean(0)[i_ch]
    for j, n in enumerate(sizes):
        pk, mn_ = [], []
        for _ in range(N_REP):
            sel = rng.choice(tgt.shape[0], n, replace=False)
            d = tgt[sel].mean(0) - std_avg
            pk.append(L3.peak_amplitude(d, None, W, times=times))
            mn_.append(L3.mean_amplitude(d, None, W, times=times))
        curve["peak amplitude"][k, j] = np.mean(pk)
        curve["mean amplitude"][k, j] = np.mean(mn_)

fig, ax = plt.subplots(figsize=(8.5, 4.2))
for name, style in (("peak amplitude", "-o"), ("mean amplitude", "-s")):
    y = curve[name].mean(0)
    e = curve[name].std(0, ddof=1) / np.sqrt(len(SUBJECTS))
    ax.errorbar(sizes, y, yerr=e, fmt=style, capsize=3, label=name)
ax.set(xlabel="Target trials averaged (random subsets, 60 draws per subject)", ylabel="Amplitude (uV)",
       title=f"Peak amplitude is biased by trial count; mean amplitude is not ({CH}, "
             f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms, uV)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

for name in ("peak amplitude", "mean amplitude"):
    y = curve[name].mean(0)
    print(f"  {name:15s}: N = {sizes[0]:2d} -> {y[0]:+.3f} uV, N = {sizes[-1]:2d} -> {y[-1]:+.3f} uV "
          f"(change {y[0] - y[-1]:+.3f} uV, {100 * (y[0] - y[-1]) / abs(y[-1]):+.1f} %)")

def split_half(measure_fn, mode_kwargs=None):
    """(r_half, r_spearman_brown) over subjects for one measurement, odd vs even target trials."""
    mode_kwargs = mode_kwargs or {}
    a, b = [], []
    for s in SUBJECTS:
        tgt = store[s]["target"][:, i_ch, :]
        std_avg = store[s]["standard"].mean(0)[i_ch]
        odd = tgt[1::2].mean(0) - std_avg
        even = tgt[0::2].mean(0) - std_avg
        a.append(measure_fn(odd, None, W, times=times, **mode_kwargs))
        b.append(measure_fn(even, None, W, times=times, **mode_kwargs))
    a, b = np.array(a), np.array(b)
    ok = np.isfinite(a) & np.isfinite(b)
    r = float(np.corrcoef(a[ok], b[ok])[0, 1])
    # Spearman-Brown is only meaningful for a positive half correlation: a negative r means the two
    # halves rank the subjects differently, i.e. the measure carries no reliable between-subject signal,
    # and "correcting" it produces a number below -1 that should not be reported as a reliability.
    rsb = 2 * r / (1 + r) if r > 0 else float("nan")
    return r, rsb


print(f"Split-half reliability over {len(SUBJECTS)} subjects (odd vs even target trials, same standard average)")
print(f"{'method':30s} {'unit':7s} {'across-subject SD':>19s} {'r(half)':>9s} {'r(Spearman-Brown)':>19s}")
reliability = {}
for k, (fn, unit) in L3.MEASURES.items():
    r, rsb = split_half(fn)
    reliability[k] = (r, rsb)
    shown = f"{rsb:19.3f}" if np.isfinite(rsb) else f"{'n/a (r <= 0)':>19s}"
    print(f"{k:30s} {unit:7s} {table[k].std(ddof=1):19.4f} {r:9.3f} {shown}")
print()
print("Read the last column as: how much of the between-subject variance in this measure is signal rather "
      "than noise.  A negative half correlation is not a small reliability, it is no reliability: the two "
      "halves of the same session rank the subjects differently, so the measure is dominated by noise.")

collapsed = np.mean([(store[s]["target"].sum(0) + store[s]["standard"].sum(0))
                     / (store[s]["target"].shape[0] + store[s]["standard"].shape[0])
                     for s in SUBJECTS], axis=0)[i_ch]
grand_diff = np.mean([diff[s] for s in SUBJECTS], axis=0)
post = (times >= 0) & (times <= 0.8)
t_collapsed = times[post][int(collapsed[post].argmax())]
t_diffpeak = times[post][int(grand_diff[post].argmax())]
HALF = 0.150
windows = {
    f"a-priori (helpers_l3.P3_WINDOW) {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms": W,
    f"collapsed localizer: peak of the condition-collapsed grand average "
    f"({t_collapsed * 1000:.0f} ms) +/- {HALF * 1000:.0f} ms": (t_collapsed - HALF, t_collapsed + HALF),
    f"circular: peak of the grand-average DIFFERENCE ({t_diffpeak * 1000:.0f} ms) "
    f"+/- {HALF * 1000:.0f} ms": (t_diffpeak - HALF, t_diffpeak + HALF),
}
from scipy import stats

print(f"Mean amplitude of the difference wave at {CH} under three ways of choosing the window "
      f"({len(SUBJECTS)} subjects):")
for label, win in windows.items():
    v = np.array([L3.mean_amplitude(diff[s], None, win, times=times) for s in SUBJECTS])
    tstat, pval = stats.ttest_1samp(v, 0)
    print(f"  {label}")
    print(f"      mean {v.mean():+.3f} uV, SD {v.std(ddof=1):.3f}, t({len(v) - 1}) = {tstat:.3f}, p = {pval:.5f}, "
          f"dz = {v.mean() / v.std(ddof=1):.3f}")

fig, ax = plt.subplots(figsize=(9, 4.2))
L3.plot_erp({"condition-collapsed grand average (the localizer)": (collapsed, {"color": "tab:blue"}),
             "grand-average difference (target - standard)": (grand_diff, {"color": "k", "lw": 1.8})},
            times, ax=ax, title=f"Collapsed localizer versus the contrast at {CH} (uV)")
colors = ["tab:orange", "tab:green", "tab:red"]
for (label, win), c in zip(windows.items(), colors):
    ax.axvspan(win[0] * 1000, win[1] * 1000, color=c, alpha=0.15, lw=0, label=label.split(":")[0])
ax.legend(fontsize=7)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

d_low_pk = L3.peak_amplitude(diff[LOW], None, W, times=times)
d_low_mn = L3.mean_amplitude(diff[LOW], None, W, times=times)
HIGH = max(SUBJECTS, key=lambda s: n_target[s])
d_high_pk = L3.peak_amplitude(diff[HIGH], None, W, times=times)
d_high_mn = L3.mean_amplitude(diff[HIGH], None, W, times=times)

print("nb-3-3-measurement -- L3.3 exercise numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, sub-001 to sub-{SUBJECTS[-1]:03d} ({len(SUBJECTS)} subjects, "
      f"helpers_l3.SUBSET_DEFAULT); CC-BY-SA-4.0 per data/directory.yaml (contested at source; spec 10.7 most-restrictive rule)")
print(f"Pipeline: helpers_l3.P3_PIPELINE (printed in section 1)")
print(f"Measure: channel {CH}, target minus standard, a-priori window "
      f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms")
print()
print(f"ANSWER KEY -- ex-3-3 (numeric), low-trial subject: {L3.erpcore_subject_id(LOW)} has "
      f"{n_target[LOW]} target trials (the fewest in the subset) and gives")
print(f"    peak amplitude = {d_low_pk:+.2f} uV   mean amplitude = {d_low_mn:+.2f} uV   "
      f"(peak exceeds mean by {d_low_pk - d_low_mn:+.2f} uV)")
print(f"    for comparison {L3.erpcore_subject_id(HIGH)}, {n_target[HIGH]} target trials: "
      f"peak {d_high_pk:+.2f} uV, mean {d_high_mn:+.2f} uV (gap {d_high_pk - d_high_mn:+.2f} uV)")
print(f"    across all {len(SUBJECTS)} subjects the peak exceeds the mean by "
      f"{(table['peak amplitude'] - table['mean amplitude']).mean():+.2f} uV on average "
      f"(range {(table['peak amplitude'] - table['mean amplitude']).min():+.2f} to "
      f"{(table['peak amplitude'] - table['mean amplitude']).max():+.2f})")
print(f"ANSWER KEY -- ex-3-3 (free response): subsampling the target trials of every subject shows the "
      f"mechanism -- averaging {sizes[0]} instead of {sizes[-1]} target trials moves the peak from "
      f"{curve['peak amplitude'].mean(0)[-1]:+.3f} to {curve['peak amplitude'].mean(0)[0]:+.3f} uV "
      f"({100 * (curve['peak amplitude'].mean(0)[0] - curve['peak amplitude'].mean(0)[-1]) / abs(curve['peak amplitude'].mean(0)[-1]):+.1f} %) "
      f"while the mean moves from {curve['mean amplitude'].mean(0)[-1]:+.3f} to "
      f"{curve['mean amplitude'].mean(0)[0]:+.3f} uV "
      f"({100 * (curve['mean amplitude'].mean(0)[0] - curve['mean amplitude'].mean(0)[-1]) / abs(curve['mean amplitude'].mean(0)[-1]):+.1f} %). "
      f"The peak is the maximum of a noisy set and the maximum grows with the noise; the mean is unbiased.")
print()
print("ANSWER KEY -- across-subject standard deviation of each method "
      f"({len(SUBJECTS)} subjects, same window, same channel):")
for k, (fn, unit) in L3.MEASURES.items():
    v = table[k]
    r, rsb = reliability[k]
    extra = f"  ({v.std(ddof=1) * 1000:.1f} ms)" if unit == "s" else ""
    sb = f"Spearman-Brown {rsb:.3f}" if np.isfinite(rsb) else "Spearman-Brown n/a (half correlation <= 0)"
    print(f"    {k:30s} SD = {v.std(ddof=1):.4f} {unit}{extra}; mean {v.mean():.4f} {unit}; "
          f"split-half r = {r:.3f}, {sb}")
print()
print(f"Supporting -- window choice: a-priori {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms gives "
      f"{table['mean amplitude'].mean():+.3f} uV; a collapsed localizer centred on the condition-collapsed "
      f"peak ({t_collapsed * 1000:.0f} ms) gives "
      f"{np.mean([L3.mean_amplitude(diff[s], None, (t_collapsed - HALF, t_collapsed + HALF), times=times) for s in SUBJECTS]):+.3f} uV; "
      f"a window centred on the peak of the difference itself ({t_diffpeak * 1000:.0f} ms) gives "
      f"{np.mean([L3.mean_amplitude(diff[s], None, (t_diffpeak - HALF, t_diffpeak + HALF), times=times) for s in SUBJECTS]):+.3f} uV -- "
      f"the last is circular and is reported only to show the size of the circularity.")
print(f"Pitfalls: pf-peak-amplitude-noise-bias, pf-post-hoc-windows.  Widget: w-measurement-explorer.")