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

SUBJECT = 1                      # the detail subject: the first of the documented subset, chosen before looking
WIDE = (-1.0, 1.0)               # wide enough to hold every baseline window this notebook tries

epochs, info = L3.load_p3_epochs(SUBJECT, tmin=WIDE[0], tmax=WIDE[1], baseline=None, verbose=True)
print()
print(epochs)
print()
print("metadata columns:", list(epochs.metadata.columns))
print(epochs.metadata.head(6).to_string(index=False))
print()
print(f"{info['subject']}: {info['n_stimulus_events']} stimulus events, {info['n_epochs']} epochs survived the "
      f"{WIDE[0]:+g} to {WIDE[1]:+g} s window, {info['n_rejected']} rejected at "
      f"{info['reject_ptp_uv']:g} uV peak-to-peak")
print(f"  kept: {info['n_kept']['target']} target, {info['n_kept']['standard']} standard "
      f"(design ratio 1:4, p(target) = .2)")
print(f"  bad channels flagged and interpolated: {info['bad_channels'] or 'none'}; "
      f"ICA components removed as ocular: {info['ica_excluded']}")
print(f"  mean reaction time: target {info['mean_rt_ms']['target']:.0f} ms, "
      f"standard {info['mean_rt_ms']['standard']:.0f} ms; accuracy {100 * info['accuracy']:.1f} %")

soa = epochs.metadata["soa_s"].to_numpy(float)
soa = soa[np.isfinite(soa)]
print(f"stimulus-onset asynchrony: median {np.median(soa):.3f} s, "
      f"range {soa.min():.3f}-{soa.max():.3f} s ({len(soa)} intervals)")
print(f"  an epoch reaching {WIDE[1]:+g} s therefore overlaps the next stimulus on "
      f"{100 * (soa < WIDE[1]).mean():.0f} % of trials, and an epoch starting at {WIDE[0]:+g} s overlaps the "
      f"previous one on {100 * (soa < -WIDE[0]).mean():.0f} %")
print(f"  the pipeline's analysis window is {L3.EPOCH_TMIN:+g} to {L3.EPOCH_TMAX:+g} s, inside the shortest "
      f"observed asynchrony ({soa.min():.3f} s), so no analysed epoch contains a second stimulus")

fig, ax = plt.subplots(figsize=(8, 3.2))
ax.hist(soa * 1000, bins=40, color="tab:blue", alpha=0.8)
for edge, label in ((L3.EPOCH_TMAX * 1000, "analysis tmax"), (-L3.EPOCH_TMIN * 1000, "|analysis tmin|")):
    ax.axvline(edge, color="tab:orange", lw=1.2, ls="--", label=f"{label} = {edge:.0f} ms")
ax.set(xlabel="Stimulus-onset asynchrony (ms)", ylabel="Trials",
       title=f"{info['subject']}: stimulus-onset asynchrony and the epoch window (ms)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

BASELINES = [
    ("none", None),
    ("-200 to 0 ms (pipeline default)", (-0.2, 0.0)),
    ("-600 to -400 ms", (-0.6, -0.4)),
    ("150 to 300 ms (the component's rising edge)", (0.15, 0.30)),
    ("300 to 600 ms (the measurement window itself)", (0.3, 0.6)),
]
W = L3.P3_WINDOW
times = epochs.times
eeg = info["eeg_channels"]
i_pz = eeg.index(L3.P3_CHANNEL)

tgt = L3.condition_epochs(epochs, "target").get_data(picks="eeg") * 1e6      # trials x 30 x times, uV
std = L3.condition_epochs(epochs, "standard").get_data(picks="eeg") * 1e6


def rebaseline(x, window):
    """Subtract each epoch's mean over `window` (None = leave the epoch alone)."""
    if window is None:
        return x
    m = (times >= window[0]) & (times <= window[1])
    return x - x[..., m].mean(axis=-1, keepdims=True)


rows = []
waves = {}
for label, win in BASELINES:
    t_avg, s_avg = rebaseline(tgt, win).mean(0), rebaseline(std, win).mean(0)
    diff = t_avg - s_avg
    waves[label] = diff[i_pz]
    rows.append((label,
                 L3.mean_amplitude(t_avg[i_pz], None, W, times=times),
                 L3.mean_amplitude(s_avg[i_pz], None, W, times=times),
                 L3.mean_amplitude(diff[i_pz], None, W, times=times)))

print(f"{info['subject']}, {L3.P3_CHANNEL}, mean amplitude over "
      f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms (uV)")
print(f"{'baseline window':48s} {'target':>8s} {'standard':>9s} {'T - S':>8s}")
for label, a, b, c in rows:
    print(f"{label:48s} {a:+8.3f} {b:+9.3f} {c:+8.3f}")

fig, axes = plt.subplots(1, 2, figsize=(13, 4.4))
L3.plot_erp({label: waves[label] for label, _ in BASELINES}, times, window=W, ax=axes[0],
            title=f"{info['subject']}: target - standard at {L3.P3_CHANNEL} under four baselines (uV)")
for label, win in BASELINES:
    if win is not None:
        axes[0].axvspan(win[0] * 1000, win[1] * 1000, color="tab:green", alpha=0.07, lw=0)
b_tgt = rebaseline(tgt, (-0.2, 0.0)).mean(0)
b_std = rebaseline(std, (-0.2, 0.0)).mean(0)
L3.plot_erp({f"target (n = {tgt.shape[0]})": b_tgt[i_pz],
             f"standard (n = {std.shape[0]})": b_std[i_pz],
             "difference": (b_tgt - b_std)[i_pz]}, times, window=W, ax=axes[1],
            title=f"{info['subject']}: the two conditions, baseline -200 to 0 ms (uV)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

SUBJECTS = list(L3.SUBSET_DEFAULT)      # the documented subset: helpers_l3.SUBSET_DEFAULT
print(f"documented subset: {len(SUBJECTS)} subjects, sub-001 to sub-{SUBJECTS[-1]:03d} "
      f"(the first N of the 40 the dataset ships; no subject was chosen by its result)")

store = {}
for s in SUBJECTS:
    ep, nfo = L3.load_p3_epochs(s, tmin=WIDE[0], tmax=WIDE[1], baseline=None, verbose=False)
    store[s] = {
        "target": L3.condition_epochs(ep, "target").get_data(picks="eeg") * 1e6,
        "standard": L3.condition_epochs(ep, "standard").get_data(picks="eeg") * 1e6,
        "info": nfo,
    }
    print(f"  {nfo['subject']}: {nfo['n_kept']['target']:3d} target, {nfo['n_kept']['standard']:3d} standard "
          f"kept; {nfo['n_rejected']:3d} rejected; bads {nfo['bad_channels'] or '[]'}; "
          f"ICA excluded {nfo['ica_excluded']}")

PRE_WINDOWS = [(-1.0, -0.8), (-0.8, -0.6), (-0.6, -0.4), (-0.4, -0.2), (-0.2, 0.0)]
print(f"target - standard at {L3.P3_CHANNEL} in pre-stimulus windows, unbaselined, "
      f"mean over {len(SUBJECTS)} subjects (uV):")
pre_means = {}
for win in PRE_WINDOWS:
    vals = np.array([L3.mean_amplitude((store[s]["target"].mean(0) - store[s]["standard"].mean(0))[i_pz],
                                       None, win, times=times) for s in SUBJECTS])
    pre_means[win] = vals
    print(f"  {win[0] * 1000:+6.0f} to {win[1] * 1000:+6.0f} ms : {vals.mean():+6.3f} uV  "
          f"(SD {vals.std(ddof=1):5.3f}, {int((vals > 0).sum())}/{len(vals)} subjects positive)")
post = np.array([L3.mean_amplitude((store[s]["target"].mean(0) - store[s]["standard"].mean(0))[i_pz],
                                   None, W, times=times) for s in SUBJECTS])
print(f"  {W[0] * 1000:+6.0f} to {W[1] * 1000:+6.0f} ms : {post.mean():+6.3f} uV   <- the measurement window")

def measure(window_baseline):
    out = []
    for s in SUBJECTS:
        t_avg = rebaseline(store[s]["target"], window_baseline).mean(0)
        s_avg = rebaseline(store[s]["standard"], window_baseline).mean(0)
        out.append(L3.mean_amplitude((t_avg - s_avg)[i_pz], None, W, times=times))
    return np.array(out)


per_baseline = {label: measure(win) for label, win in BASELINES}
default = per_baseline["-200 to 0 ms (pipeline default)"]
contaminated = per_baseline["-600 to -400 ms"]
shift = contaminated - default

print(f"P3 (target - standard, {L3.P3_CHANNEL}, {W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms mean amplitude), "
      f"{len(SUBJECTS)} subjects")
for label, vals in per_baseline.items():
    print(f"  baseline {label:48s} {vals.mean():+7.3f} uV   (SD {vals.std(ddof=1):5.3f}, "
          f"SEM {vals.std(ddof=1) / np.sqrt(len(vals)):5.3f})")
print()
print(f"moving the baseline from -200..0 ms to -600..-400 ms changes the P3 by "
      f"{shift.mean():+.3f} uV (SD {shift.std(ddof=1):.3f}), i.e. "
      f"{100 * shift.mean() / default.mean():+.1f} % of the P3 measured with the default baseline")
print(f"  per subject: " + ", ".join(f"{v:+.2f}" for v in shift))
predicted = pre_means[(-0.2, 0.0)].mean() - pre_means[(-0.6, -0.4)].mean()
print(f"  the arithmetic closes: the two baseline windows differ in their own condition difference by "
      f"{pre_means[(-0.2, 0.0)].mean():+.3f} - {pre_means[(-0.6, -0.4)].mean():+.3f} = {predicted:+.3f} uV, "
      f"and the P3 moves by {shift.mean():+.3f} uV.  A baseline is subtracted one for one.")

fig, ax = plt.subplots(figsize=(8.5, 4))
x = np.arange(len(SUBJECTS))
ax.bar(x - 0.2, default, 0.4, label="baseline -200 to 0 ms")
ax.bar(x + 0.2, contaminated, 0.4, label="baseline -600 to -400 ms")
ax.set_xticks(x)
ax.set_xticklabels([f"{s:03d}" for s in SUBJECTS], fontsize=8)
ax.axhline(0, color="gray", lw=0.6)
ax.set(xlabel="Subject", ylabel="Amplitude (uV)",
       title=f"P3 per subject under two baselines ({L3.P3_CHANNEL}, "
             f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms mean, target - standard, uV)")
ax.grid(alpha=0.3, axis="y")
ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

b_pre = []
for s in SUBJECTS:
    x = np.concatenate([store[s]["target"][:, i_pz, :], store[s]["standard"][:, i_pz, :]])
    cond = np.r_[np.ones(store[s]["target"].shape[0]), np.zeros(store[s]["standard"].shape[0])]
    m_pre = (times >= -0.2) & (times <= 0.0)
    m_win = (times >= W[0]) & (times <= W[1])
    pre, post_ = x[:, m_pre].mean(1), x[:, m_win].mean(1)
    design = np.c_[np.ones_like(cond), cond, pre - pre.mean()]
    beta, *_ = np.linalg.lstsq(design, post_, rcond=None)
    b_pre.append(beta[2])
b_pre = np.array(b_pre)
print(f"coefficient on the pre-stimulus mean, single-trial regression per subject "
      f"(amplitude ~ 1 + condition + pre-stimulus mean):")
print(f"  mean {b_pre.mean():+.3f}, SD {b_pre.std(ddof=1):.3f}, range {b_pre.min():+.3f} to {b_pre.max():+.3f}")
print(f"  mean subtraction assumes this coefficient is exactly +1.000; the estimate is "
      f"{b_pre.mean():+.3f}, so subtraction removes "
      f"{'more' if b_pre.mean() < 1 else 'less'} than the data support here (TODO(confirm) at author review)")

print("nb-3-1-epochs -- L3.1 exercise numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, subjects sub-001 to sub-{SUBJECTS[-1]:03d} ({len(SUBJECTS)} subjects, the "
      "documented subset 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); epochs {WIDE[0]:+g} to {WIDE[1]:+g} s, "
      f"baseline applied afterwards per variant; {L3.REJECT_PTP_UV:g} uV peak-to-peak rejection")
print(f"Measure: mean amplitude of the target-minus-standard difference wave, "
      f"{W[0] * 1000:.0f}-{W[1] * 1000:.0f} ms (a-priori window), channel {L3.P3_CHANNEL}")
print()
for label, vals in per_baseline.items():
    print(f"  baseline {label:48s} P3 = {vals.mean():+6.3f} uV  (SD {vals.std(ddof=1):5.3f})")
print()
print(f"ANSWER KEY -- ex-3-1 (numeric): moving the baseline window from -200..0 ms to -600..-400 ms, "
      f"a window that does contain a condition difference, changes the P3 from "
      f"{default.mean():+.3f} uV to {contaminated.mean():+.3f} uV, a change of {shift.mean():+.3f} uV "
      f"({100 * shift.mean() / default.mean():+.1f} %).")
print(f"ANSWER KEY -- ex-3-1 (free response): the pre-stimulus target-minus-standard difference at "
      f"{L3.P3_CHANNEL} is {pre_means[(-0.6, -0.4)].mean():+.3f} uV over -600..-400 ms and "
      f"{pre_means[(-0.2, 0.0)].mean():+.3f} uV over -200..0 ms, so a baseline taken there is subtracted "
      f"from the effect one for one; candidates for the pre-stimulus difference are trial overlap at a "
      f"{np.median(soa):.2f} s stimulus-onset asynchrony and the symmetric smearing of a zero-phase high-pass "
      f"(TODO(confirm)).")
print(f"ANSWER KEY -- degenerate case: a baseline equal to the measurement window gives exactly "
      f"{per_baseline['300 to 600 ms (the measurement window itself)'].mean():+.3f} uV by construction.")
print(f"Supporting: no baseline at all gives {per_baseline['none'].mean():+.3f} uV "
      f"(SD {per_baseline['none'].std(ddof=1):.3f}) against {default.mean():+.3f} uV "
      f"(SD {default.std(ddof=1):.3f}) with the default baseline -- on these already high-passed, "
      f"average-referenced epochs the short baseline window adds across-subject variance rather than "
      f"removing it.")
print(f"Pitfall: pf-baseline-contamination.  Widget: w-epoch-builder.")