# 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
ALPHA = 0.05
SEED = 20260918

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
i_ch = eeg.index(CH)

# subject-level difference waves: one per subject, which is the unit of analysis for every test below
X_all = np.stack([(store[s]["target"].copy().pick("eeg").data
                   - store[s]["standard"].copy().pick("eeg").data) * 1e6 for s in SUBJECTS])   # subj x ch x time
X = X_all[:, i_ch, :]                                                                          # subj x time
n = len(SUBJECTS)
epochs_sfreq = float(ep.info["sfreq"])
print(f"{n} subject-level difference waves (target minus standard), {X_all.shape[1]} channels, "
      f"{X_all.shape[2]} time points at {epochs_sfreq:.0f} Hz")
print(f"trials behind them: target "
      f"{[store[s]['info']['n_kept']['target'] for s in SUBJECTS]}, standard "
      f"{[store[s]['info']['n_kept']['standard'] for s in SUBJECTS]}")

from scipy import stats

amp = np.array([L3.mean_amplitude(X[k], None, W, times=times) for k in range(n)])
t_stat, p_val = stats.ttest_1samp(amp, 0)
dz = amp.mean() / amp.std(ddof=1)
ci = stats.t.ppf(1 - ALPHA / 2, n - 1) * amp.std(ddof=1) / np.sqrt(n)
w_stat, w_p = stats.wilcoxon(amp)

print(f"P3 mean amplitude at {CH}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms, target minus standard, {n} subjects")
print(f"  per subject: " + ", ".join(f"{v:+.2f}" for v in amp))
print(f"  mean {amp.mean():+.3f} uV, SD {amp.std(ddof=1):.3f}, SEM {amp.std(ddof=1) / np.sqrt(n):.3f}, "
      f"95% CI [{amp.mean() - ci:+.3f}, {amp.mean() + ci:+.3f}]")
print(f"  one-sample t({n - 1}) = {t_stat:.3f}, p = {p_val:.5f}, Cohen dz = {dz:.3f}")
print(f"  Wilcoxon signed rank (no normality assumption): W = {w_stat:.1f}, p = {w_p:.5f}")
print(f"  {int((amp > 0).sum())} of {n} subjects show a positive effect")

t_point, p_point = stats.ttest_1samp(X, 0, axis=0)
sig = p_point < ALPHA
pre = times < 0
post = times >= 0

print(f"uncorrected t tests at every time point ({len(times)} tests at alpha = {ALPHA})")
print(f"  significant anywhere      : {int(sig.sum()):3d} of {len(times)} ({100 * sig.mean():.1f} %)")
print(f"  significant before 0 ms   : {int((sig & pre).sum()):3d} of {int(pre.sum())} "
      f"({100 * sig[pre].mean():.1f} %) -- a LOWER bound, not a null rate: baseline correction forced the "
      f"mean of -200..0 ms to zero in every epoch")
print(f"  significant after 0 ms    : {int((sig & post).sum()):3d} of {int(post.sum())} "
      f"({100 * sig[post].mean():.1f} %)")
runs = [len(r) for r in "".join("1" if v else "0" for v in sig[pre]).split("0") if r]
print(f"  longest run of consecutive significant points before 0 ms: {max(runs or [0])} samples")
print(f"  if the {len(times)} tests were independent, alpha = {ALPHA} would give about "
      f"{ALPHA * len(times):.0f} false positives; at {epochs_sfreq:.0f} Hz they are nowhere near "
      f"independent, so the family-wise error rate is neither {100 * ALPHA:.0f} % nor "
      f"1 - {1 - ALPHA:g}^{len(times)} = {1 - (1 - ALPHA) ** len(times):.3f}; it is unknown, which is the "
      f"problem.")
print(f"  the {int(sig[post].sum())} significant points after 0 ms are also not {int(sig[post].sum())} "
      f"independent findings: they form "
      f"{len([r for r in ''.join('1' if v else '0' for v in sig[post]).split('0') if r])} contiguous "
      f"run(s), which is what the cluster test in section 4 is built to exploit.")

fig, axes = plt.subplots(2, 1, figsize=(10, 6.5), sharex=True)
L3.plot_erp({f"grand-average difference (n = {n} subjects)": (X.mean(0), {"color": "k", "lw": 1.8})},
            times, window=W, ax=axes[0], title=f"Subject-level difference waves at {CH} (uV)")
for k in range(n):
    axes[0].plot(times * 1000, X[k], lw=0.7, alpha=0.5)
axes[1].plot(times * 1000, t_point, "k", lw=1.2)
axes[1].fill_between(times * 1000, t_point, 0, where=sig, color="tab:red", alpha=0.35,
                     label=f"p < {ALPHA} uncorrected ({int(sig.sum())} points)")
for s_ in (1, -1):
    axes[1].axhline(s_ * stats.t.ppf(1 - ALPHA / 2, n - 1), color="tab:red", ls=":", lw=0.9)
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=f"t (one-sample, {n - 1} df)",
            title="The same data as 256 uncorrected tests")
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

threshold = float(stats.t.ppf(1 - ALPHA / 2, n - 1))     # two-tailed cluster-forming threshold at alpha = 0.05
N_PERM = 10000                                            # MNE reduces this to the exhaustive count when smaller

t_obs, clusters, cluster_p, H0 = mne.stats.permutation_cluster_1samp_test(
    X, threshold=threshold, n_permutations=N_PERM, tail=0, out_type="indices",
    seed=SEED, verbose=False)

print("call: mne.stats.permutation_cluster_1samp_test(")
print(f"          X,                      # {X.shape[0]} subjects x {X.shape[1]} time points, uV, {CH}")
print(f"          threshold={threshold:.4f},  # two-tailed t threshold at alpha = {ALPHA} with {n - 1} df")
print(f"          n_permutations={N_PERM}, tail=0, out_type='indices', seed={SEED})")
print(f"permutations actually used: {len(H0)}  (2^{n} = {2 ** n} sign flips exist; MNE enumerates them when "
      f"that is fewer than n_permutations)")
print(f"smallest p-value this test can return: {1 / len(H0):.5f}")
print()
order = np.argsort(cluster_p)
print(f"{'#':>2s} {'time range (ms)':>20s} {'samples':>8s} {'t-sum':>10s} {'max |t|':>8s} {'p':>8s} {'':>12s}")
cluster_rows = []
for rank, i in enumerate(order):
    idx = np.asarray(clusters[i][0])
    row = {"start_ms": float(times[idx[0]] * 1000), "end_ms": float(times[idx[-1]] * 1000),
           "n": int(len(idx)), "t_sum": float(t_obs[idx].sum()),
           "max_t": float(np.abs(t_obs[idx]).max()), "p": float(cluster_p[i]),
           "sign": "positive" if t_obs[idx].sum() > 0 else "negative"}
    cluster_rows.append(row)
    mark = "SIGNIFICANT" if row["p"] <= ALPHA else ""
    print(f"{rank:2d} {row['start_ms']:8.1f} to {row['end_ms']:8.1f} {row['n']:8d} {row['t_sum']:+10.2f} "
          f"{row['max_t']:8.2f} {row['p']:8.4f} {mark:>12s}")
sig_clusters = [r for r in cluster_rows if r["p"] <= ALPHA]
print(f"\n{len(sig_clusters)} of {len(cluster_rows)} clusters survive at alpha = {ALPHA}")

fig, axes = plt.subplots(2, 1, figsize=(10, 7))
L3.plot_cluster_test(times, t_obs, clusters, cluster_p, alpha=ALPHA, threshold=threshold, ax=axes[0],
                     title=f"Cluster permutation test at {CH}: observed t and the surviving clusters "
                           f"({n} subjects)")
axes[1].hist(H0, bins=50, color="tab:blue", alpha=0.8)
for r in sig_clusters:
    axes[1].axvline(r["t_sum"], color="tab:orange", lw=2,
                    label=f"observed t-sum {r['t_sum']:+.0f} (p = {r['p']:.4f})")
axes[1].set(xlabel="Largest cluster t-sum per permutation", ylabel="Permutations",
            title=f"The permutation distribution ({len(H0)} sign-flip permutations, seed {SEED})")
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"the observed t-sum sits at the {100 * np.mean(np.abs(H0) < abs(sig_clusters[0]['t_sum'])):.2f}nd "
      f"percentile of the permutation distribution" if sig_clusters else
      "no cluster survived, so there is nothing to locate in the permutation distribution")

print(f"cluster-forming threshold sweep ({CH}, {n} subjects, {len(H0)} permutations, seed {SEED})")
print(f"{'alpha_form':>11s} {'t threshold':>12s} {'clusters':>9s} {'best p':>8s} {'its range (ms)':>22s} "
      f"{'its t-sum':>10s}")
for a_form in (0.20, 0.10, 0.05, 0.01, 0.005):
    thr = float(stats.t.ppf(1 - a_form / 2, n - 1))
    to, cl, cp, h0 = mne.stats.permutation_cluster_1samp_test(
        X, threshold=thr, n_permutations=N_PERM, tail=0, out_type="indices", seed=SEED, verbose=False)
    if len(cp) == 0:
        print(f"{a_form:11.3f} {thr:12.3f} {0:9d} {'-':>8s} {'-':>22s} {'-':>10s}")
        continue
    b = int(np.argmin(cp))
    idx = np.asarray(cl[b][0])
    print(f"{a_form:11.3f} {thr:12.3f} {len(cp):9d} {cp[b]:8.4f} "
          f"{times[idx[0]] * 1000:9.1f} to {times[idx[-1]] * 1000:8.1f} {to[idx].sum():+10.2f}")

adjacency, ch_names_adj = mne.channels.find_ch_adjacency(proto.info, ch_type="eeg")
print(f"adjacency: {adjacency.shape[0]} channels, "
      f"{int(adjacency.sum() - adjacency.shape[0]) // 2} neighbour pairs "
      f"(mne.channels.find_ch_adjacency, Delaunay triangulation of the montage)")

X_st = np.transpose(X_all, (0, 2, 1))          # MNE wants subjects x times x channels
t_st, cl_st, p_st, H0_st = mne.stats.spatio_temporal_cluster_1samp_test(
    X_st, threshold=threshold, n_permutations=N_PERM, tail=0, adjacency=adjacency,
    out_type="mask", seed=SEED, verbose=False)

print(f"permutations used: {len(H0_st)}; smallest possible p-value {1 / len(H0_st):.5f}")
print(f"{'#':>2s} {'time range (ms)':>20s} {'channels':>9s} {'t-sum':>10s} {'p':>8s}  channels")
st_rows = []
for rank, i in enumerate(np.argsort(p_st)):
    mask = cl_st[i]
    ti = np.where(mask.any(axis=1))[0]
    chan_idx = np.where(mask.any(axis=0))[0]
    row = {"start_ms": float(times[ti[0]] * 1000), "end_ms": float(times[ti[-1]] * 1000),
           "n_ch": int(len(chan_idx)), "t_sum": float(t_st[mask].sum()), "p": float(p_st[i]),
           "channels": [eeg[j] for j in chan_idx]}
    st_rows.append(row)
    if row["p"] <= ALPHA or rank < 4:
        print(f"{rank:2d} {row['start_ms']:8.1f} to {row['end_ms']:8.1f} {row['n_ch']:9d} "
              f"{row['t_sum']:+10.2f} {row['p']:8.4f}  {', '.join(row['channels'][:10])}"
              + (" ..." if len(row["channels"]) > 10 else ""))
st_sig = [r for r in st_rows if r["p"] <= ALPHA]
print(f"\n{len(st_sig)} of {len(st_rows)} spatio-temporal clusters survive at alpha = {ALPHA}")

if st_sig:
    fig, axes = plt.subplots(1, len(st_sig) + 1, figsize=(3.4 * len(st_sig) + 1.2, 3.4),
                             gridspec_kw={"width_ratios": [1] * len(st_sig) + [0.09]})
    axes = np.atleast_1d(axes)
    v = float(np.abs(t_st).max())
    im = None
    for ax, (i, row) in zip(axes[:-1], zip(np.argsort(p_st), st_sig)):
        mask = cl_st[i]
        t_in_window = t_st[mask.any(axis=1)].mean(axis=0)     # mean t over the cluster's own time range
        im, _ = mne.viz.plot_topomap(t_in_window, proto.info, axes=ax, show=False, contours=4,
                                     vlim=(-v, v), sensors=True, mask=mask.any(axis=0),
                                     mask_params=dict(marker="o", markerfacecolor="k", markersize=5))
        ax.set_title(f"{row['start_ms']:.0f}-{row['end_ms']:.0f} ms\nt-sum {row['t_sum']:+.0f}, "
                     f"p = {row['p']:.4f}", fontsize=9)
    cb = fig.colorbar(im, cax=axes[-1])
    cb.set_label("mean t over the cluster's time range")
    fig.suptitle("Surviving spatio-temporal clusters: black dots are the channels in the cluster",
                 y=1.04, fontsize=10)
    fig.tight_layout()
    plt.show()   # render the static figure(s) of this cell inline

print("nb-3-7-cluster-test -- L3.7 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, sub-001 to sub-{SUBJECTS[-1]:03d} ({n} subjects, helpers_l3.SUBSET_DEFAULT); "
      f"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); unit of analysis = one difference wave "
      f"(target minus standard) per subject")
print()
print(f"1. PAIRED TEST ON ONE MEASURE ({CH}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms mean amplitude):")
print(f"     mean {amp.mean():+.3f} uV, 95% CI [{amp.mean() - ci:+.3f}, {amp.mean() + ci:+.3f}], "
      f"t({n - 1}) = {t_stat:.3f}, p = {p_val:.5f}, dz = {dz:.3f}")
print()
print(f"2. UNCORRECTED TIME-POINT TESTS: {int(sig.sum())} of {len(times)} points reach p < {ALPHA} "
      f"({100 * sig.mean():.1f} %), of which {int((sig & pre).sum())} lie in the pre-stimulus baseline where "
      f"the epochs were zeroed and nothing can be happening.")
print()
print(f"3. CLUSTER PERMUTATION TEST at {CH} -- exact call:")
print(f"     mne.stats.permutation_cluster_1samp_test(X, threshold={threshold:.4f}, "
      f"n_permutations={N_PERM}, tail=0, out_type='indices', seed={SEED})")
print(f"     X is {X.shape[0]} subjects x {X.shape[1]} time points of the {CH} difference wave in uV; "
      f"threshold = two-tailed t at alpha = {ALPHA} with {n - 1} df")
print(f"     permutations used {len(H0)} (exhaustive: 2^{n} = {2 ** n} sign flips), so the smallest "
      f"attainable p-value is {1 / len(H0):.5f}")
print(f"     CLUSTERS ({len(cluster_rows)} formed, {len(sig_clusters)} significant at alpha = {ALPHA}):")
for r in cluster_rows:
    print(f"       {r['sign']:8s} cluster {r['start_ms']:7.1f} to {r['end_ms']:7.1f} ms "
          f"({r['n']:3d} samples): t-sum {r['t_sum']:+9.2f}, max |t| {r['max_t']:.2f}, p = {r['p']:.4f}"
          + ("   <- significant" if r["p"] <= ALPHA else ""))
print()
print(f"4. SPATIO-TEMPORAL CLUSTER TEST over {len(eeg)} channels and {len(times)} time points "
      f"(mne.stats.spatio_temporal_cluster_1samp_test, same threshold, seed and permutation count, "
      f"adjacency from mne.channels.find_ch_adjacency):")
for r in st_rows[:4]:
    print(f"       {r['start_ms']:7.1f} to {r['end_ms']:7.1f} ms, {r['n_ch']:2d} channels: "
          f"t-sum {r['t_sum']:+9.1f}, p = {r['p']:.4f}"
          + ("   <- significant" if r["p"] <= ALPHA else ""))
    print(f"           channels: {', '.join(r['channels'])}")
print()
print("ANSWER KEY -- ex-3-7 (multiple select), what the result licenses:")
if sig_clusters:
    best = sig_clusters[0]
    print(f"  LICENSED: 'target and standard differ somewhere in the tested window and channel set' "
          f"(cluster p = {best['p']:.4f}, family-wise error rate controlled at alpha = {ALPHA} over the "
          f"whole tested space).")
    print(f"  LICENSED: reporting the effect size separately -- {amp.mean():+.3f} uV "
          f"(95% CI [{amp.mean() - ci:+.3f}, {amp.mean() + ci:+.3f}], dz = {dz:.3f}) from the a-priori "
          f"window, which the cluster test does not provide.")
    print(f"  NOT LICENSED: 'the effect starts at {best['start_ms']:.0f} ms and ends at "
          f"{best['end_ms']:.0f} ms'. Those boundaries are where |t| crossed an arbitrary threshold; the "
          f"sweep in section 4 moves them by changing only that threshold.")
    print(f"  NOT LICENSED: 'the effect is significant at {CH} and not elsewhere' -- the {CH}-only test "
          f"never looked elsewhere, and the spatio-temporal test found "
          f"{len(st_sig)} cluster(s) spanning up to {max((r['n_ch'] for r in st_sig), default=0)} channels.")
    print(f"  NOT LICENSED: comparing the t-sums of two clusters, or calling a cluster with p just above "
          f"{ALPHA} 'a trend'. The t-sum grows with cluster length and is not an effect size.")
else:
    print("  no cluster survived; the licensed statement is that this test did not reject the null, which "
          "is not evidence that the conditions are the same.")
print()
print(f"Pitfalls: pf-uncorrected-timepoint-tests, pf-cluster-inference-misread.  "
      f"Widget: w-cluster-permutation-viz (mode erp).")