# 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, rows = {}, []
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}
    rows.append((nfo["subject"], nfo["n_kept"]["target"], nfo["n_kept"]["standard"], nfo["n_rejected"],
                 nfo["bad_channels"], nfo["ica_excluded"],
                 nfo["mean_rt_ms"]["target"], nfo["mean_rt_ms"]["standard"], nfo["accuracy"]))

times = ep.times
eeg = store[SUBJECTS[0]]["info"]["eeg_channels"]
i_ch = eeg.index(CH)

print(f"{'subject':9s} {'target':>7s} {'standard':>9s} {'rejected':>9s}  {'bads':12s} {'ICA':8s} "
      f"{'RT tgt':>7s} {'RT std':>7s} {'acc':>6s}")
for r in rows:
    print(f"{r[0]:9s} {r[1]:7d} {r[2]:9d} {r[3]:9d}  {str(r[4]):12s} {str(r[5]):8s} "
          f"{r[6]:7.0f} {r[7]:7.0f} {100 * r[8]:5.1f}%")
n_t = sum(r[1] for r in rows)
n_s = sum(r[2] for r in rows)
print(f"{'TOTAL':9s} {n_t:7d} {n_s:9d} {sum(r[3] for r in rows):9d}   "
      f"({n_t / len(SUBJECTS):.1f} and {n_s / len(SUBJECTS):.1f} per subject; ratio 1:{n_s / n_t:.1f})")

per_subject = {c: np.stack([store[s][c].mean(0) for s in SUBJECTS]) for c in ("target", "standard")}
grand = {c: per_subject[c].mean(0) for c in ("target", "standard")}
grand["difference"] = grand["target"] - grand["standard"]

fig, axes = plt.subplots(1, 2, figsize=(13, 4.4))
L3.plot_erp({f"target (n = {n_t} trials, {len(SUBJECTS)} subjects)": grand["target"][i_ch],
             f"standard (n = {n_s} trials, {len(SUBJECTS)} subjects)": grand["standard"][i_ch],
             "difference (target - standard)": (grand["difference"][i_ch], {"color": "k", "lw": 1.8})},
            times, window=W, ax=axes[0],
            title=f"Grand-average P3 at {CH}, {len(SUBJECTS)} ERP CORE subjects (uV, positive up)")
for s in SUBJECTS:
    axes[1].plot(times * 1000, (store[s]["target"].mean(0) - store[s]["standard"].mean(0))[i_ch],
                 lw=0.8, alpha=0.65)
axes[1].plot(times * 1000, grand["difference"][i_ch], color="k", lw=2.2, label="grand average")
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].axvline(0, color="gray", lw=0.6)
axes[1].set(xlabel="Time from stimulus (ms)", ylabel="Amplitude (uV)",
            title=f"Every subject's difference wave at {CH} (uV) -- the spread is the story")
axes[1].grid(alpha=0.3)
axes[1].legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"Grand average over {len(SUBJECTS)} subjects at {CH}, window {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms:")
for name in ("target", "standard", "difference"):
    g = grand[name][i_ch]
    print(f"  {name:11s} mean amplitude {L3.mean_amplitude(g, None, W, times=times):+7.3f} uV | "
          f"peak amplitude {L3.peak_amplitude(g, None, W, times=times):+7.3f} uV at "
          f"{1000 * L3.peak_latency(g, None, W, times=times):6.1f} ms | "
          f"50% area latency {1000 * L3.fractional_area_latency(g, None, W, times=times):6.1f} ms")
print()
win_mean = np.array([L3.mean_amplitude(grand["difference"][i], None, W, times=times) for i in range(len(eeg))])
order = np.argsort(-win_mean)
print("largest window-mean difference by channel (uV): "
      + ", ".join(f"{eeg[i]} {win_mean[i]:+.2f}" for i in order[:5]))
print("most negative: " + ", ".join(f"{eeg[i]} {win_mean[i]:+.2f}" for i in order[-3:]))
print(f"  -- the negative sites are the other half of an average reference: the 30 channel values at any "
      f"time point sum to zero by construction, so a centro-parietal positivity forces a negativity elsewhere "
      f"(L2.3, nb-3-4).")

rng = np.random.default_rng(20260918)
N_ORDERS = 40
m_win = (times >= W[0]) & (times <= W[1])
m_pre = (times >= -0.2) & (times < 0.0)
counts = np.arange(2, min(store[s]["target"].shape[0] for s in SUBJECTS) + 1)

snr = np.zeros((len(SUBJECTS), len(counts)))
for k, s in enumerate(SUBJECTS):
    x = store[s]["target"][:, i_ch, :]
    for _ in range(N_ORDERS):
        idx = rng.permutation(x.shape[0])
        run = np.cumsum(x[idx], axis=0) / np.arange(1, x.shape[0] + 1)[:, None]
        sel = run[counts - 1]
        snr[k] += np.abs(sel[:, m_win].mean(1)) / sel[:, m_pre].std(axis=1, ddof=1)
    snr[k] /= N_ORDERS

mean_snr = snr.mean(0)
fig, ax = plt.subplots(figsize=(8.5, 4.2))
for k, s in enumerate(SUBJECTS):
    ax.plot(counts, snr[k], lw=0.8, alpha=0.5)
ax.plot(counts, mean_snr, "k", lw=2.2, label=f"mean over {len(SUBJECTS)} subjects")
ax.plot(counts, mean_snr[-1] * np.sqrt(counts / counts[-1]), "k--", lw=1.2,
        label="sqrt(N), scaled to the last point")
ax.set(xlabel="Trials averaged (target condition)",
       ylabel="SNR (|window mean| / pre-stimulus SD of the average)",
       title=f"Averaging: SNR of the {CH} P3 against trial count "
             f"({N_ORDERS} random orders per subject, seed 20260918)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

for n in (4, 8, 16, 32):
    if n <= counts[-1]:
        j = int(np.argmin(np.abs(counts - n)))
        print(f"  N = {n:3d} trials: mean SNR {mean_snr[j]:5.2f}"
              + (f"  (x{mean_snr[j] / mean_snr[0]:.2f} relative to N = {counts[0]}; "
                 f"sqrt law predicts x{np.sqrt(n / counts[0]):.2f})"))

facts = L3.DATASETS_L3["ds-erpcore"]
print("ds-erpcore, from data/directory.yaml:")
for key in ("name", "citation", "device", "sfreq", "n_channels", "n_subjects", "reference", "mains_hz",
            "online_filters", "license", "access"):
    print(f"  {key:16s}: {facts[key]}")
print()
print("paradigm P3, TaskDescription from the subject's own eeg.json:")
print("  " + facts["p3_task"])
print()
print("the seven ERP CORE paradigms (OSF component ids, CONTRACTS.md Phase 2 addendum): "
      + ", ".join(f"{k} {v}" for k, v in L3.ERPCORE_COMPONENTS.items()))
print("  TODO(confirm): which component each paradigm targets is stated in the lesson from the dataset's "
      "own documentation; this notebook downloads and measures only P3.")

import os

os.environ.setdefault("MOABB_DOWNLOAD_PROVIDER", "upstream")   # pin moabb to the dataset's Zenodo record
raw_bi = helpers.load_spine("ds-brain-invaders", 1)
print(raw_bi)
print("channels:", raw_bi.ch_names)
raw_bi.filter(L3.HP_HZ, L3.LP_HZ, picks="eeg", method="fir", fir_design="firwin", phase="zero", verbose=False)
events_bi = mne.find_events(raw_bi, stim_channel="STI 014", shortest_event=1, verbose=False)
ep_bi = mne.Epochs(raw_bi, events_bi, helpers.BI2014A_EVENT_ID, tmin=L3.EPOCH_TMIN, tmax=L3.EPOCH_TMAX,
                   baseline=L3.EPOCH_BASELINE, picks="eeg", preload=True, verbose=False)
d_bi = ep_bi.get_data() * 1e6
ptp_bi = (d_bi.max(2) - d_bi.min(2)).max(1)
keep_bi = ptp_bi <= L3.REJECT_PTP_UV
ep_bi = ep_bi[np.where(keep_bi)[0]]
ev_t, ev_n = ep_bi["Target"].average(), ep_bi["NonTarget"].average()
diff_bi = mne.combine_evoked([ev_t, ev_n], weights=[1, -1])

print(f"\nbi2014a subject 1, {raw_bi.info['sfreq']:.0f} Hz, recording reference kept (right earlobe, catalog)")
print(f"  {len(events_bi)} flashes; the same {L3.REJECT_PTP_UV:g} uV peak-to-peak criterion rejects "
      f"{int((~keep_bi).sum())} of {len(keep_bi)} epochs ({100 * (~keep_bi).mean():.0f} %) -- against "
      f"{100 * sum(r[3] for r in rows) / (len(SUBJECTS) * 200):.0f} % on the ERP CORE gel caps")
print(f"  kept: {ev_t.nave} Target, {ev_n.nave} NonTarget (ratio 1:{ev_n.nave / ev_t.nave:.1f})")
for chn in ("Pz", "P3", "P4", "Cz"):
    print(f"  {chn}: target {L3.mean_amplitude(ev_t, chn, W):+6.2f} | non-target "
          f"{L3.mean_amplitude(ev_n, chn, W):+6.2f} | T - N {L3.mean_amplitude(diff_bi, chn, W):+6.2f} uV, "
          f"peak {L3.peak_amplitude(diff_bi, chn, W):+6.2f} uV at "
          f"{1000 * L3.peak_latency(diff_bi, chn, W):.0f} ms")

fig, axes = plt.subplots(1, 2, figsize=(13, 4.4), sharey=False)
L3.plot_erp({f"target (n = {n_t})": grand["target"][i_ch], f"standard (n = {n_s})": grand["standard"][i_ch],
             "difference": (grand["difference"][i_ch], {"color": "k", "lw": 1.8})}, times, window=W, ax=axes[0],
            title=f"ds-erpcore P3, {len(SUBJECTS)} subjects, {CH} (uV)")
i_bi = diff_bi.ch_names.index("Pz")
L3.plot_erp({f"Target (n = {ev_t.nave})": ev_t.data[i_bi] * 1e6,
             f"NonTarget (n = {ev_n.nave})": ev_n.data[i_bi] * 1e6,
             "difference": (diff_bi.data[i_bi] * 1e6, {"color": "k", "lw": 1.8})},
            diff_bi.times, window=W, ax=axes[1],
            title="ds-brain-invaders bi2014a subject 1, dry electrodes, Pz (uV)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

pre_bi = (diff_bi.times >= -0.2) & (diff_bi.times < 0)
pre_ec = (times >= -0.2) & (times < 0)
noise_bi = float(diff_bi.data[i_bi, pre_bi].std(ddof=1) * 1e6)
noise_ec = float(grand["difference"][i_ch, pre_ec].std(ddof=1))
print(f"pre-stimulus SD of the difference average (a noise estimate after averaging):")
print(f"  ds-erpcore grand average, {CH}: {noise_ec:.3f} uV over {n_t} target trials and {len(SUBJECTS)} subjects")
print(f"  bi2014a subject 1, Pz:        {noise_bi:.3f} uV over {ev_t.nave} target trials, one subject")

g = grand["difference"][i_ch]
print("nb-3-2-erp-core-p3 -- L3.2 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, sub-001 to sub-{SUBJECTS[-1]:03d} ({len(SUBJECTS)} subjects, the documented subset "
      f"helpers_l3.SUBSET_DEFAULT); CC-BY-SA-4.0 per data/directory.yaml (contested at source; spec 10.7 most-restrictive rule). Second dataset: ds-brain-invaders "
      f"bi2014a subject 1 (Zenodo DOI 10.5281/zenodo.3266223; CC BY 4.0).")
print(f"Pipeline: helpers_l3.P3_PIPELINE (printed in section 1).")
print(f"Measure: channel {CH}, a-priori window {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms.")
print()
print(f"  trials per condition (ds-erpcore, after rejection): target {n_t} total, "
      f"{n_t / len(SUBJECTS):.1f} per subject (range {min(r[1] for r in rows)}-{max(r[1] for r in rows)}); "
      f"standard {n_s} total, {n_s / len(SUBJECTS):.1f} per subject "
      f"(range {min(r[2] for r in rows)}-{max(r[2] for r in rows)})")
for name in ("target", "standard", "difference"):
    gg = grand[name][i_ch]
    print(f"  grand-average {name:11s}: mean amplitude {L3.mean_amplitude(gg, None, W, times=times):+.3f} uV, "
          f"peak amplitude {L3.peak_amplitude(gg, None, W, times=times):+.3f} uV at "
          f"{1000 * L3.peak_latency(gg, None, W, times=times):.1f} ms")
sub_amp = np.array([L3.mean_amplitude((store[s]["target"].mean(0) - store[s]["standard"].mean(0))[i_ch],
                                      None, W, times=times) for s in SUBJECTS])
print(f"  per-subject P3 (mean amplitude of the difference): mean {sub_amp.mean():+.3f} uV, "
      f"SD {sub_amp.std(ddof=1):.3f}, range {sub_amp.min():+.3f} to {sub_amp.max():+.3f}")
print()
print(f"ANSWER KEY -- SNR and trial count ({CH} P3, SNR = |window mean| / pre-stimulus SD of the same "
      f"average, {N_ORDERS} random trial orders per subject, seed 20260918, mean over "
      f"{len(SUBJECTS)} subjects):")
print(f"    measured mean SNR by trial count: "
      + ", ".join(f"N={nn} -> {mean_snr[int(np.argmin(np.abs(counts - nn)))]:.2f}" for nn in (2, 4, 8, 16, 32)
                  if nn <= counts[-1]))
print(f"    trials needed to reach a target SNR: "
      + ", ".join((f"SNR {tgt:g} -> {int(counts[int(np.argmax(mean_snr >= tgt))])} trials"
                   if (mean_snr >= tgt).any()
                   else f"SNR {tgt:g} -> more than {counts[-1]} trials (the largest count every subject has)")
                  for tgt in (2, 3, 5, 8)))
print(f"    the sqrt law: going from N={counts[0]} to N={counts[-1]} is a factor "
      f"{np.sqrt(counts[-1] / counts[0]):.2f} in theory and {mean_snr[-1] / mean_snr[0]:.2f} measured")
print(f"ANSWER KEY -- dry versus gel: the same component, measured the same way, gives "
      f"{L3.mean_amplitude(diff_bi, 'Pz', W):+.2f} uV at Pz on bi2014a subject 1 "
      f"({ev_t.nave} target, {ev_n.nave} non-target trials, 1:{ev_n.nave / ev_t.nave:.1f}) against "
      f"{L3.mean_amplitude(g, None, W, times=times):+.2f} uV for the ERP CORE grand average; the "
      f"{L3.REJECT_PTP_UV:g} uV criterion costs {100 * (~keep_bi).mean():.0f} % of the dry-electrode epochs "
      f"and {100 * sum(r[3] for r in rows) / (len(SUBJECTS) * 200):.0f} % of the ERP CORE epochs.")
print(f"ANSWER KEY -- paradigm to component (multiple choice): TODO(confirm) -- the mapping is stated in the "
      f"lesson from the dataset's documentation; this notebook verifies only P3 from the active visual oddball.")
print(f"Widget: w-erp-averager (modes averager, sme).")