# 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", "pooch", "sklearn")
_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_l5.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L5/ (or notebooks/) so that _shared/helpers_l5.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l5 as L5

# 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
import pooch

mne.set_log_level("WARNING")
pooch.get_logger().setLevel("WARNING")   # no download chatter: it would print local paths
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers_l5 imported from notebooks/_shared")
print("mne-connectivity available:", L5.have_mne_connectivity())
print(L5.disk_line("disk at the start"))

print("Licences, from data/directory.yaml (never from memory):")
L5.print_licences("ds-eegbci", notes=True)

import warnings
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import Pipeline
from mne.decoding import CSP

PIPELINE = {
    "dataset": "ds-eegbci",
    "runs": (4, 8, 12),
    "runs_note": "motor imagery of the left (T1) or right (T2) fist",
    "band_hz": (8.0, 30.0),
    "filter": "FIR, zero-phase, MNE defaults",
    "epoch_s": (-1.0, 4.0),
    "window_s": (0.5, 2.5),
    "window_note": "the analysis window, inside the epoch, from cue onset",
    "n_components": 6,
    "estimator": "Pipeline([('csp', mne.decoding.CSP(n_components=6, log=True, norm_trace=False)), "
                 "('lda', LinearDiscriminantAnalysis())])",
    "cv": "StratifiedKFold(n_splits=5, shuffle=True, random_state=20260918)",
    "seed": L5.SEED,
}
CANDIDATES = tuple(range(1, 13))
FEATURED = 2
BAND, EPOCH_S, WINDOW_S = PIPELINE["band_hz"], PIPELINE["epoch_s"], PIPELINE["window_s"]
for k, v in PIPELINE.items():
    print(f"  {k:14s}: {v}")
print()
print(f"candidate subjects: {['S%03d' % s for s in CANDIDATES]}")
print(f"featured subject (chosen by accuracy -- see section 6): S{FEATURED:03d}")
print()
print(L5.disk_line("disk before the downloads"))


def load_subject(subject, runs=PIPELINE["runs"]):
    """One subject's band-passed imagery epochs; the EDF files are deleted before this returns."""
    paths = []
    try:
        raw = helpers.load_spine("ds-eegbci", subject, list(runs), preload=True)
        paths = L5.eegbci_files(subject, runs)
        raw.filter(BAND[0], BAND[1], verbose=False)
        events, event_id = mne.events_from_annotations(raw, verbose=False)
        wanted = {k: v for k, v in event_id.items() if k in ("T1", "T2")}
        epochs = mne.Epochs(raw, events, wanted, tmin=EPOCH_S[0], tmax=EPOCH_S[1], baseline=None,
                            picks="eeg", preload=True, verbose=False)
        y = np.array([1 if e == wanted["T1"] else 2 for e in epochs.events[:, 2]])
        return epochs, y
    finally:
        L5.delete_files(paths, verbose=False)


data, survey_meta = {}, []
for s in CANDIDATES:
    ep, y = load_subject(s)
    X = ep.copy().crop(*WINDOW_S).get_data(copy=True)
    data[s] = (X, y, ep.info.copy(), ep.ch_names)
    survey_meta.append((s, len(y), int((y == 1).sum()), int((y == 2).sum())))
    print(f"  S{s:03d}: {len(y)} epochs ({int((y == 1).sum())} left, {int((y == 2).sum())} right), "
          f"{X.shape[1]} channels x {X.shape[2]} samples in the {WINDOW_S[0]:g}-{WINDOW_S[1]:g} s window")
print()
print(L5.disk_line("disk after every download was deleted"))

X_f, y_f, info_f, ch_f = data[FEATURED]
print()
print(f"featured subject S{FEATURED:03d}: {X_f.shape[0]} epochs x {X_f.shape[1]} channels x "
      f"{X_f.shape[2]} samples at {info_f['sfreq']:g} Hz")

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    csp_display = CSP(n_components=PIPELINE["n_components"], log=True, norm_trace=False, reg=None)
    csp_display.fit(X_f, y_f)

# Recompute the eigenvalues from the stored filters rather than quoting them, so the two cannot disagree.
S1 = np.mean([np.cov(x) for x in X_f[y_f == 1]], axis=0)
S2 = np.mean([np.cov(x) for x in X_f[y_f == 2]], axis=0)
eig = []
for w in csp_display.filters_[:PIPELINE["n_components"]]:
    eig.append(float(w @ S1 @ w / (w @ (S1 + S2) @ w)))
print(f"CSP eigenvalues of the {PIPELINE['n_components']} stored filters "
      f"(lambda = w'S_left w / w'(S_left + S_right) w):")
print("  " + "  ".join(f"{e:.4f}" for e in eig))
print("  1 = all the variance is in left imagery, 0 = all in right, 0.5 = the filter does not separate them")
print()
print(f"w-csp-explorer ships {[0.176603, 0.245932, 0.263226, 0.681189, 0.319435, 0.321245]} for the same")
print("subject, runs, band and window.  CSP eigenvectors are defined up to sign and their ORDER depends on")
print("the library's sorting convention, so the two lists are compared as sets of distances from 0.5, not")
print("element by element:")
mine = sorted(abs(e - 0.5) for e in eig)
theirs = sorted(abs(e - 0.5) for e in [0.176603, 0.245932, 0.263226, 0.681189, 0.319435, 0.321245])
print(f"  this notebook : {[round(v, 4) for v in mine]}")
print(f"  the widget    : {[round(v, 4) for v in theirs]}")
print(f"  largest difference: {max(abs(a - b) for a, b in zip(mine, theirs)):.4f}")

n_show = 4
fig, axes = plt.subplots(2, n_show, figsize=(3.0 * n_show, 6.0))
for k in range(n_show):
    im0, _ = mne.viz.plot_topomap(csp_display.filters_[k], info_f, axes=axes[0, k], show=False,
                                  contours=4, sensors=True)
    im1, _ = mne.viz.plot_topomap(csp_display.patterns_[k], info_f, axes=axes[1, k], show=False,
                                  contours=4, sensors=True)
    axes[0, k].set_title(f"component {k + 1}\nlambda = {eig[k]:.3f}", fontsize=9)
axes[0, 0].set_ylabel("FILTER (not interpretable)", fontsize=9)
axes[1, 0].set_ylabel("PATTERN (interpretable)", fontsize=9)
cb0 = fig.colorbar(im0, ax=axes[0, :], shrink=0.8); cb0.set_label("filter weight (arbitrary units)")
cb1 = fig.colorbar(im1, ax=axes[1, :], shrink=0.8); cb1.set_label("pattern (arbitrary units)")
fig.suptitle(f"ds-eegbci S{FEATURED:03d}, {BAND[0]:g}-{BAND[1]:g} Hz: CSP filters above, patterns below -- "
             f"the same components, and they do not look alike", y=1.0, fontsize=11)
plt.show()   # render the static figure(s) of this cell inline

print("Agreement between each filter and its own pattern, as an across-channel correlation:")
for k in range(PIPELINE["n_components"]):
    r = float(np.corrcoef(csp_display.filters_[k], csp_display.patterns_[k])[0, 1])
    fp = ch_f[int(np.argmax(np.abs(csp_display.filters_[k])))]
    pp = ch_f[int(np.argmax(np.abs(csp_display.patterns_[k])))]
    print(f"  component {k + 1}: r = {r:+.3f}; largest |filter| weight at {fp:>4s}, "
          f"largest |pattern| value at {pp:>4s}")
print()
print("Where the two disagree about which channel matters most, the FILTER is the one that is not answering")
print("the question a reader will assume it answers.  TODO(confirm): whether the physiologically plausible")
print("pattern for this subject is the one over left or right sensorimotor cortex is a judgement for the")
print("author; the notebook reports which channels carry the extrema and does not label them.")

CV = StratifiedKFold(n_splits=5, shuffle=True, random_state=L5.SEED)


def csp_lda(n_components):
    return Pipeline([("csp", CSP(n_components=n_components, log=True, norm_trace=False, reg=None)),
                     ("lda", LinearDiscriminantAnalysis())])


def cv_scores(X, y, estimator=None, n_components=PIPELINE["n_components"]):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        return cross_val_score(estimator or csp_lda(n_components), X, y, cv=CV, n_jobs=1)


scores = cv_scores(X_f, y_f)
ACC = float(scores.mean())
print(f"ds-eegbci S{FEATURED:03d}, CSP-{PIPELINE['n_components']} + LDA, "
      f"{CV.get_n_splits()}-fold stratified cross-validation:")
print(f"  fold scores: {[f'{s:.4f}' for s in scores]}")
print(f"  mean {ACC:.4f}, SD across folds {scores.std():.4f}")
print()
n = len(y_f)
from scipy import stats
lo, hi = stats.binom.ppf([0.025, 0.975], n, 0.5) / n
print(f"  chance is 0.5 with two balanced classes, but with {n} epochs the binomial 95 % interval around")
print(f"  chance runs from {lo:.3f} to {hi:.3f}.  A single subject inside that band is not evidence of "
      f"anything.")
print()
print(f"  w-csp-explorer reports 0.8667 for the same subject, runs, band, window, folds and seed.")
print(f"  This notebook: {ACC:.4f}.  Difference: {ACC - 0.8667:+.4f}")

class BandPower:
    """Log variance of a fixed channel subset: the simplest possible 'spatial filter', namely none."""

    def __init__(self, picks):
        self.picks = picks

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        return np.log(X[:, self.picks, :].var(axis=-1))

    def fit_transform(self, X, y=None):
        return self.transform(X)

    def get_params(self, deep=True):
        return {"picks": self.picks}

    def set_params(self, **kw):
        self.picks = kw.get("picks", self.picks)
        return self


c3c4 = [ch_f.index("C3"), ch_f.index("C4")]
stages = [("log band power at C3 and C4 only (no spatial filter)",
           Pipeline([("bp", BandPower(c3c4)), ("lda", LinearDiscriminantAnalysis())]), 2)]
for k in (2, 4, 6, 8):
    stages.append((f"CSP, {k} components -> LDA", csp_lda(k), k))

WIDGET_STAGES = {"log band power at C3 and C4 only (no spatial filter)": 0.7333,
                 "CSP, 2 components -> LDA": 0.5111, "CSP, 4 components -> LDA": 0.8667,
                 "CSP, 6 components -> LDA": 0.8667, "CSP, 8 components -> LDA": 0.8889}
print(f"{'feature set':>52s} {'features':>9s} {'mean':>7s} {'SD':>7s} {'widget':>8s} {'difference':>11s}")
stage_rows = []
for label, est, n_feat in stages:
    sc = cv_scores(X_f, y_f, estimator=est)
    w = WIDGET_STAGES[label]
    stage_rows.append((label, n_feat, float(sc.mean()), float(sc.std()), w))
    print(f"{label:>52s} {n_feat:9d} {sc.mean():7.4f} {sc.std():7.4f} {w:8.4f} {sc.mean() - w:+11.4f}")
print()
print("CSP with two components is WORSE than no spatial filter at all on this subject.  That is a real result")
print("and it is worth keeping: two components means one filter per class, and one filter per class is not")
print("enough to describe this subject's sensorimotor rhythms.  More components is not monotonically better")
print("either -- it is a hyperparameter, and choosing it by looking at the cross-validated score is itself a")
print("selection (Level 6).")

# The leak: fit CSP once on everything, cross-validate only the classifier.
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    leaked_csp = CSP(n_components=PIPELINE["n_components"], log=True, norm_trace=False, reg=None)
    F = leaked_csp.fit_transform(X_f, y_f)           # <- the labels of every epoch are used here
    leaked = cross_val_score(LinearDiscriminantAnalysis(), F, y_f, cv=CV, n_jobs=1)
LEAKED = float(leaked.mean())
print(f"honest  (CSP refitted inside each training fold): {ACC:.4f}   folds {[f'{s:.3f}' for s in scores]}")
print(f"leaked  (CSP fitted once on all {len(y_f)} epochs) : {LEAKED:.4f}   "
      f"folds {[f'{s:.3f}' for s in leaked]}")
print(f"inflation: {LEAKED - ACC:+.4f} ({100 * (LEAKED - ACC):+.1f} percentage points)")
print()
print("Same data, same folds, same seed, same classifier.  The only difference is where the fold boundary")
print("sits relative to the spatial filter.  w-csp-explorer reports 0.8667 honest against 0.9778 leaked, an")
print(f"inflation of 0.1111; this notebook gets {ACC:.4f} against {LEAKED:.4f}, an inflation of "
      f"{LEAKED - ACC:.4f}.")
print()
print("Note what it does NOT look like: the leaked number is not absurd, it does not fail any internal check,")
print("and nothing in the output says 'leak'.  It is simply 11 points too high.")

survey = []
for s in CANDIDATES:
    X, y, _, _ = data[s]
    sc = cv_scores(X, y)
    survey.append((s, len(y), float(sc.mean()), float(sc.std()), [float(v) for v in sc]))
means = np.array([r[2] for r in survey])
print(f"CSP-{PIPELINE['n_components']} + LDA, identical pipeline, {len(survey)} candidate subjects:")
print(f"{'subject':>8s} {'epochs':>7s} {'mean':>7s} {'SD':>7s}  fold scores")
for s, n_ep, m, sd, sc in survey:
    mark = "  <- the featured subject" if s == FEATURED else ""
    print(f"{'S%03d' % s:>8s} {n_ep:7d} {m:7.4f} {sd:7.4f}  {[f'{v:.3f}' for v in sc]}{mark}")
print()
print(f"  median {np.median(means):.4f}   mean {means.mean():.4f}   range {means.min():.4f}-{means.max():.4f}")
print(f"  subjects above the upper edge of the chance band ({hi:.3f}): "
      f"{int((means > hi).sum())} of {len(means)}")
print()
print("A FINDING, reported rather than reconciled.  The Phase 3 integration notes record the Level 5 decoding")
print("distribution as 'a median of 60.0 %' with a range of 44.4-95.6 %.  The range matches this run and the")
print("widget's own stored survey exactly.  The median does not: the twelve subject means stored in")
print("w-csp-explorer/csp.json have a median of 0.5556 and a MEAN of 0.6037, and 60.0 % is the mean, not the")
print(f"median.  This notebook's own twelve subjects give median {np.median(means):.4f} and mean "
      f"{means.mean():.4f}.")
print("Both statistics are worth quoting and they are not interchangeable; the lesson should say which it")
print("means.  TODO(confirm) for the author.")

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2))
o = np.argsort(means)
axes[0].bar(range(len(survey)), means[o], color=["C1" if survey[i][0] == FEATURED else "C0" for i in o])
axes[0].axhline(0.5, color="k", lw=1, ls="--")
axes[0].axhspan(lo, hi, color="grey", alpha=0.25)
axes[0].axhline(np.median(means), color="C2", lw=1.5, label=f"median {np.median(means):.3f}")
axes[0].axhline(means.mean(), color="C3", lw=1.5, ls=":", label=f"mean {means.mean():.3f}")
axes[0].set_xticks(range(len(survey)))
axes[0].set_xticklabels([f"S{survey[i][0]:03d}" for i in o], rotation=45, fontsize=7)
axes[0].set_ylabel("cross-validated accuracy (proportion correct)")
axes[0].set_title(f"twelve candidates; the grey band is the 95 % chance interval for {n} epochs", fontsize=9)
axes[0].legend(fontsize=8); axes[0].set_ylim(0, 1.05)

labels = [r[0].replace(" -> LDA", "").replace(" (no spatial filter)", "") for r in stage_rows]
axes[1].bar(range(len(stage_rows)), [r[2] for r in stage_rows],
            yerr=[r[3] for r in stage_rows], capsize=3, color="C0")
axes[1].bar([len(stage_rows)], [LEAKED], color="C3")
axes[1].set_xticks(range(len(stage_rows) + 1))
axes[1].set_xticklabels(labels + ["CSP-6 LEAKED"], rotation=30, ha="right", fontsize=7)
axes[1].axhline(0.5, color="k", lw=1, ls="--")
axes[1].axhspan(lo, hi, color="grey", alpha=0.25)
axes[1].set_ylabel("cross-validated accuracy (proportion correct)")
axes[1].set_title(f"S{FEATURED:03d}: feature set, and the cost of one misplaced fold boundary", fontsize=9)
axes[1].set_ylim(0, 1.05)
fig.suptitle("ds-eegbci motor imagery, 8-30 Hz, 0.5-2.5 s window, 5-fold stratified cross-validation",
             y=1.02, fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print("nb-5-7-csp -- L5.7 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-eegbci runs R04+R08+R12, {BAND[0]:g}-{BAND[1]:g} Hz, epochs "
      f"{EPOCH_S[0]:g} to {EPOCH_S[1]:g} s, analysis window {WINDOW_S[0]:g}-{WINDOW_S[1]:g} s, "
      f"{X_f.shape[1]} channels at {info_f['sfreq']:g} Hz")
print(f"      {L5.licence_line('ds-eegbci')}")
print(f"Estimator: {PIPELINE['estimator']}")
print(f"Cross-validation: {PIPELINE['cv']}")
print(f"Downloads: 3 EDF runs x {len(CANDIDATES)} subjects, about 7.7 MB each, each deleted as soon as its "
      f"epochs were in memory")
print()
print(f"ANSWER KEY -- ex-5-7 (cross-validated accuracy for one subject): S{FEATURED:03d}, "
      f"{ACC:.4f} ({100 * ACC:.1f} %)")
print(f"    fold scores {[f'{s:.4f}' for s in scores]}, SD across folds {scores.std():.4f}, "
      f"{len(y_f)} epochs")
print(f"    THE SUBJECT WAS SELECTED BY ACCURACY.  Across the {len(survey)} candidates the range is "
      f"{means.min():.3f}-{means.max():.3f},")
print(f"    the median is {np.median(means):.4f} and the mean is {means.mean():.4f}.  Quoting "
      f"{100 * ACC:.1f} % on its own is the selection")
print("    effect Level 6 exists to warn about, so the key must carry the median and the range beside it.")
print(f"    Chance band for {len(y_f)} epochs: {lo:.3f} to {hi:.3f}; "
      f"{int((means > hi).sum())} of {len(means)} subjects beat it.")
print()
print("    per-subject distribution:")
for s, n_ep, m, sd, sc in survey:
    print(f"      S{s:03d}  {m:.4f}  (SD {sd:.4f})" + ("   <- featured" if s == FEATURED else ""))
print()
print("ANSWER KEY -- feature set (same subject, same folds, same seed):")
print(f"{'':6s}{'feature set':>52s} {'here':>8s} {'widget':>8s}")
for label, n_feat, m, sd, w in stage_rows:
    print(f"{'':6s}{label:>52s} {m:8.4f} {w:8.4f}")
print(f"{'':6s}CSP-2 scores {dict((r[0], r[2]) for r in stage_rows)['CSP, 2 components -> LDA']:.4f}, "
      f"which is WORSE than the "
      f"{dict((r[0], r[2]) for r in stage_rows)['log band power at C3 and C4 only (no spatial filter)']:.4f} "
      f"of two raw channels.")
print()
print(f"ANSWER KEY -- the cost of leakage (ex-5-7 / ex-6-5): honest {ACC:.4f} against leaked {LEAKED:.4f}, "
      f"an inflation of")
print(f"    {LEAKED - ACC:+.4f} ({100 * (LEAKED - ACC):+.1f} percentage points).  Same subject, same epochs, "
      f"same folds, same seed; the only")
print("    change is whether CSP is fitted inside the cross-validation or once on everything.")
print(f"    w-csp-explorer reports 0.8667 against 0.9778, an inflation of 0.1111.")
print()
print("ANSWER KEY -- filters against patterns:")
for k in range(PIPELINE["n_components"]):
    r = float(np.corrcoef(csp_display.filters_[k], csp_display.patterns_[k])[0, 1])
    print(f"      component {k + 1}: lambda {eig[k]:.4f}, filter/pattern correlation r = {r:+.3f}, "
          f"|filter| peak {ch_f[int(np.argmax(np.abs(csp_display.filters_[k])))]}, "
          f"|pattern| peak {ch_f[int(np.argmax(np.abs(csp_display.patterns_[k])))]}")
print("      Only the patterns may be read as topographies.  The filters and patterns of the same component")
print("      disagree about which channel matters most, which is exactly why reading a filter is a mistake.")
print()
print("Pitfall: pf-decoding-leakage.  Widget: w-csp-explorer (mode patterns).")
print(L5.disk_line("disk at the end"))