# Setup: dependencies, the shared helpers, non-interactive plotting.
import hashlib
import importlib.util
import subprocess
import sys
import time
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_l6.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L6/ (or notebooks/) so that _shared/helpers_l6.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l3 as L3
import helpers_l6 as L6

# 3. Plotting: static PNGs through the inline backend; every figure cell ends with plt.show().
import matplotlib.pyplot as plt
import numpy as np
import mne
from scipy import stats

mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72

# 4. Iteration counts.  Everything expensive here is a null draw, and a draw is cheap because the
#    per-trial measurements are precomputed; FULL_RUN raises the draw count where it buys precision.
FULL_RUN = False
N_DRAWS = 200 if not FULL_RUN else 2000         # null datasets for the circularity and garden sections
SEED = L6.SEED
ALPHA = L6.ALPHA

print(f"MNE {mne.__version__}; helpers_l6 imported from notebooks/_shared")
print(f"FULL_RUN = {FULL_RUN}: {N_DRAWS} null datasets")
print(f"ERP CORE cache: {L3.erpcore_root().name}/ (resolved relative to the working directory, or "
      f"$EEG_COURSE_ERPCORE); each subject's EEGLAB pair is deleted as soon as its measurements exist")

SUBJECTS = list(L6.SUBSET_DEFAULT)
before = L6.disk_report("before any download")
t0 = time.time()
cohort = L6.load_cohort(SUBJECTS)
print(f"{len(cohort)} subjects in {time.time() - t0:.0f} s")
after = L6.disk_report("after the download loop")
print(f"the ~58 MB EEGLAB pair of each subject is deleted before the next is fetched, so the folder sizes "
      f"above are the whole footprint; free space itself moves for reasons a notebook knows nothing about")

paths = L6.all_paths()
subs = sorted(cohort)
times = cohort[subs[0]]["times"].astype(float)
print(f"\n{len(paths)} paths from {len(L6.CHOICES)} decisions:")
for c in L6.CHOICES:
    print(f"   {c['id']:11s} ({len(c['options'])}) {', '.join(c['options'])}")
    print(f"               why each is defensible: {L6.CHOICE_RATIONALE[c['id']]}")

# The nine window x electrode measurements, with the other four decisions at their pre-specified values.
combos = [dict(L6.PRESPECIFIED, window=w, electrode=e)
          for w in ("300-500", "300-600", "350-650") for e in ("Pz", "CPz", "Pz+CPz+Cz")]
combo_idx = np.array([L6.path_index(c, paths) for c in combos])
i_prespec = L6.path_index(L6.PRESPECIFIED, paths)
j_prespec = int(np.where(combo_idx == i_prespec)[0][0])
print(f"{len(combos)} window x electrode measurements (the other four decisions fixed at "
      f"{L6.PRESPECIFIED['reference']} / {L6.PRESPECIFIED['highpass']} Hz / "
      f"{L6.PRESPECIFIED['baseline']} / {L6.PRESPECIFIED['rejection']}):")
for j, c in enumerate(combos):
    print(f"   {j}: {c['window']:8s} ms at {c['electrode']:10s}" + ("   <- pre-specified" if j == j_prespec else ""))

MIN_TRIALS = int(L6.FIXED_CHOICES["min_trials_per_condition"])
usable = [s for s in subs if min(cohort[s]["n_kept"].values()) >= MIN_TRIALS]
print(f"\n{len(usable)} of {len(subs)} subjects meet the stated minimum of {MIN_TRIALS} surviving trials per "
      f"condition under this rejection setting"
      + (f" (excluded: {', '.join(s for s in subs if s not in usable)})" if len(usable) < len(subs) else ""))

# collapsed-localizer statistic: each subject's mean amplitude over ALL trials, per measurement.
# It is computed once because it does not depend on the condition labels at all -- which is the
# property that makes it a legitimate selector.
collapsed = np.stack([
    np.array([float(cohort[s]["amp"][i][cohort[s]["keep"][i]].mean()) for i in combo_idx])
    for s in usable])                                                    # subjects x 9
print(f"\ncollapsed-localizer statistic (group mean over all trials, uV): "
      + ", ".join(f"{v:+.2f}" for v in collapsed.mean(0)))
j_collapsed = int(np.argmax(np.abs(collapsed.mean(0))))
print(f"   it selects measurement {j_collapsed}: {combos[j_collapsed]['window']} ms at "
      f"{combos[j_collapsed]['electrode']}.  Under label permutation this choice never changes, because "
      f"nothing it looks at depends on the labels.")

def selection_rates(label_source, n_draws, seed_offset):
    """False-positive (or hit) rate of each selection rule over n_draws label assignments.

    A permuted label vector can leave a subject below the 15-trial minimum on some measurements
    even when its true labels do not, so `helpers_l6.subject_diffs` returns NaN there and the
    group test is the NaN-aware `helpers_l6.paired_t` rather than `scipy.stats.ttest_1samp`.
    `n_nan` counts how often that happened, because silently turning a NaN into "not significant"
    would bias every rate downwards.
    """
    rng = {s: np.random.default_rng(SEED + seed_offset + int(s.split("-")[1])) for s in usable}
    out = {k: np.zeros(n_draws, bool) for k in ("pre-specified", "best-of-9", "collapsed", "loso")}
    tvals = {k: np.zeros(n_draws) for k in out}
    n_nan = 0
    for d in range(n_draws):
        D = np.empty((len(usable), len(combos)))                          # subjects x 9 estimates
        for k, s in enumerate(usable):
            lab = (cohort[s]["is_target"] if label_source == "real"
                   else rng[s].permutation(cohort[s]["is_target"]))
            D[k] = L6.subject_diffs(cohort[s], lab)[combo_idx]
        n_nan += int((~np.isfinite(D)).any())
        t, p, _m, _n = L6.paired_t(D.T)                                   # 9 group tests, NaN-aware

        out["pre-specified"][d] = p[j_prespec] < ALPHA
        tvals["pre-specified"][d] = t[j_prespec]

        j_best = int(np.nanargmax(np.abs(t)))                             # <- selected ON the tested contrast
        out["best-of-9"][d] = p[j_best] < ALPHA
        tvals["best-of-9"][d] = t[j_best]

        out["collapsed"][d] = p[j_collapsed] < ALPHA
        tvals["collapsed"][d] = t[j_collapsed]

        scores = np.empty(len(usable))                                    # leave-one-subject-out selection
        for k in range(len(usable)):
            others = np.delete(np.arange(len(usable)), k)
            t_o, _p, _m, _n = L6.paired_t(D[others].T)
            scores[k] = D[k, int(np.nanargmax(np.abs(t_o)))]
        t_l, p_l, _m, _n = L6.paired_t(scores)
        out["loso"][d] = bool(p_l < ALPHA)
        tvals["loso"][d] = float(t_l)
    return out, tvals, n_nan

t0 = time.time()
null_rates, null_t, n_nan = selection_rates("null", N_DRAWS, seed_offset=0)
print(f"{N_DRAWS} null datasets x 4 selection rules in {time.time() - t0:.1f} s")
print(f"draws in which at least one subject fell below the {MIN_TRIALS}-trial minimum on at least one "
      f"measurement, and was dropped from it: {n_nan} of {N_DRAWS}\n")
lo, hi = stats.binom.interval(0.95, N_DRAWS, ALPHA)
print(f"{'selection rule':>22s} {'false-positive rate':>20s} {'mean |t|':>9s}  verdict")
for k in null_rates:
    r = null_rates[k].mean()
    inside = lo / N_DRAWS <= r <= hi / N_DRAWS
    verdict = ("consistent with alpha" if inside else
               f"{r / ALPHA:.1f}x alpha -- OUTSIDE the binomial interval")
    print(f"{k:>22s} {r:19.1%} {np.abs(null_t[k]).mean():9.2f}  {verdict}")
print(f"\nbinomial 95% interval for a true {ALPHA:.0%} rate over {N_DRAWS} draws: "
      f"{lo / N_DRAWS:.1%} to {hi / N_DRAWS:.1%}")

real_rates, real_t, _ = selection_rates("real", 1, seed_offset=0)
print("The same four rules on the REAL contrast (one dataset, so these are single results, not rates):")
for k in real_rates:
    print(f"   {k:>16s}: t = {real_t[k][0]:+.3f}, "
          f"{'reaches' if real_rates[k][0] else 'does not reach'} p < {ALPHA}")
print(f"\nOn a real effect this size every rule agrees, which is exactly why the null simulation above is the "
      f"only way to tell them apart.  A selection rule that is wrong 4 times in 10 under the null still looks "
      f"fine on data with a large effect in it -- and most published contrasts are nowhere near this large.")

diffs = L6.cohort_diffs(cohort)                  # (486 paths, 20 subjects), NaN where a path runs out of trials
grid = L6.grid_summary(paths, diffs, alpha=ALPHA)
print(f"MULTIVERSE: {grid['n_significant']} of {grid['n_paths']} paths reach p < {ALPHA} "
      f"({grid['n_significant'] / grid['n_paths']:.1%})")
print(f"   estimate  {grid['effect_min']:+.4f} to {grid['effect_max']:+.4f} uV, median "
      f"{grid['effect_median']:+.4f}, IQR {grid['effect_iqr']:.4f}")
print(f"   all paths point the same way: {grid['all_same_sign']}")
print(f"   p         min {grid['p_min']:.3g}, median {grid['p_median']:.4f}, max {grid['p_max']:.4f}")
print(f"   ratio of the largest to the smallest estimate: "
      f"{grid['effect_max'] / grid['effect_min']:.1f}x")
print(f"\n   the pre-specified path: {grid['effect_uv'][i_prespec]:+.4f} uV, p = {grid['p'][i_prespec]:.3g}, "
      f"{int(grid['n_subjects'][i_prespec])} subjects")
print(f"   subjects contributing per path: min {int(np.min(grid['n_subjects']))}, "
      f"max {int(np.max(grid['n_subjects']))} (a path with a tight rejection threshold loses subjects who "
      f"fall below the {MIN_TRIALS}-trial minimum)")

fig, axes = L6.plot_specification_curve(
    grid["effect_uv"], grid["p"], paths, alpha=ALPHA, highlight=L6.PRESPECIFIED,
    title=f"Specification curve: ERP CORE P3 target minus standard, {len(paths)} defensible pipelines (µV)")
plt.show()   # render the static figure(s) of this cell inline

rows = L6.marginals(paths, grid["effect_uv"], grid["significant"])
print("Marginals: each option's median estimate with every other decision averaged over.")
print(f"{'decision':>11s} {'option':>17s} {'paths':>6s} {'median (uV)':>12s} {'range (uV)':>20s} "
      f"{'p < .05':>9s}")
spread = {}
for r in rows:
    print(f"{r['choice']:>11s} {r['option']:>17s} {r['n_paths']:6d} {r['effect_median']:+12.4f} "
          f"{r['effect_min']:+9.2f} to {r['effect_max']:+6.2f} {r['n_significant']:4d}/{r['n_paths']:<4d}")
    spread.setdefault(r["choice"], []).append(r["effect_median"])
print()
print("Which single decision moves the estimate most (spread of the option medians):")
for k, v in sorted(spread.items(), key=lambda kv: -(max(kv[1]) - min(kv[1]))):
    print(f"   {k:11s} {max(v) - min(v):6.3f} uV   ({min(v):+.3f} to {max(v):+.3f})")
worst = max(spread, key=lambda k: max(spread[k]) - min(spread[k]))
print(f"\nThe {worst} decision alone moves the median estimate by "
      f"{max(spread[worst]) - min(spread[worst]):.2f} uV.  Two analysts who agree about every other step and "
      f"differ only there will report numbers that far apart from the same recordings, and both methods "
      f"sections will be correct.")

fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.4))
order = [c["id"] for c in L6.CHOICES]
pos, labels, data = [], [], []
y = 0
for cid in order:
    for option in next(c["options"] for c in L6.CHOICES if c["id"] == cid):
        m = np.array([p[cid] == option for p in paths])
        data.append(grid["effect_uv"][m][np.isfinite(grid["effect_uv"][m])])
        labels.append(f"{cid}: {option}")
        pos.append(y)
        y += 1
    y += 0.6
bp = axes[0].boxplot(data, positions=pos, vert=False, widths=0.6, showfliers=False, patch_artist=True)
for patch in bp["boxes"]:
    patch.set_facecolor("tab:blue")
    patch.set_alpha(0.45)
axes[0].axvline(grid["effect_median"], color="k", ls=":", lw=1,
                label=f"overall median {grid['effect_median']:+.2f} µV")
axes[0].axvline(0, color="gray", lw=0.8)
axes[0].set_yticks(pos, labels, fontsize=8)
axes[0].set(xlabel="Estimate (µV)", title="Each option's estimates, everything else averaged over")
axes[0].grid(alpha=0.3, axis="x")
axes[0].legend(fontsize=8, loc="lower right")

axes[1].hist(grid["effect_uv"][np.isfinite(grid["effect_uv"])], bins=30, color="tab:blue", alpha=0.85)
axes[1].axvline(grid["effect_uv"][i_prespec], color="tab:orange", lw=2,
                label=f"pre-specified path {grid['effect_uv'][i_prespec]:+.2f} µV")
axes[1].axvline(0, color="gray", lw=0.8)
axes[1].set(xlabel="Estimate (µV)", ylabel="Paths",
            title=f"{grid['n_paths']} defensible pipelines, one dataset")
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

PREREGISTRATION = """
1. HYPOTHESIS   Rare target letters elicit a larger positive deflection than frequent standards over
                centro-parietal scalp, 300-600 ms after stimulus onset (the P3b).
2. DATA         ds-erpcore, paradigm P3, sub-001 to sub-020.  No subject is inspected before the plan is fixed.
3. EXCLUSION    A subject is excluded if fewer than 15 trials survive artifact rejection in either condition.
                Decided before analysis; the count and the reason are reported for every exclusion.
4. PREPROCESS   Resample 1024 -> 256 Hz; FIR band-pass 0.1-30 Hz; bipolar HEOG/VEOG; ocular correction by
                regression fitted on epochs with each condition's evoked response removed.
5. REFERENCE    Average of the 30 EEG channels.
6. EPOCHS       -200 to +800 ms, baseline -200 to 0 ms.  Reject an epoch whose peak-to-peak amplitude
                exceeds 100 uV on any EEG channel.
7. MEASURE      Mean amplitude 300-600 ms at Pz, target minus standard, one value per subject.
8. TEST         Two-sided one-sample t-test against zero, alpha = .05.  Report the estimate with a 95 %
                confidence interval and Cohen's dz.
9. CORRECTION   One test, so none is required.  Any additional time point, channel or window reported is
                exploratory and labelled as such, with a cluster-based permutation test over the tested space.
10. DEVIATIONS  Every departure from lines 1-9 is reported with the reason and the stage at which it was made.
"""
digest = hashlib.sha256(PREREGISTRATION.encode("utf-8")).hexdigest()
print(PREREGISTRATION)
print(f"SHA-256 of the plan above: {digest}")
print(f"length: 10 numbered items over {len(PREREGISTRATION.strip().splitlines())} physical lines, "
      f"{len(PREREGISTRATION)} characters")
print()
print("What each line buys, in the terms of this notebook:")
print(f"   line 3  removes the exclusion fork      -- and it fires: "
      f"{len(subs) - len(usable)} of {len(subs)} subjects excluded under this pipeline")
print(f"   line 5  removes the reference fork      -- worth {max(spread['reference']) - min(spread['reference']):.2f} uV "
      f"of spread in the multiverse above")
print(f"   line 6  removes baseline and rejection  -- worth "
      f"{max(spread['baseline']) - min(spread['baseline']):.2f} and "
      f"{max(spread['rejection']) - min(spread['rejection']):.2f} uV")
print(f"   line 7  removes window and electrode    -- worth "
      f"{max(spread['window']) - min(spread['window']):.2f} and "
      f"{max(spread['electrode']) - min(spread['electrode']):.2f} uV, and takes the false-positive rate from "
      f"{null_rates['best-of-9'].mean():.1%} back to {null_rates['pre-specified'].mean():.1%}")
print(f"   line 4  removes the high-pass fork      -- worth {max(spread['highpass']) - min(spread['highpass']):.2f} uV")
print(f"   Together lines 3-7 turn {len(paths)} pipelines into 1.")

print("nb-6-4-multiverse -- L6.4 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, sub-001 to sub-{SUBJECTS[-1]:03d} ({len(subs)} loaded, {len(usable)} meeting the "
      f"{MIN_TRIALS}-trial minimum under the pre-specified rejection); CC-BY-SA-4.0 per data/directory.yaml")
print(f"Pipeline: helpers_l6.PRESPECIFIED = {L6.describe_path(L6.PRESPECIFIED)}")
print()
print(f"1. CIRCULARITY -- false-positive rate of four selection rules over {N_DRAWS} label-permuted datasets "
      f"(9 window x electrode measurements, everything else fixed):")
for k in null_rates:
    print(f"     {k:>16s}: {null_rates[k].mean():6.1%}"
          + ("   <- nominal alpha" if k == "pre-specified" else "")
          + ("   <- double dipping" if k == "best-of-9" else ""))
print(f"     inflation from choosing the best of {len(combos)} after the fact: "
      f"{null_rates['best-of-9'].mean() - null_rates['pre-specified'].mean():+.1%} "
      f"({null_rates['best-of-9'].mean() / max(null_rates['pre-specified'].mean(), 1e-9):.1f}x)")
print()
print(f"2. MULTIVERSE -- {len(paths)} defensible pipelines on the real contrast:")
print(f"     significant      : {grid['n_significant']} of {grid['n_paths']} "
      f"({grid['n_significant'] / grid['n_paths']:.1%})")
print(f"     estimate         : {grid['effect_min']:+.4f} to {grid['effect_max']:+.4f} uV, median "
      f"{grid['effect_median']:+.4f}, IQR {grid['effect_iqr']:.4f}")
print(f"     all same sign    : {grid['all_same_sign']}")
print(f"     p                : min {grid['p_min']:.3g}, median {grid['p_median']:.4f}, "
      f"max {grid['p_max']:.4f}")
print(f"     pre-specified    : {grid['effect_uv'][i_prespec]:+.4f} uV, p = {grid['p'][i_prespec]:.3g}")
print(f"     reference marginal (median estimate): "
      + " | ".join(f"{r['option']} {r['effect_median']:+.4f} uV {r['n_significant']}/{r['n_paths']}"
                   for r in rows if r["choice"] == "reference"))
print(f"     electrode marginal: "
      + " | ".join(f"{r['option']} {r['effect_median']:+.4f} uV {r['n_significant']}/{r['n_paths']}"
                   for r in rows if r["choice"] == "electrode"))
print()
print("3. PREREGISTRATION: 10 numbered items, SHA-256 " + digest)
print()
print(f"CROSS-CHECK against w-garden-of-forking-paths (paths.json, mode multiverse, 20 subjects):")
print(f"     widget:   476/486 significant; effect +0.40 to +6.05 uV, median +2.32; all positive;")
print(f"               reference marginal Cz +0.96, average +2.32, linked mastoids +5.26 uV.")
print(f"     notebook: {grid['n_significant']}/{grid['n_paths']} significant; effect "
      f"{grid['effect_min']:+.2f} to {grid['effect_max']:+.2f} uV, median {grid['effect_median']:+.2f}; "
      f"all positive = {grid['all_same_sign']};")
print(f"               reference marginal "
      + ", ".join(f"{r['option']} {r['effect_median']:+.2f}" for r in rows if r["choice"] == "reference")
      + " uV.")
print(f"     The notebook re-derives the grid from the raw files with its own code; agreement to the printed "
      f"precision means the two implementations of the same 486 pipelines agree.")
print()
print(f"Pitfall: pf-post-hoc-windows.  Widget: w-garden-of-forking-paths (mode multiverse).")
print(f"Statistics thread: L3.7 -> L4.7 -> L6.1 -> L6.4.")