# 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

store = {}
for s in SUBJECTS:
    ep, nfo = L3.load_p3_epochs(s, verbose=False)
    store[s] = {"target": L3.condition_epochs(ep, "target").average(),
                "standard": L3.condition_epochs(ep, "standard").average(), "info": nfo}
times = ep.times
proto = store[SUBJECTS[0]]["target"].copy().pick("eeg")
eeg = proto.ch_names
print(f"{len(SUBJECTS)} subjects; {len(eeg)} EEG channels in the dataset's own order:")
print("  " + ", ".join(eeg))

grand = {}
for cond in ("target", "standard"):
    grand[cond] = np.mean([store[s][cond].copy().pick("eeg").data for s in SUBJECTS], axis=0)
grand["difference"] = grand["target"] - grand["standard"]

i9, i10, icz = eeg.index("P9"), eeg.index("P10"), eeg.index("Cz")
REFS = {
    "average of the 30 EEG channels": lambda X: X,
    "linked mastoids (mean of P9 and P10; stand-in, no mastoid electrode exists here)":
        lambda X: X - 0.5 * (X[[i9]] + X[[i10]]),
    "single electrode Cz": lambda X: X - X[[icz]],
}
print()
print("check that re-referencing conserves the differences between channels (it must, it is a subtraction):")
X = grand["difference"]
for name, fn in REFS.items():
    Y = fn(X)
    d_pz_cpz = np.abs((X[eeg.index("Pz")] - X[eeg.index("CPz")]) - (Y[eeg.index("Pz")] - Y[eeg.index("CPz")])).max()
    print(f"  {name[:46]:46s}: max |(Pz - CPz) change| = {d_pz_cpz:.3e} uV")

TIMES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
TWO = ["average of the 30 EEG channels",
       "linked mastoids (mean of P9 and P10; stand-in, no mastoid electrode exists here)"]

fig, axes = plt.subplots(len(TWO), len(TIMES) + 1, figsize=(2.05 * len(TIMES) + 1.5, 2.55 * len(TWO)),
                         gridspec_kw={"width_ratios": [1] * len(TIMES) + [0.09]})
for row, name in enumerate(TWO):
    Y = REFS[name](grand["difference"]) * 1e6
    v = float(np.abs(Y[:, (times >= 0) & (times <= 0.8)]).max())
    im = None
    for col, t_ in enumerate(TIMES):
        i = int(np.argmin(np.abs(times - t_)))
        im, _ = mne.viz.plot_topomap(Y[:, i], proto.info, axes=axes[row, col], show=False,
                                     contours=4, vlim=(-v, v), sensors=True)
        if row == 0:
            axes[row, col].set_title(f"{times[i] * 1000:.0f} ms", fontsize=9)
    cb = fig.colorbar(im, cax=axes[row, -1])
    cb.set_label("uV")
    axes[row, 0].set_ylabel(name.split("(")[0].strip(), fontsize=8)
    axes[row, -1].set_title(f"+/-{v:.1f}", fontsize=8)
fig.suptitle("Grand-average P3 (target - standard) every 100 ms, two references -- "
             "one colour scale per row, in uV", y=1.0, fontsize=11)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"Peak channel and peak time of the grand-average difference wave, 0-800 ms "
      f"({len(SUBJECTS)} subjects, uV)")
post = (times >= 0) & (times <= 0.8)
t_post = times[post]
peaks = {}
for name, fn in REFS.items():
    Y = fn(grand["difference"])[:, post] * 1e6
    i_pos, j_pos = np.unravel_index(int(Y.argmax()), Y.shape)
    i_abs, j_abs = np.unravel_index(int(np.abs(Y).argmax()), Y.shape)
    win = (times >= W[0]) & (times <= W[1])
    Yw = fn(grand["difference"])[:, win].mean(1) * 1e6
    peaks[name] = {"pos_ch": eeg[i_pos], "pos_t": float(t_post[j_pos]), "pos_v": float(Y[i_pos, j_pos]),
                   "abs_ch": eeg[i_abs], "abs_t": float(t_post[j_abs]), "abs_v": float(Y[i_abs, j_abs]),
                   "win_ch": eeg[int(Yw.argmax())], "win_v": float(Yw.max()),
                   "win_min_ch": eeg[int(Yw.argmin())], "win_min_v": float(Yw.min()),
                   "map": Yw}
    p = peaks[name]
    print(f"  {name}")
    print(f"      most positive sample : {p['pos_ch']:>4s} at {p['pos_t'] * 1000:6.1f} ms = {p['pos_v']:+7.3f} uV")
    print(f"      largest |amplitude|  : {p['abs_ch']:>4s} at {p['abs_t'] * 1000:6.1f} ms = {p['abs_v']:+7.3f} uV")
    print(f"      window mean {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms: peak channel {p['win_ch']:>4s} "
          f"{p['win_v']:+7.3f} uV, most negative {p['win_min_ch']:>4s} {p['win_min_v']:+7.3f} uV")

def centre(v):
    return np.asarray(v, float) - np.mean(v)


maps = {name: peaks[name]["map"] for name in REFS}
names = list(maps)
print(f"Window-mean maps ({W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms) under three references, across 30 channels:")
print(f"{'':52s} {'RMS (uV)':>10s} {'r with avg ref':>15s} {'max |diff| mean-centred':>25s}")
for name in names:
    r = float(np.corrcoef(maps[name], maps[names[0]])[0, 1])
    d = float(np.abs(centre(maps[name]) - centre(maps[names[0]])).max())
    print(f"  {name[:50]:50s} {np.sqrt(np.mean(maps[name] ** 2)):10.3f} {r:15.3f} {d:25.2e}")
print()
print("A reference change cannot change the shape of a map, only where zero sits on it -- which is enough to "
      "move the most extreme value from one side of the head to the other, as section 3 printed.")

fig, axes = plt.subplots(2, len(names) + 1, figsize=(3.1 * len(names) + 1.4, 5.6),
                         gridspec_kw={"width_ratios": [1] * len(names) + [0.09]})
v_raw = max(np.abs(m).max() for m in maps.values())
for col, name in enumerate(names):
    im0, _ = mne.viz.plot_topomap(maps[name], proto.info, axes=axes[0, col], show=False, contours=4,
                                  vlim=(-v_raw, v_raw), sensors=True)
    axes[0, col].set_title(name.split("(")[0].strip()[:28], fontsize=8)
    im1, _ = mne.viz.plot_topomap(centre(maps[name]), proto.info, axes=axes[1, col], show=False,
                                  contours=4, vlim=(-v_raw, v_raw), sensors=True)
cb0 = fig.colorbar(im0, cax=axes[0, -1]); cb0.set_label("uV")
cb1 = fig.colorbar(im1, cax=axes[1, -1]); cb1.set_label("uV")
axes[0, 0].set_ylabel("as measured (uV)", fontsize=9)
axes[1, 0].set_ylabel("mean-centred across channels", fontsize=9)
fig.suptitle(f"P3 window mean {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms under three references: "
             "the bottom row is the same map three times", y=1.0, fontsize=11)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

def window_map(X, a, b):
    m = (times >= a) & (times <= b)
    return X[:, m].mean(1) * 1e6


CANDIDATES = {
    "A  difference 400-500 ms": window_map(grand["difference"], 0.40, 0.50),
    "B  difference 500-600 ms": window_map(grand["difference"], 0.50, 0.60),
    "C  difference 150-250 ms": window_map(grand["difference"], 0.15, 0.25),
    "D  target 100-200 ms": window_map(grand["target"], 0.10, 0.20),
    "E  standard 100-200 ms": window_map(grand["standard"], 0.10, 0.20),
}
keys = list(CANDIDATES)
print(f"Across-channel correlation (shape) and RMS (amplitude), average reference, {len(SUBJECTS)} subjects")
print(f"{'':26s}" + "".join(f"{k.split()[0]:>8s}" for k in keys) + f"{'RMS (uV)':>11s}")
corr = np.zeros((len(keys), len(keys)))
for i, a in enumerate(keys):
    for j, b in enumerate(keys):
        corr[i, j] = float(np.corrcoef(CANDIDATES[a], CANDIDATES[b])[0, 1])
    print(f"{a:26s}" + "".join(f"{corr[i, j]:8.3f}" for j in range(len(keys)))
          + f"{np.sqrt(np.mean(CANDIDATES[a] ** 2)):11.3f}")
print()
print(f"  A vs B: r = {corr[0, 1]:.3f}, RMS {np.sqrt(np.mean(CANDIDATES[keys[0]] ** 2)):.2f} against "
      f"{np.sqrt(np.mean(CANDIDATES[keys[1]] ** 2)):.2f} uV -- the same field, decaying: AMPLITUDE")
print(f"  A vs C: r = {corr[0, 2]:.3f} -- a different field: SHAPE")
print(f"  D vs E: r = {corr[3, 4]:.3f}, RMS {np.sqrt(np.mean(CANDIDATES[keys[3]] ** 2)):.2f} against "
      f"{np.sqrt(np.mean(CANDIDATES[keys[4]] ** 2)):.2f} uV -- the early sensory response the two conditions "
      f"share: neither differs")

fig, axes = plt.subplots(2, len(keys), figsize=(2.3 * len(keys), 5.2))
v_all = max(np.abs(m).max() for m in CANDIDATES.values())
for col, k in enumerate(keys):
    im0, _ = mne.viz.plot_topomap(CANDIDATES[k], proto.info, axes=axes[0, col], show=False, contours=4,
                                  vlim=(-v_all, v_all), sensors=True)
    axes[0, col].set_title(k, fontsize=8)
    im1, _ = mne.viz.plot_topomap(CANDIDATES[k] / np.sqrt(np.mean(CANDIDATES[k] ** 2)), proto.info,
                                  axes=axes[1, col], show=False, contours=4, vlim=(-2.5, 2.5), sensors=True)
axes[0, 0].set_ylabel("as measured (uV)", fontsize=9)
axes[1, 0].set_ylabel("normalised: shape only", fontsize=9)
fig.colorbar(im0, ax=axes[0, :], shrink=0.75, label="uV")
fig.colorbar(im1, ax=axes[1, :], shrink=0.75, label="normalised (RMS = 1)")
fig.suptitle("Five maps: which pairs differ only in amplitude, and which differ in shape?", y=1.0, fontsize=11)
plt.show()   # render the static figure(s) of this cell inline

m_avg = maps[names[0]]
fig, axes = plt.subplots(1, 3, figsize=(11, 3.4))
v = float(np.abs(m_avg).max())
for ax, (extrap, title) in zip(axes, [("head", "extrapolate to the head outline (MNE default for EEG)"),
                                      ("local", "extrapolate only near the sensors ('local')"),
                                      ("box", "extrapolate to a bounding box ('box')")]):
    im, _ = mne.viz.plot_topomap(m_avg, proto.info, axes=ax, show=False, contours=4, vlim=(-v, v),
                                 sensors=True, extrapolate=extrap)
    ax.set_title(title, fontsize=8)
cb = fig.colorbar(im, ax=axes, shrink=0.8)
cb.set_label("uV")
fig.suptitle(f"The same 30 numbers, three extrapolation rules "
             f"(P3 window mean, average reference, uV)", y=1.03, fontsize=10)
plt.show()   # render the static figure(s) of this cell inline

hull = [c for c in eeg]
print(f"the map is drawn from {len(hull)} electrode values; everything between them is interpolated and "
      f"everything outside their hull is extrapolated")
print(f"  the most extreme *measured* values in this window are "
      f"{eeg[int(m_avg.argmax())]} {m_avg.max():+.3f} uV and "
      f"{eeg[int(m_avg.argmin())]} {m_avg.min():+.3f} uV; any colour more extreme than that on the map "
      f"is the interpolator, not the data")

print("nb-3-4-topomaps -- L3.4 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); grand average of per-subject averages, "
      f"unweighted")
print(f"Maps: grand-average difference wave (target - standard), spherical-spline interpolation (MNE "
      f"plot_topomap), 30 channels, standard_1005 positions")
print()
for name in TWO:
    p = peaks[name]
    print(f"ANSWER KEY -- peak channel and time, reference = {name}:")
    print(f"    most positive sample over 0-800 ms : {p['pos_ch']} at {p['pos_t'] * 1000:.1f} ms, "
          f"{p['pos_v']:+.3f} uV")
    print(f"    largest |amplitude| over 0-800 ms  : {p['abs_ch']} at {p['abs_t'] * 1000:.1f} ms, "
          f"{p['abs_v']:+.3f} uV")
    print(f"    peak channel of the {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms window mean: {p['win_ch']} "
          f"({p['win_v']:+.3f} uV); most negative {p['win_min_ch']} ({p['win_min_v']:+.3f} uV)")
r_lm = float(np.corrcoef(maps[TWO[0]], maps[TWO[1]])[0, 1])
print()
print(f"ANSWER KEY -- what a reference change does: the three references give window-mean maps whose "
      f"across-channel correlation is r = {r_lm:.3f} (exactly 1 -- a reference change is an additive "
      f"constant across channels) but whose RMS amplitudes are "
      + ", ".join(f"{np.sqrt(np.mean(maps[n] ** 2)):.3f} uV ({n.split('(')[0].strip()})" for n in names)
      + ".  Mean-centring the maps makes them identical to machine precision, so a reference change moves "
        "the peak channel and the zero line but never the shape.")
print(f"ANSWER KEY -- ex-3-4 (which maps differ only in amplitude, which in shape), average reference, "
      f"{len(SUBJECTS)} subjects:")
for i, a in enumerate(keys):
    print(f"    {a:26s} RMS {np.sqrt(np.mean(CANDIDATES[a] ** 2)):6.3f} uV; r with A = {corr[0, i]:6.3f}")
print(f"    amplitude-only pair: A (difference 400-500 ms) and B (difference 500-600 ms), "
      f"r = {corr[0, 1]:.3f}, RMS {np.sqrt(np.mean(CANDIDATES[keys[0]] ** 2)):.3f} vs "
      f"{np.sqrt(np.mean(CANDIDATES[keys[1]] ** 2)):.3f} uV")
print(f"    shape pair: A and C (difference 150-250 ms), r = {corr[0, 2]:.3f}")
print(f"    D and E (the two conditions' early sensory maps) agree in both: r = {corr[3, 4]:.3f}, "
      f"RMS {np.sqrt(np.mean(CANDIDATES[keys[3]] ** 2)):.3f} vs "
      f"{np.sqrt(np.mean(CANDIDATES[keys[4]] ** 2)):.3f} uV")
print(f"Caveat printed with every linked-mastoid number: the ERP CORE 30-channel montage has no mastoid "
      f"electrode; P9 and P10 are the nearest inferior posterior sites and stand in for M1/M2 "
      f"(TODO(confirm)).")
print(f"Pitfall: pf-reference-changes-everything.  Widget: w-reference-explorer (mode topo).")