# 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", "specparam", "pandas")
_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"]
    if "specparam" in _missing:
        _cmd += ["specparam==2.0.0rc4"]
    if "pandas" in _missing:
        _cmd += ["pandas>=2"]
    subprocess.check_call(_cmd)

# 2. Shared helpers (notebooks/_shared/helpers.py and helpers_l1.py), 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_l1.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/capstones/ (or notebooks/) so that _shared/helpers_l1.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l1

# 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 plt.show() renders each cell's
#    figures in place.
import matplotlib.pyplot as plt
import numpy as np
import mne
import pooch
import pandas as pd
import specparam
from scipy import stats

mne.viz.set_browser_backend("matplotlib", verbose=False)
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__}, specparam {specparam.__version__}; helpers and helpers_l1 imported from notebooks/_shared; downloads go to the course "
      "data directory (EEG_COURSE_DOWNLOADS or EEG_COURSE_DATA if set, else MNE's data directory under eeg-course/)")

FULL_COHORT = False                 # True: all 109 subjects minus the exclusions (local runs; ~250 MB of downloads)
SUBSET_LAST = 20                    # the documented subset is S001..S020
EXCLUDE_OPTIONAL = True             # also refuse S038 and S104 (catalog: 'several reports additionally drop')
SUBJECTS = helpers_l1.eegbci_subjects(1, 109 if FULL_COHORT else SUBSET_LAST, exclude_optional=EXCLUDE_OPTIONAL)
RUNS = {"EO": "R01", "EC": "R02"}   # eyes open, eyes closed (catalog run table)
PICKS = ["O1", "Oz", "O2"]          # occipital average
WELCH = dict(seg_s=2.0, overlap=0.5, window="hann")
FIT_RANGE = helpers_l1.FIT_RANGE    # (1, 40) Hz
SPECPARAM = helpers_l1.SPECPARAM_SETTINGS
IAF_RANGE = helpers_l1.IAF_RANGE    # (7, 13) Hz
ALPHA_BAND, BROAD_BAND = (8.0, 12.0), (1.0, 40.0)

excluded = list(helpers.EEGBCI_EXCLUDE_DEFAULT) + (list(helpers.EEGBCI_EXCLUDE_OPTIONAL) if EXCLUDE_OPTIONAL else [])
print(f"cohort: {'FULL (S001-S109 minus exclusions)' if FULL_COHORT else f'subset S001-S{SUBSET_LAST:03d}'} -> {len(SUBJECTS)} subjects")
print(f"excluded before loading: {', '.join(helpers.EEGBCI_EXCLUDE_DEFAULT)} (catalog: inconsistent/overlapping event timestamps)"
      + (f"; {', '.join(helpers.EEGBCI_EXCLUDE_OPTIONAL)} (catalog: several reports additionally drop them)" if EXCLUDE_OPTIONAL else ""))
print(f"  of which inside the requested range: {[s for s in excluded if int(s[1:]) <= (109 if FULL_COHORT else SUBSET_LAST)] or 'none'}")
print(f"Welch: {WELCH['seg_s']:g}-s {WELCH['window']} segments, {int(100 * WELCH['overlap'])} % overlap, resolution {1 / WELCH['seg_s']:g} Hz, "
      f"{helpers_l1.n_welch_segments(int(61 * 160), int(WELCH['seg_s'] * 160), int(WELCH['seg_s'] * 160 * WELCH['overlap']))} segments in a 61-s run; "
      f"channels {PICKS} averaged; specparam {specparam.__version__} fixed mode {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz, {SPECPARAM}; IAF window {IAF_RANGE} Hz")

import time


def band_power(freqs, psd, band):
    m = (freqs >= band[0]) & (freqs <= band[1])
    return float(psd[m].sum() * (freqs[1] - freqs[0]))


t_start = time.time()
rows, psds, fits, skipped = [], {}, {}, []
for s in SUBJECTS:
    for cond, run in RUNS.items():
        try:
            raw = helpers.load_spine("ds-eegbci", s, run)
        except Exception as e:                       # a refused or missing run: report, do not hide
            skipped.append(f"{s} {run}: {e}")
            continue
        if raw.info["sfreq"] != 160.0:
            skipped.append(f"{s} {run}: sampling rate {raw.info['sfreq']:g} Hz, not 160 Hz")
            continue
        f, p_ch = helpers_l1.welch_psd(raw.get_data(picks=PICKS) * 1e6, raw.info["sfreq"], **WELCH)
        p = p_ch.mean(axis=0)                                            # occipital average, linear units
        r = helpers_l1.fit_specparam(f, p, freq_range=FIT_RANGE, settings=SPECPARAM, iaf_range=IAF_RANGE)
        psds[(s, cond)], fits[(s, cond)] = p, r
        rows.append(dict(subject=s, condition=cond, duration_s=raw.times[-1] + 1 / raw.info["sfreq"], offset=r["offset"], exponent=r["exponent"],
                         iaf=r["iaf"], iaf_power=r["iaf_power"], iaf_bw=r["iaf_bandwidth"], n_peaks=r["n_peaks"], r_squared=r["r_squared"],
                         error=r["error"], alpha_uV2=band_power(f, p, ALPHA_BAND), broad_uV2=band_power(f, p, BROAD_BAND)))
FREQS = f
df = pd.DataFrame(rows)
ec = df[df.condition == "EC"].set_index("subject")
eo = df[df.condition == "EO"].set_index("subject")
done = sorted(set(ec.index) & set(eo.index))
print(f"{len(done)} subjects with both runs in {time.time() - t_start:.0f} s; skipped: {skipped or 'none'}")
print(f"run durations: median {ec.loc[done, 'duration_s'].median():.1f} s (range {ec.loc[done, 'duration_s'].min():.1f}-{ec.loc[done, 'duration_s'].max():.1f} s)")
with pd.option_context("display.float_format", "{:.2f}".format, "display.width", 200, "display.max_columns", 20, "display.max_rows", 200):
    print(ec.loc[done, ["offset", "exponent", "iaf", "iaf_power", "r_squared", "alpha_uV2"]].rename(columns=lambda c: f"EC_{c}")
          .join(eo.loc[done, ["offset", "exponent", "iaf", "iaf_power", "r_squared", "alpha_uV2"]].rename(columns=lambda c: f"EO_{c}")))

fig, axes = plt.subplots(2, 3, figsize=(16, 9))
for ax, cond, ttl in zip(axes[0, :2], ("EC", "EO"), ("(a) eyes closed (R02)", "(b) eyes open (R01)")):
    spectra = {s: (FREQS, psds[(s, cond)], dict(lw=0.5, color="0.6")) for s in done}
    spectra["log mean"] = (FREQS, 10 ** np.mean([np.log10(psds[(s, cond)]) for s in done], axis=0), dict(lw=2, color="tab:blue"))
    helpers_l1.plot_spectra(spectra, ax=ax, log_x=True, xlim=(1, 40), title=f"{ttl}, {len(done)} subjects", legend=False)
    ax.legend(handles=ax.get_lines()[-1:], fontsize=8)
ax = axes[0, 2]
rng = np.random.default_rng(0)
for i, (cond, color) in enumerate((("EO", "tab:orange"), ("EC", "tab:blue"))):
    vals = (eo if cond == "EO" else ec).loc[done, "iaf"]
    ax.plot(i + rng.uniform(-0.1, 0.1, len(vals)), vals, "o", color=color, alpha=0.7)
    ax.hlines(np.nanmedian(vals), i - 0.25, i + 0.25, color="k", lw=2)
ax.set(xticks=[0, 1], xticklabels=["eyes open", "eyes closed"], xlim=(-0.5, 1.5), ylabel="IAF (Hz)", title="(c) individual alpha frequency (Hz); bar = median")
ax.grid(alpha=0.3, axis="y")
ax = axes[1, 0]
for s in done:
    ax.plot([0, 1], [eo.loc[s, "exponent"], ec.loc[s, "exponent"]], "-o", color="tab:blue", alpha=0.5, ms=4)
ax.set(xticks=[0, 1], xticklabels=["eyes open", "eyes closed"], xlim=(-0.4, 1.4), ylabel="exponent (log10 uV^2/Hz per log10 Hz)", title="(d) aperiodic exponent, paired (log10 uV^2/Hz per log10 Hz)")
ax.grid(alpha=0.3, axis="y")
ax = axes[1, 1]
d_pw = (ec.loc[done, "iaf_power"].fillna(0) - eo.loc[done, "iaf_power"].fillna(0))
d_ex = ec.loc[done, "exponent"] - eo.loc[done, "exponent"]
ax.axhline(0, color="gray", lw=0.6); ax.axvline(0, color="gray", lw=0.6)
ax.plot(d_ex, d_pw, "o", color="tab:blue", alpha=0.7)
for s in done:
    ax.annotate(s[1:], (d_ex[s], d_pw[s]), fontsize=6, textcoords="offset points", xytext=(3, 2))
ax.set(xlabel="exponent EC - EO", ylabel="alpha peak power EC - EO (log10; 0 = no peak)", title="(e) eyes-closed effect: peak vs slope change (log10 units)")
ax.grid(alpha=0.3)
ax = axes[1, 2]
ratio = ec.loc[done, "alpha_uV2"] / eo.loc[done, "alpha_uV2"]
ax.bar(range(len(done)), 10 * np.log10(ratio), color=["tab:blue" if r >= 1 else "tab:orange" for r in ratio])
ax.set(xticks=range(len(done)), xticklabels=[s[1:] for s in done], xlabel="subject", ylabel="8-12 Hz power EC / EO (dB)", title="(f) 8-12 Hz band power, eyes closed / eyes open (dB)")
ax.tick_params(axis="x", labelsize=7)
ax.grid(alpha=0.3, axis="y")
fig.suptitle(f"C1 spectral fingerprint: ds-eegbci {done[0]}-{done[-1]} ({len(done)} subjects), occipital average, Welch {WELCH['seg_s']:g}-s Hann {int(100 * WELCH['overlap'])} %, specparam fixed {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz", y=1.0)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

def q(x):
    x = pd.Series(x).dropna()
    return f"{x.median():6.2f} [{x.quantile(0.25):6.2f}, {x.quantile(0.75):6.2f}] (n={len(x)})"


measures = [("IAF (Hz)", "iaf"), ("alpha peak power above the fit (log10)", "iaf_power"), ("exponent", "exponent"), ("offset (log10 uV^2/Hz at 1 Hz)", "offset"),
            ("8-12 Hz band power (uV^2)", "alpha_uV2"), ("1-40 Hz power (uV^2)", "broad_uV2"), ("R^2", "r_squared")]
print(f"{'measure':40s} {'eyes open: median [IQR]':>36s} {'eyes closed: median [IQR]':>36s}")
for name, col in measures:
    print(f"{name:40s} {q(eo.loc[done, col]):>36s} {q(ec.loc[done, col]):>36s}")
print()
print(f"{'paired EC - EO':40s} {'median [IQR]':>30s} {'EC > EO':>9s} {'Wilcoxon p':>11s}")
summary = {}
for name, col in measures[:5]:
    d = (ec.loc[done, col] - eo.loc[done, col]).dropna()
    if col == "alpha_uV2":
        d = np.log10(ec.loc[done, col] / eo.loc[done, col])
        name = "log10(8-12 Hz power EC / EO)"
    p_w = stats.wilcoxon(d).pvalue if len(d) >= 6 and np.any(d != 0) else np.nan
    summary[col] = dict(median=float(d.median()), n_pos=int((d > 0).sum()), n=int(len(d)), p=float(p_w))
    print(f"{name:40s} {q(d):>30s} {int((d > 0).sum()):>4d}/{len(d):<4d} {p_w:11.3g}")

def robust_z(x):
    x = pd.Series(x, dtype=float)
    mad = 1.4826 * np.nanmedian(np.abs(x - np.nanmedian(x)))
    return (x - np.nanmedian(x)) / mad if mad > 0 else x * 0


z = pd.DataFrame({"exponent_EC": robust_z(ec.loc[done, "exponent"]), "offset_EC": robust_z(ec.loc[done, "offset"]),
                  "log_broad_EC": robust_z(np.log10(ec.loc[done, "broad_uV2"])), "log_broad_EO": robust_z(np.log10(eo.loc[done, "broad_uV2"]))})
flagged = {}
for s in done:
    reasons = []
    if np.isnan(ec.loc[s, "iaf"]):
        reasons.append("A: no alpha peak (7-13 Hz) with eyes closed")
    for cond, frame in (("EC", ec), ("EO", eo)):
        if frame.loc[s, "r_squared"] < 0.9:
            reasons.append(f"B: R^2 {frame.loc[s, 'r_squared']:.2f} ({cond})")
    for col in z.columns:
        if abs(z.loc[s, col]) > 3:
            reasons.append(f"C: {col} robust z {z.loc[s, col]:+.1f}")
    if ec.loc[s, "alpha_uV2"] <= eo.loc[s, "alpha_uV2"]:
        reasons.append(f"D: 8-12 Hz power EC/EO = {ec.loc[s, 'alpha_uV2'] / eo.loc[s, 'alpha_uV2']:.2f}")
    if reasons:
        flagged[s] = reasons
if flagged:
    for s, reasons in flagged.items():
        print(f"{s}: " + "; ".join(reasons))
else:
    s = z.abs().max(axis=1).idxmax()
    flagged[s] = [f"watch: largest |robust z| in the cohort ({z.abs().max(axis=1)[s]:.1f}), no rule fired"]
    print(f"no rule fired; {s} has the largest |robust z| ({z.abs().max(axis=1)[s]:.1f})")
print(f"flagged {len(flagged)} of {len(done)} subjects")

first = next(iter(flagged))
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
for ax, cond in zip(axes, ("EO", "EC")):
    r = fits[(first, cond)]
    ax.plot(r["freqs"], r["log_power"], "k", lw=1.2, label="PSD")
    ax.plot(r["freqs"], r["aperiodic_log"], "--", color="tab:blue", lw=1.2, label=f"aperiodic: offset {r['offset']:.2f}, exponent {r['exponent']:.2f}")
    ax.plot(r["freqs"], r["model_log"], color="tab:orange", lw=1.2, label=f"model, R^2 {r['r_squared']:.2f}")
    ax.set_xscale("log")
    ax.set(xlabel="Frequency (Hz) [log]", ylabel="log10 PSD (log10 uV^2/Hz)", title=f"{first} {'eyes closed' if cond == 'EC' else 'eyes open'}: {'; '.join(flagged[first])[:60]} (log10 uV^2/Hz)")
    ax.grid(alpha=0.3, which="both"); ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

dur = ec.loc[done, "duration_s"].median()
n_seg = helpers_l1.n_welch_segments(int(dur * 160), int(WELCH["seg_s"] * 160), int(WELCH["seg_s"] * 160 * WELCH["overlap"]))
paragraph = (
    f"Resting EEG from the {len(done)} subjects {done[0]}-{done[-1]} of the EEG Motor Movement/Imagery Dataset (PhysioNet, 160 Hz, 64 channels, no hardware filters) "
    f"was analysed in its two one-minute baseline runs (R01 eyes open, R02 eyes closed; median duration {dur:.0f} s). Subjects S088, S089, S092 and S100 were "
    f"excluded because of documented inconsistent event timestamps, and S038 and S104 because several reports drop them; none of these lies in the analysed range. "
    f"Power spectral densities were estimated for O1, Oz and O2 with Welch's method ({WELCH['seg_s']:g}-s Hann segments, {int(100 * WELCH['overlap'])} % overlap, "
    f"{1 / WELCH['seg_s']:g} Hz resolution, {n_seg} segments per run, units uV^2/Hz) and averaged over the three channels in linear units. No filtering was applied "
    f"before spectral estimation: the fit range of {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz excludes the slow drift below 1 Hz and the 60 Hz mains line, so neither a "
    f"high-pass nor a notch would change the fitted quantities [judgement: a cleaned pipeline is deferred to Level 2]; the recording reference was kept as stored "
    f"(TODO(confirm): undocumented). Spectra were parameterized with {fits[(done[0], 'EC')]['tool']} in fixed aperiodic mode over {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz "
    f"(at most {SPECPARAM['max_n_peaks']} peaks, peak widths {SPECPARAM['peak_width_limits'][0]:g}-{SPECPARAM['peak_width_limits'][1]:g} Hz, minimum peak height "
    f"{SPECPARAM['min_peak_height']:g}); the individual alpha frequency was defined as the centre frequency of the highest peak between {IAF_RANGE[0]:g} and "
    f"{IAF_RANGE[1]:g} Hz (undefined when no peak was found: {int(ec.loc[done, 'iaf'].isna().sum())} subject(s) eyes closed, {int(eo.loc[done, 'iaf'].isna().sum())} eyes open). "
    f"Median goodness of fit was R^2 = {ec.loc[done, 'r_squared'].median():.3f} (eyes closed) and {eo.loc[done, 'r_squared'].median():.3f} (eyes open). "
    f"Eyes-closed minus eyes-open differences were summarised as paired medians: alpha peak power {summary['iaf_power']['median']:+.2f} log10 units "
    f"(larger with eyes closed in {summary['iaf_power']['n_pos']} of {summary['iaf_power']['n']}), exponent {summary['exponent']['median']:+.2f} "
    f"(larger in {summary['exponent']['n_pos']} of {summary['exponent']['n']}), IAF {summary['iaf']['median']:+.2f} Hz (n = {summary['iaf']['n']}). "
    f"Anomalies were flagged by a pre-stated rule (no eyes-closed alpha peak; R^2 < 0.9; |robust z| > 3 on exponent, offset or broadband power; no alpha reactivity): "
    f"{', '.join(f'{s} ({'; '.join(r)})' for s, r in flagged.items())}."
)
print(paragraph)

print("nb-c1-spectral-fingerprint -- C1 numbers (draft; TODO(confirm) at author review)")
print(f"cohort: {'FULL' if FULL_COHORT else 'subset'} {done[0]}-{done[-1]}, N = {len(done)}; skipped: {skipped or 'none'}; exclusions applied before loading: {', '.join(excluded)}")
print(f"Welch: {WELCH['seg_s']:g}-s {WELCH['window']}, {int(100 * WELCH['overlap'])} % overlap, {1 / WELCH['seg_s']:g} Hz resolution; channels {PICKS} averaged; "
      f"specparam fixed {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz, {SPECPARAM}")
print(f"eyes closed: median IAF {ec.loc[done, 'iaf'].median():.2f} Hz (n = {int(ec.loc[done, 'iaf'].notna().sum())}), exponent {ec.loc[done, 'exponent'].median():.2f}, "
      f"offset {ec.loc[done, 'offset'].median():.2f}, R^2 {ec.loc[done, 'r_squared'].median():.3f}")
print(f"eyes open  : median IAF {eo.loc[done, 'iaf'].median():.2f} Hz (n = {int(eo.loc[done, 'iaf'].notna().sum())}), exponent {eo.loc[done, 'exponent'].median():.2f}, "
      f"offset {eo.loc[done, 'offset'].median():.2f}, R^2 {eo.loc[done, 'r_squared'].median():.3f}")
for col, name in (("iaf_power", "alpha peak power (log10)"), ("exponent", "exponent"), ("offset", "offset"), ("iaf", "IAF (Hz)"), ("alpha_uV2", "log10 8-12 Hz power ratio")):
    d = summary[col]
    print(f"EC - EO {name:26s}: median {d['median']:+.3f}, EC > EO in {d['n_pos']}/{d['n']}, Wilcoxon p = {d['p']:.3g}")
print(f"flagged: " + "; ".join(f"{s} [{'; '.join(r)}]" for s, r in flagged.items()))
print("Peak vs slope: see the two rows above -- the effect is a peak change if 'alpha peak power' moves consistently, a slope change if 'exponent' does, both if both.")