# 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")

DETAIL = 1                       # the detail subject: the first of the documented subset, fixed before looking
SUBJECTS = list(L3.SUBSET_DEFAULT)
W = L3.P3_WINDOW
CH = L3.P3_CHANNEL
LAT_WINDOW = (0.20, 0.80)        # a-priori window for a SINGLE-TRIAL latency: wider than the measurement window,
                                 # because single trials scatter in latency far more than the average does
SMOOTH_HZ = 10.0                 # single trials are low-passed before any latency is read off them

epochs, info = L3.load_p3_epochs(DETAIL, verbose=True)
times = epochs.times
eeg = info["eeg_channels"]
i_ch = eeg.index(CH)
tgt = L3.condition_epochs(epochs, "target")
std = L3.condition_epochs(epochs, "standard")
x = tgt.get_data(picks="eeg")[:, i_ch, :] * 1e6           # target single trials at Pz, uV
rt = tgt.metadata["rt_ms"].to_numpy(float)
acc = tgt.metadata["correct"].to_numpy(float)

print()
print(f"{info['subject']}: {x.shape[0]} target trials kept at {CH} ({info['n_rejected']} of "
      f"{info['n_stimulus_events']} epochs rejected at {L3.REJECT_PTP_UV:g} uV peak-to-peak)")
print(f"  reaction time: {np.sum(np.isfinite(rt))} trials have one; median {np.nanmedian(rt):.0f} ms, "
      f"range {np.nanmin(rt):.0f}-{np.nanmax(rt):.0f} ms, IQR "
      f"{np.nanpercentile(rt, 25):.0f}-{np.nanpercentile(rt, 75):.0f} ms")
print(f"  accuracy on those trials: {100 * np.nanmean(acc):.1f} %")

ok = np.isfinite(rt)
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
L3.plot_erp_image(x[ok], times, title=f"{info['subject']} {CH}: target trials, unsorted (uV)", ax=axes[0])
L3.plot_erp_image(x[ok], times, sort_values=rt[ok], sort_label="reaction time", overlay=rt[ok] / 1000,
                  title=f"{info['subject']} {CH}: sorted by reaction time (uV)", ax=axes[1])
L3.plot_erp_image(x[ok], times, sort_values=rt[ok], sort_label="reaction time", smooth=5,
                  overlay=rt[ok] / 1000,
                  title=f"{info['subject']} {CH}: sorted by RT, 5-trial moving average (uV)", ax=axes[2])
for ax in axes:
    ax.axvspan(W[0] * 1000, W[1] * 1000, color="none", ec="k", lw=0.8, ls="--")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

fast = rt[ok] <= np.nanmedian(rt[ok])
fig, ax = plt.subplots(figsize=(9, 4.2))
L3.plot_erp({f"fast half (n = {int(fast.sum())}, median RT {np.median(rt[ok][fast]):.0f} ms)": x[ok][fast].mean(0),
             f"slow half (n = {int((~fast).sum())}, median RT {np.median(rt[ok][~fast]):.0f} ms)":
                 x[ok][~fast].mean(0),
             "all target trials": (x[ok].mean(0), {"color": "k", "lw": 1.8})},
            times, window=W, ax=ax,
            title=f"{info['subject']} {CH}: target average split at the median reaction time (uV)")
for half, c in ((fast, "tab:blue"), (~fast, "tab:orange")):
    ax.axvline(np.median(rt[ok][half]), color=c, lw=1.0, ls=":")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

from scipy import signal, stats

sos = signal.butter(4, SMOOTH_HZ, btype="lowpass", fs=float(epochs.info["sfreq"]), output="sos")


def single_trial_latency(trials, smooth=True):
    y = signal.sosfiltfilt(sos, trials, axis=1) if smooth else trials
    return np.array([L3.fractional_area_latency(y[i], None, LAT_WINDOW, times=times) for i in range(y.shape[0])])


lat = single_trial_latency(x)
lat_raw = single_trial_latency(x, smooth=False)
amp = x[:, (times >= W[0]) & (times <= W[1])].mean(1)

good = np.isfinite(lat) & np.isfinite(rt)
r, p = stats.pearsonr(lat[good] * 1000, rt[good])
rho, p_rho = stats.spearmanr(lat[good] * 1000, rt[good])
r_raw, p_raw = stats.pearsonr(lat_raw[good] * 1000, rt[good])
r_amp, p_amp = stats.pearsonr(amp[good], rt[good])

print(f"{info['subject']}, {CH}, {int(good.sum())} target trials with a reaction time")
print(f"  single-trial 50% fractional-area latency over "
      f"{LAT_WINDOW[0] * 1000:.0f}-{LAT_WINDOW[1] * 1000:.0f} ms, trials low-passed at {SMOOTH_HZ:g} Hz:")
print(f"      median {1000 * np.nanmedian(lat[good]):.0f} ms, SD {1000 * np.nanstd(lat[good], ddof=1):.0f} ms")
print(f"  correlation with reaction time: Pearson r = {r:+.3f}, p = {p:.4f}; "
      f"Spearman rho = {rho:+.3f}, p = {p_rho:.4f}")
print(f"  without the {SMOOTH_HZ:g} Hz smoothing:  Pearson r = {r_raw:+.3f}, p = {p_raw:.4f}")
print(f"  single-trial mean AMPLITUDE against reaction time: r = {r_amp:+.3f}, p = {p_amp:.4f}")

fig, axes = plt.subplots(1, 2, figsize=(12, 4.4))
axes[0].plot(rt[good], lat[good] * 1000, "o", ms=5, alpha=0.75)
b1, b0 = np.polyfit(rt[good], lat[good] * 1000, 1)
xs = np.linspace(np.nanmin(rt[good]), np.nanmax(rt[good]), 10)
axes[0].plot(xs, b0 + b1 * xs, "k-", lw=1.4,
             label=f"slope {b1:+.3f} ms of latency per ms of RT\nr = {r:+.3f}, p = {p:.4f}")
axes[0].set(xlabel="Reaction time (ms)", ylabel=f"Single-trial 50% area latency at {CH} (ms)",
            title=f"{info['subject']}: does P3 latency track reaction time?")
axes[0].grid(alpha=0.3)
axes[0].legend(fontsize=8)
axes[1].plot(rt[good], amp[good], "o", ms=5, alpha=0.75, color="tab:orange")
b1a, b0a = np.polyfit(rt[good], amp[good], 1)
axes[1].plot(xs, b0a + b1a * xs, "k-", lw=1.4,
             label=f"slope {b1a:+.4f} uV per ms of RT\nr = {r_amp:+.3f}, p = {p_amp:.4f}")
axes[1].set(xlabel="Reaction time (ms)",
            ylabel=f"Single-trial mean amplitude at {CH}, "
                   f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms (uV)",
            title=f"{info['subject']}: does P3 amplitude track reaction time?")
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

import pandas as pd

all_ep = epochs[np.where(np.isfinite(epochs.metadata["rt_ms"].to_numpy(float)))[0]].copy().pick("eeg")
md_ = all_ep.metadata
rt_all = md_["rt_ms"].to_numpy(float)
is_target = (md_["condition"].to_numpy() == "target").astype(float)
rt_centred = (rt_all - rt_all.mean()) / 100.0        # units of 100 ms, mean-centred

design = pd.DataFrame({"intercept": np.ones(len(md_)),
                       "target": is_target - is_target.mean(),
                       "rt_100ms": rt_centred})
res = mne.stats.linear_regression(all_ep, design, names=list(design.columns))

print(f"regression ERP on {len(md_)} trials of {info['subject']} "
      f"(both conditions, only trials with a reaction time)")
print(f"  design: {list(design.columns)}; 'target' and 'rt_100ms' are mean-centred, so the intercept is the "
      f"average trial")
print(f"  correlation between the two centred regressors: "
      f"{np.corrcoef(design['target'], design['rt_100ms'])[0, 1]:+.3f} "
      f"(targets are responded to more slowly here, so the two are not independent -- which is exactly why "
      f"they belong in the same model)")

fig, axes = plt.subplots(1, 2, figsize=(13, 4.4))
for name, ax in (("target", axes[0]), ("rt_100ms", axes[1])):
    beta = res[name].beta.data[all_ep.ch_names.index(CH)] * 1e6
    tval = res[name].t_val.data[all_ep.ch_names.index(CH)]
    ax.plot(times * 1000, beta, "k", lw=1.6, label="beta")
    ax.axvspan(W[0] * 1000, W[1] * 1000, color="tab:orange", alpha=0.18, lw=0, label="measurement window")
    ax.axhline(0, color="gray", lw=0.6)
    ax.axvline(0, color="gray", lw=0.6)
    sig = np.abs(tval) > 2.0
    ax.fill_between(times * 1000, beta, 0, where=sig, color="tab:blue", alpha=0.3,
                    label="|t| > 2 (uncorrected, see nb-3-7)")
    unit = "uV (target minus standard)" if name == "target" else "uV per 100 ms of reaction time"
    ax.set(xlabel="Time from stimulus (ms)", ylabel=unit,
           title=f"{info['subject']} {CH}: regression ERP, beta for '{name}'")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

for name in ("target", "rt_100ms"):
    b = res[name].beta.data[all_ep.ch_names.index(CH)] * 1e6
    t_ = res[name].t_val.data[all_ep.ch_names.index(CH)]
    m = (times >= W[0]) & (times <= W[1])
    j = int(np.argmax(np.abs(b[m])))
    unit = "uV" if name == "target" else "uV per 100 ms RT"
    print(f"  beta '{name}': window mean {b[m].mean():+.4f} {unit}; largest |beta| in the window "
          f"{b[m][j]:+.4f} {unit} at {1000 * times[m][j]:.0f} ms (t = {t_[m][j]:+.2f})")

rows = []
for s in SUBJECTS:
    ep_s, nfo = L3.load_p3_epochs(s, verbose=False)
    t_s = L3.condition_epochs(ep_s, "target")
    xs = t_s.get_data(picks="eeg")[:, i_ch, :] * 1e6
    rts = t_s.metadata["rt_ms"].to_numpy(float)
    ls = single_trial_latency(xs)
    g = np.isfinite(ls) & np.isfinite(rts)
    if g.sum() < 10:
        rows.append((nfo["subject"], int(g.sum()), np.nan, np.nan, np.nan, np.nan))
        continue
    rr, pp = stats.pearsonr(ls[g] * 1000, rts[g])
    ra, pa = stats.pearsonr(xs[:, (times >= W[0]) & (times <= W[1])].mean(1)[g], rts[g])
    rows.append((nfo["subject"], int(g.sum()), rr, pp, ra, pa))

print(f"single-trial P3 latency and amplitude against reaction time, {CH}, per subject")
print(f"{'subject':9s} {'n':>4s} {'r(latency, RT)':>15s} {'p':>9s} {'r(amplitude, RT)':>17s} {'p':>9s}")
for r_ in rows:
    print(f"{r_[0]:9s} {r_[1]:4d} {r_[2]:15.3f} {r_[3]:9.4f} {r_[4]:17.3f} {r_[5]:9.4f}")
rs = np.array([r_[2] for r_ in rows], float)
ps = np.array([r_[3] for r_ in rows], float)
t_group, p_group = stats.ttest_1samp(rs[np.isfinite(rs)], 0)
print(f"\nacross the {len(SUBJECTS)} subjects: mean r = {np.nanmean(rs):+.3f} "
      f"(SD {np.nanstd(rs, ddof=1):.3f}), "
      f"{int(np.sum(ps < 0.05))} subjects reach p < 0.05 individually "
      f"({int(np.sum((ps < 0.05) & (rs > 0)))} positive, {int(np.sum((ps < 0.05) & (rs < 0)))} negative)")
print(f"  one-sample t test on the {len(rs[np.isfinite(rs)])} correlations: t = {t_group:.3f}, "
      f"p = {p_group:.4f}")

print("nb-3-6-erp-image -- L3.6 exercise numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3; detail subject {info['subject']} (the first of the documented subset, fixed "
      f"before looking at any result); group table over sub-001 to sub-{SUBJECTS[-1]:03d} "
      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)")
print(f"Single-trial latency: 50% fractional-area latency over "
      f"{LAT_WINDOW[0] * 1000:.0f}-{LAT_WINDOW[1] * 1000:.0f} ms at {CH}, trials low-passed at "
      f"{SMOOTH_HZ:g} Hz; both choices fixed before the correlation was computed")
print(f"Reaction time: first response event after the stimulus and before the next one, from the subject's "
      f"own events.tsv")
print()
print(f"ANSWER KEY -- ex-3-6 (numeric): in {info['subject']}, over {int(good.sum())} target trials, "
      f"single-trial P3 latency correlates with reaction time at")
print(f"    Pearson r = {r:+.3f}, p = {p:.4f}   (Spearman rho = {rho:+.3f}, p = {p_rho:.4f})")
print(f"    slope {b1:+.3f} ms of P3 latency per ms of reaction time")
verdict = ("does NOT track" if p >= 0.05 else ("tracks" if r > 0 else "tracks INVERSELY"))
print(f"ANSWER KEY -- ex-3-6 (free response): in this subject P3 latency {verdict} reaction time "
      f"(p = {p:.4f} against alpha = 0.05). Single-trial amplitude against reaction time in the same "
      f"subject: r = {r_amp:+.3f}, p = {p_amp:.4f}.")
print(f"    The group table is the context the single-subject answer needs: mean r = {np.nanmean(rs):+.3f} "
      f"across {len(SUBJECTS)} subjects, {int(np.sum(ps < 0.05))} of them individually significant, "
      f"one-sample t({len(rs[np.isfinite(rs)]) - 1}) = {t_group:.3f}, p = {p_group:.4f}. A single subject's "
      f"correlation on ~{int(good.sum())} trials has a wide confidence interval; the honest answer names the "
      f"trial count.")
print(f"    Without the {SMOOTH_HZ:g} Hz single-trial smoothing the same measurement gives r = {r_raw:+.3f}, "
      f"p = {p_raw:.4f}, so the conclusion "
      f"{'does not depend' if (p < 0.05) == (p_raw < 0.05) else 'DOES depend'} on that choice.")
print()
print(f"Supporting -- regression ERP ({CH}, {len(md_)} trials, both conditions, reaction time mean-centred "
      f"in units of 100 ms):")
for name in ("target", "rt_100ms"):
    b = res[name].beta.data[all_ep.ch_names.index(CH)] * 1e6
    mm = (times >= W[0]) & (times <= W[1])
    unit = "uV" if name == "target" else "uV per 100 ms of reaction time"
    print(f"    beta '{name}' window mean {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms = {b[mm].mean():+.4f} {unit}")
print(f"Pitfalls: none named for L3.6 in spec section 6.  Widget: none.")