# 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
SME_THRESHOLD_UV = 1.0      # the stated threshold; see the header for why

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)
m_win = (times >= W[0]) & (times <= W[1])

# Single-trial scores: the mean amplitude of each individual trial in the measurement window.
scores = {s: {c: store[s][c][:, i_ch, :][:, m_win].mean(1) for c in ("target", "standard")} for s in SUBJECTS}
print(f"single-trial mean amplitudes at {CH}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms, "
      f"{len(SUBJECTS)} subjects")
print(f"{'subject':9s} {'nT':>4s} {'nS':>4s} {'trial SD (T)':>13s} {'SME (T)':>9s} "
      f"{'trial SD (S)':>13s} {'SME (S)':>9s} {'score (T-S)':>12s}")
for s in SUBJECTS:
    t_, s_ = scores[s]["target"], scores[s]["standard"]
    print(f"{L3.erpcore_subject_id(s):9s} {len(t_):4d} {len(s_):4d} {t_.std(ddof=1):13.3f} "
          f"{L3.sme(t_):9.3f} {s_.std(ddof=1):13.3f} {L3.sme(s_):9.3f} {t_.mean() - s_.mean():12.3f}")

rng = np.random.default_rng(20260918)
N_BOOT = 400
CHECK = SUBJECTS[0]
ns = np.array([4, 8, 12, 16, 20, 24, 28, 32, 36])
ns = ns[ns <= len(scores[CHECK]["target"])]

emp, formula = [], []
x = scores[CHECK]["target"]
for n in ns:
    draws = np.array([x[rng.choice(len(x), n, replace=False)].mean() for _ in range(N_BOOT)])
    emp.append(draws.std(ddof=1))
    formula.append(x.std(ddof=1) / np.sqrt(n))

fig, ax = plt.subplots(figsize=(8.5, 4.2))
ax.plot(ns, emp, "o-", label=f"resampled: SD of {N_BOOT} re-measurements")
ax.plot(ns, formula, "s--", label="formula: trial SD / sqrt(N)")
ax.set(xlabel="Trials averaged", ylabel="SME (uV)",
       title=f"{L3.erpcore_subject_id(CHECK)}: the SME formula against resampling "
             f"({CH} mean amplitude, {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

print(f"largest relative disagreement between resampling and the formula: "
      f"{100 * np.max(np.abs(np.array(emp) - np.array(formula)) / np.array(formula)):.1f} %")
print("  (resampling without replacement from a finite pool is slightly optimistic at large N -- the finite "
      "population correction; the formula is the estimate you would report)")

n_grid = np.arange(4, 201)
fig, ax = plt.subplots(figsize=(9, 4.6))
min_n = {}
for s in SUBJECTS:
    sd = scores[s]["target"].std(ddof=1)
    curve = sd / np.sqrt(n_grid)
    need = int(np.ceil((sd / SME_THRESHOLD_UV) ** 2))
    min_n[s] = need
    have = len(scores[s]["target"])
    ax.plot(n_grid, curve, lw=1.0, alpha=0.8,
            label=f"{L3.erpcore_subject_id(s)} (SD {sd:.1f} uV, needs {need}, has {have})")
ax.axhline(SME_THRESHOLD_UV, color="k", ls="--", lw=1.4, label=f"threshold {SME_THRESHOLD_UV:g} uV")
ax.set(xlabel="Target trials averaged", ylabel="SME of the mean amplitude (uV)", xscale="log",
       title=f"SME against trial count per subject ({CH}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms mean "
             f"amplitude, target condition, uV)")
ax.set_ylim(0, 4)
ax.grid(alpha=0.3, which="both")
ax.legend(fontsize=7, ncol=2)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

have = {s: len(scores[s]["target"]) for s in SUBJECTS}
reach = [s for s in SUBJECTS if min_n[s] <= have[s]]
print(f"threshold SME <= {SME_THRESHOLD_UV:g} uV on the target-condition mean amplitude at {CH}:")
print(f"{'subject':9s} {'trial SD':>9s} {'trials held':>12s} {'SME held':>9s} {'trials needed':>14s} {'verdict':>10s}")
for s in SUBJECTS:
    sd = scores[s]["target"].std(ddof=1)
    print(f"{L3.erpcore_subject_id(s):9s} {sd:9.3f} {have[s]:12d} {L3.sme(scores[s]['target']):9.3f} "
          f"{min_n[s]:14d} {'reaches' if min_n[s] <= have[s] else 'short':>10s}")
needs = np.array([min_n[s] for s in SUBJECTS])
print(f"\nminimum trials for SME <= {SME_THRESHOLD_UV:g} uV: median {int(np.median(needs))}, "
      f"range {needs.min()}-{needs.max()}; "
      f"{len(reach)}/{len(SUBJECTS)} subjects reach it with the {have[SUBJECTS[0]]}-ish trials the paradigm gives")
print(f"to cover 80 % of these subjects you would need {int(np.percentile(needs, 80))} target trials; "
      f"to cover all of them, {needs.max()}")

effect = np.array([scores[s]["target"].mean() - scores[s]["standard"].mean() for s in SUBJECTS])
sme_t = np.array([L3.sme(scores[s]["target"]) for s in SUBJECTS])
sme_s = np.array([L3.sme(scores[s]["standard"]) for s in SUBJECTS])
sme_d = np.sqrt(sme_t ** 2 + sme_s ** 2)     # the two averages are independent, so the errors add in quadrature

print(f"per subject: P3 (target - standard) against the SME of that difference")
print(f"{'subject':9s} {'P3 (uV)':>9s} {'SME(T)':>8s} {'SME(S)':>8s} {'SME(diff)':>10s} {'P3 / SME':>9s}")
for k, s in enumerate(SUBJECTS):
    print(f"{L3.erpcore_subject_id(s):9s} {effect[k]:9.3f} {sme_t[k]:8.3f} {sme_s[k]:8.3f} "
          f"{sme_d[k]:10.3f} {effect[k] / sme_d[k]:9.2f}")
print(f"\ngroup: P3 {effect.mean():+.3f} uV, between-subject SD {effect.std(ddof=1):.3f} uV, "
      f"median within-subject SME of the difference {np.median(sme_d):.3f} uV")
print(f"  between-subject variance {effect.var(ddof=1):.3f} uV^2 = true between-subject variance "
      f"{max(effect.var(ddof=1) - np.mean(sme_d ** 2), 0):.3f} + mean measurement variance "
      f"{np.mean(sme_d ** 2):.3f}")
print(f"  so about {100 * np.mean(sme_d ** 2) / effect.var(ddof=1):.0f} % of the apparent spread between "
      f"these subjects is measurement error, not difference between people")

fig, ax = plt.subplots(figsize=(8.5, 4.2))
order = np.argsort(effect)
ax.errorbar(np.arange(len(SUBJECTS)), effect[order], yerr=sme_d[order], fmt="o", capsize=4)
ax.axhline(0, color="gray", lw=0.8)
ax.axhline(effect.mean(), color="tab:orange", lw=1.2, ls="--",
           label=f"group mean {effect.mean():+.2f} uV")
ax.set_xticks(np.arange(len(SUBJECTS)))
ax.set_xticklabels([L3.erpcore_subject_id(SUBJECTS[i])[-3:] for i in order], fontsize=8)
ax.set(xlabel="Subject (sorted by effect)", ylabel="P3, target - standard (uV)",
       title=f"Each subject's P3 with its standardized measurement error ({CH}, "
             f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms, uV)")
ax.grid(alpha=0.3, axis="y")
ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"measurement error of four scores on the same target trials "
      f"({CH}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms, 500 bootstrap resamples, seed 20260918)")
print(f"{'subject':9s} {'nT':>4s} {'mean amp':>10s} {'SME analytic':>13s} {'SME boot':>9s} "
      f"{'peak amp':>9s} {'SME boot':>9s} {'peak lat (ms)':>14s} {'SME boot (ms)':>14s} "
      f"{'50% FAL (ms)':>13s} {'SME boot (ms)':>14s}")
boot = {}
for s in SUBJECTS:
    x = store[s]["target"][:, i_ch, :]
    avg = x.mean(0)
    b_mean = L3.bootstrap_sme(x, times, W, measure="mean amplitude")
    b_peak = L3.bootstrap_sme(x, times, W, measure="peak amplitude")
    b_plat = L3.bootstrap_sme(x, times, W, measure="peak latency")
    b_fal = L3.bootstrap_sme(x, times, W, measure="50% fractional-area latency")
    boot[s] = (b_mean, b_peak, b_plat, b_fal)
    print(f"{L3.erpcore_subject_id(s):9s} {x.shape[0]:4d} "
          f"{L3.mean_amplitude(avg, None, W, times=times):10.3f} {L3.sme(scores[s]['target']):13.3f} "
          f"{b_mean:9.3f} {L3.peak_amplitude(avg, None, W, times=times):9.3f} {b_peak:9.3f} "
          f"{1000 * L3.peak_latency(avg, None, W, times=times):14.1f} {1000 * b_plat:14.1f} "
          f"{1000 * L3.fractional_area_latency(avg, None, W, times=times):13.1f} {1000 * b_fal:14.1f}")
b = np.array([boot[s] for s in SUBJECTS])
print(f"\nmedian measurement error: mean amplitude {np.median(b[:, 0]):.3f} uV (analytic "
      f"{np.median([L3.sme(scores[s]['target']) for s in SUBJECTS]):.3f} uV), "
      f"peak amplitude {np.median(b[:, 1]):.3f} uV, peak latency {1000 * np.median(b[:, 2]):.1f} ms, "
      f"50% fractional-area latency {1000 * np.median(b[:, 3]):.1f} ms")
print(f"  a peak amplitude costs {np.median(b[:, 1]) / np.median(b[:, 0]):.1f}x the measurement error of a "
      f"mean amplitude on the same trials, and a peak latency costs "
      f"{np.median(b[:, 2]) / np.median(b[:, 3]):.1f}x the error of a fractional-area latency")

print("unequal trial counts, by design and after rejection:")
rej_t, rej_s = [], []
for s in SUBJECTS:
    nfo = store[s]["info"]
    md_ = None
    n_t, n_s = nfo["n_kept"]["target"], nfo["n_kept"]["standard"]
    rej_t.append(n_t)
    rej_s.append(n_s)
    print(f"  {nfo['subject']}: {n_t:3d} target / {n_s:3d} standard = 1:{n_s / max(n_t, 1):.2f}; "
          f"SME(target) {L3.sme(scores[s]['target']):.3f} uV vs SME(standard) "
          f"{L3.sme(scores[s]['standard']):.3f} uV "
          f"(ratio {L3.sme(scores[s]['target']) / L3.sme(scores[s]['standard']):.2f})")
print(f"\nthe target average carries "
      f"{np.mean([L3.sme(scores[s]['target']) / L3.sme(scores[s]['standard']) for s in SUBJECTS]):.2f}x the "
      f"measurement error of the standard average on average, which is close to the "
      f"sqrt({np.mean(rej_s) / np.mean(rej_t):.2f}) = {np.sqrt(np.mean(rej_s) / np.mean(rej_t)):.2f} the trial "
      f"counts alone predict")
print()
print("condition-biased rejection -- did the criterion take the same share from each condition?")
for s in SUBJECTS:
    nfo = store[s]["info"]
    n_t, n_s = nfo["n_kept"]["target"], nfo["n_kept"]["standard"]
    # the paradigm presents 40 target and 160 standard stimuli per subject
    pct_t = 100 * (1 - n_t / 40)
    pct_s = 100 * (1 - n_s / 160)
    flag = "  <- imbalance > 10 points" if abs(pct_t - pct_s) > 10 else ""
    print(f"  {nfo['subject']}: rejected {pct_t:5.1f} % of target, {pct_s:5.1f} % of standard "
          f"(difference {pct_t - pct_s:+5.1f} points){flag}")

print("nb-3-5-sme -- L3.5 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"Score: mean amplitude of a single trial at {CH} over {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms; "
      f"SME = SD of those single-trial scores / sqrt(N), target condition")
print()
print(f"THRESHOLD CHOSEN AND STATED: SME <= {SME_THRESHOLD_UV:g} uV.  Reason: it is about a quarter of the "
      f"grand-average P3 (nb-3-2 measures {effect.mean():+.2f} uV here), so a measurement at that precision "
      f"separates a typical P3 from zero without pretending to separate two typical P3s from each other. "
      f"It is a course decision, not a value from the literature (TODO(confirm)).")
print()
print(f"ANSWER KEY -- ex-3-5 (numeric), minimum trials for SME <= {SME_THRESHOLD_UV:g} uV:")
print(f"    median across the {len(SUBJECTS)} subjects: {int(np.median(needs))} target trials")
print(f"    range: {needs.min()} (best subject, {L3.erpcore_subject_id(min(SUBJECTS, key=lambda s: min_n[s]))}) "
      f"to {needs.max()} (worst subject, {L3.erpcore_subject_id(max(SUBJECTS, key=lambda s: min_n[s]))})")
print(f"    80th percentile: {int(np.percentile(needs, 80))} trials; "
      f"{len(reach)} of {len(SUBJECTS)} subjects reach the threshold with the ~40 target trials the "
      f"paradigm actually gives them")
print(f"    per subject: " + ", ".join(f"{L3.erpcore_subject_id(s)[-3:]}:{min_n[s]}" for s in SUBJECTS))
print()
print(f"ANSWER KEY -- what the SME is worth: the measured SME of the target average is "
      f"{np.median(sme_t):.3f} uV (median), of the standard average {np.median(sme_s):.3f} uV, and of the "
      f"difference {np.median(sme_d):.3f} uV; about "
      f"{100 * np.mean(sme_d ** 2) / effect.var(ddof=1):.0f} % of the between-subject variance in the P3 is "
      f"measurement error rather than real between-subject difference.")
print(f"ANSWER KEY -- scores without a closed form (bootstrapped, 500 resamples, seed 20260918): median "
      f"measurement error is {np.median(b[:, 0]):.3f} uV for a mean amplitude, {np.median(b[:, 1]):.3f} uV "
      f"for a peak amplitude, {1000 * np.median(b[:, 2]):.1f} ms for a peak latency and "
      f"{1000 * np.median(b[:, 3]):.1f} ms for a 50 % fractional-area latency.")
print(f"Pitfalls: pf-peak-amplitude-noise-bias, pf-eye-movements-lateralized.  "
      f"Widget: w-erp-averager (mode sme).")