# 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/L1/ (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

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

DATASET, CH = "ds-eegbci", "O1"
WELCH = dict(seg_s=2.0, overlap=0.5, window="hann")
FIT_RANGE = helpers_l1.FIT_RANGE
print("specparam settings:", helpers_l1.SPECPARAM_SETTINGS, "| fit range", FIT_RANGE, "Hz | IAF window", helpers_l1.IAF_RANGE, "Hz")


def psd_of(subject, run):
    raw = helpers.load_spine(DATASET, subject, run)
    assert raw.info["sfreq"] == 160.0
    f, p = helpers_l1.welch_psd(raw.get_data(picks=CH)[0] * 1e6, raw.info["sfreq"], **WELCH)
    return f, p


f, p_ec = psd_of("S001", "R02")
fit = helpers_l1.fit_specparam(f, p_ec, freq_range=FIT_RANGE)
print(f"S001 R02 {CH}: offset {fit['offset']:.2f} (log10 uV^2/Hz at 1 Hz), exponent {fit['exponent']:.2f}, {fit['n_peaks']} peaks, "
      f"R^2 {fit['r_squared']:.3f}, mean absolute error {fit['error']:.3f} log10 units; {fit['tool']}")
for cf, pw, bw in fit["peaks"]:
    print(f"  peak: centre {cf:5.2f} Hz, power {pw:.2f} log10 units above the aperiodic fit, bandwidth {bw:.2f} Hz")
print(f"  IAF (highest peak in 7-13 Hz): {fit['iaf']:.2f} Hz")

fig, axes = plt.subplots(1, 2, figsize=(13, 4))
ax = axes[0]
ax.plot(fit["freqs"], fit["log_power"], "k", lw=1.2, label="PSD (log10 uV^2/Hz)")
ax.plot(fit["freqs"], fit["aperiodic_log"], "--", color="tab:blue", lw=1.2, label=f"aperiodic fit: {fit['offset']:.2f} - {fit['exponent']:.2f} log10 f")
ax.plot(fit["freqs"], fit["model_log"], color="tab:orange", lw=1.2, label="full model (aperiodic + peaks)")
ax.set_xscale("log")
ax.set(xlabel="Frequency (Hz) [log]", ylabel="log10 PSD (log10 uV^2/Hz)", title=f"S001 eyes closed {CH}: PSD and specparam model, {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz (log10 uV^2/Hz)")
ax.grid(alpha=0.3, which="both"); ax.legend(fontsize=8)
ax = axes[1]
ax.plot(fit["freqs"], fit["log_power"] - fit["aperiodic_log"], "k", lw=1.2, label="PSD minus aperiodic fit")
ax.plot(fit["freqs"], fit["model_log"] - fit["aperiodic_log"], color="tab:orange", lw=1.2, label="fitted peaks")
ax.axhline(0, color="gray", lw=0.6)
ax.set(xlabel="Frequency (Hz)", ylabel="log10 power above the aperiodic fit", title="The periodic part: what is left after the line is removed (log10 units)")
ax.grid(alpha=0.3); ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

def model_psd(freqs, offset, exponent, peaks):
    """Linear PSD (uV^2/Hz) from fixed-mode aperiodic parameters plus Gaussian peaks (cf Hz, pw log10, bw Hz; sigma = bw / 2)."""
    logp = offset - exponent * np.log10(freqs)
    for cf, pw, bw in peaks:
        logp = logp + pw * np.exp(-((freqs - cf) ** 2) / (2 * (bw / 2) ** 2))
    return 10 ** logp


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


ff = fit["freqs"]
alpha_peaks = [pk for pk in fit["peaks"] if 7 <= pk[0] <= 13]
other_peaks = [pk for pk in fit["peaks"] if not 7 <= pk[0] <= 13]
variants = {
    "as fitted": model_psd(ff, fit["offset"], fit["exponent"], fit["peaks"]),
    "exponent + 0.5 (steeper), same peaks": model_psd(ff, fit["offset"], fit["exponent"] + 0.5, fit["peaks"]),
    "same line, alpha peak removed": model_psd(ff, fit["offset"], fit["exponent"], other_peaks),
    "aperiodic line only": model_psd(ff, fit["offset"], fit["exponent"], []),
}
print(f"{'variant':40s} {'8-12 Hz (uV^2)':>15s} {'13-30 Hz (uV^2)':>16s}")
for label, psd in variants.items():
    print(f"{label:40s} {band_power(ff, psd, (8, 12)):15.1f} {band_power(ff, psd, (13, 30)):16.1f}")
ax = helpers_l1.plot_spectra({k: (ff, v) for k, v in variants.items()}, log_x=True, title="Spectra rebuilt from the S001 parameters: slope and peak both move 'band power'")
plt.show()   # render the static figure(s) of this cell inline

SUBJECTS = helpers_l1.eegbci_subjects(1, 10)
rows, psds, fits = [], {}, {}
for s in SUBJECTS:
    for run, cond in (("R02", "EC"), ("R01", "EO")):
        f, p = psd_of(s, run)
        r = helpers_l1.fit_specparam(f, p, freq_range=FIT_RANGE)
        psds[(s, cond)], fits[(s, cond)] = p, r
        rows.append(dict(subject=s, condition=cond, offset=r["offset"], exponent=r["exponent"], iaf=r["iaf"], iaf_power=r["iaf_power"],
                         n_peaks=r["n_peaks"], r_squared=r["r_squared"], error=r["error"],
                         alpha_8_12_uV2=band_power(f, p, (8, 12)), beta_13_30_uV2=band_power(f, p, (13, 30))))
df = pd.DataFrame(rows)
ec = df[df.condition == "EC"].set_index("subject")
eo = df[df.condition == "EO"].set_index("subject")
with pd.option_context("display.float_format", "{:.2f}".format, "display.width", 140):
    print("Eyes closed (R02), O1:")
    print(ec[["offset", "exponent", "iaf", "iaf_power", "n_peaks", "r_squared", "error", "alpha_8_12_uV2", "beta_13_30_uV2"]])
flags = []
for s in SUBJECTS:
    if np.isnan(ec.loc[s, "iaf"]):
        flags.append(f"{s}: no peak between 7 and 13 Hz eyes closed (no IAF)")
    if ec.loc[s, "r_squared"] < 0.9:
        flags.append(f"{s}: R^2 {ec.loc[s, 'r_squared']:.2f} < 0.9 eyes closed")
    if eo.loc[s, "r_squared"] < 0.9:
        flags.append(f"{s}: R^2 {eo.loc[s, 'r_squared']:.2f} < 0.9 eyes open")
print("flags:", "; ".join(flags) if flags else "none")

fig, axes = plt.subplots(1, 4, figsize=(14, 3.6))
for ax, (col, unit) in zip(axes, (("iaf", "Hz"), ("exponent", "log10 uV^2/Hz per log10 Hz"), ("offset", "log10 uV^2/Hz at 1 Hz"), ("r_squared", "fraction"))):
    vals = ec[col].dropna()
    jitter = np.random.default_rng(0).uniform(-0.08, 0.08, len(vals))
    ax.plot(jitter, vals, "o", color="tab:blue", alpha=0.7)
    ax.hlines(np.median(vals), -0.25, 0.25, color="tab:orange", lw=2, label=f"median {np.median(vals):.2f}")
    ax.set(xlim=(-0.5, 0.5), xticks=[], title=f"{col} ({unit})", ylabel=unit)
    ax.grid(alpha=0.3, axis="y"); ax.legend(fontsize=8)
fig.suptitle(f"{len(SUBJECTS)} subjects, eyes closed, {CH}: distribution of the fitted parameters", y=1.03)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

paired = pd.DataFrame({"exponent_EC": ec.exponent, "exponent_EO": eo.exponent, "offset_EC": ec.offset, "offset_EO": eo.offset,
                       "iaf_EC": ec.iaf, "iaf_EO": eo.iaf, "alpha_pw_EC": ec.iaf_power, "alpha_pw_EO": eo.iaf_power,
                       "alpha_8_12_EC": ec.alpha_8_12_uV2, "alpha_8_12_EO": eo.alpha_8_12_uV2})
for name, a, b in (("exponent", "exponent_EC", "exponent_EO"), ("offset", "offset_EC", "offset_EO"), ("IAF (Hz)", "iaf_EC", "iaf_EO"),
                   ("alpha peak power (log10)", "alpha_pw_EC", "alpha_pw_EO"), ("8-12 Hz band power (uV^2)", "alpha_8_12_EC", "alpha_8_12_EO")):
    d = (paired[a] - paired[b]).dropna()
    print(f"{name:28s}: median EC {paired[a].median():6.2f}, EO {paired[b].median():6.2f}; EC - EO median {d.median():+6.2f}, "
          f"positive in {int((d > 0).sum())} of {len(d)} subjects")
print(f"subjects with an alpha peak (7-13 Hz): eyes closed {int(ec.iaf.notna().sum())}/{len(ec)}, eyes open {int(eo.iaf.notna().sum())}/{len(eo)}")

fig, axes = plt.subplots(1, 3, figsize=(13, 3.8))
for ax, (a, b, ttl, unit) in zip(axes, (("exponent_EO", "exponent_EC", "aperiodic exponent", "log10 uV^2/Hz per log10 Hz"),
                                        ("alpha_pw_EO", "alpha_pw_EC", "alpha peak power above the aperiodic fit", "log10 units"),
                                        ("iaf_EO", "iaf_EC", "individual alpha frequency", "Hz"))):
    for s in SUBJECTS:
        ax.plot([0, 1], [paired.loc[s, a], paired.loc[s, b]], "-o", color="tab:blue", alpha=0.6, ms=4)
    ax.set(xticks=[0, 1], xticklabels=["eyes open", "eyes closed"], xlim=(-0.4, 1.4), title=f"{ttl} ({unit})", ylabel=unit)
    ax.grid(alpha=0.3, axis="y")
fig.suptitle(f"{len(SUBJECTS)} subjects, {CH}: paired eyes-open / eyes-closed parameters (missing points = no alpha peak found)", y=1.03)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

bp = {}
for s in SUBJECTS:
    r = fits[(s, "EC")]
    raw_bp = band_power(r["freqs"], 10 ** r["log_power"], (13, 30))
    ap_bp = band_power(r["freqs"], 10 ** r["aperiodic_log"], (13, 30))
    bp[s] = (raw_bp, ap_bp, raw_bp - ap_bp)
best, margin = None, 1.0
for a in SUBJECTS:                      # a pair whose 'more beta' answer flips between the raw integral and the part above the fit
    for b in SUBJECTS:
        if bp[a][0] > bp[b][0] and bp[a][2] < bp[b][2]:
            m = min(bp[a][0] / bp[b][0], bp[b][2] / bp[a][2])
            if m > margin:
                best, margin = (a, b), m
if best is None:
    s_lo, s_hi = ec.exponent.idxmin(), ec.exponent.idxmax()
    print(f"no pair flips in this group; showing the flattest ({s_lo}) and steepest ({s_hi}) spectra instead")
else:
    s_lo, s_hi = best
    print(f"{s_lo} has more raw 13-30 Hz power than {s_hi}, but less 13-30 Hz power above the aperiodic fit (each by a factor of at least {margin:.2f})")
print(f"{'subject':8s} {'exponent':>8s} {'13-30 Hz raw (uV^2)':>20s} {'13-30 Hz aperiodic (uV^2)':>26s} {'13-30 Hz above fit (uV^2)':>26s} {'aperiodic share':>15s}")
for s in (s_lo, s_hi):
    print(f"{s:8s} {ec.loc[s, 'exponent']:8.2f} {bp[s][0]:20.1f} {bp[s][1]:26.1f} {bp[s][2]:26.1f} {100 * bp[s][1] / bp[s][0]:14.0f} %")
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
for ax, s in zip(axes, (s_lo, s_hi)):
    r = fits[(s, "EC")]
    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: exponent {r['exponent']:.2f}")
    ax.fill_between(r["freqs"], r["aperiodic_log"], r["log_power"], where=(r["freqs"] >= 13) & (r["freqs"] <= 30), color="tab:orange", alpha=0.3, label="13-30 Hz above the fit")
    ax.set_xscale("log")
    ax.set(xlabel="Frequency (Hz) [log]", ylabel="log10 PSD (log10 uV^2/Hz)", title=f"{s} eyes closed {CH} (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

print("nb-1-7-specparam -- L1.7 numbers (draft; TODO(confirm) at author review)")
print(f"Data: {DATASET} {SUBJECTS[0]}-{SUBJECTS[-1]} ({len(SUBJECTS)} subjects), {CH}, R02 eyes closed (and R01 eyes open); Welch {WELCH['seg_s']:g}-s {WELCH['window']} "
      f"{int(100 * WELCH['overlap'])} % overlap; {fit['tool']}, fixed mode, fit {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz, settings {helpers_l1.SPECPARAM_SETTINGS}")
print(f"{'subject':8s} {'IAF (Hz)':>9s} {'exponent':>9s} {'offset':>8s} {'R^2':>6s}   | eyes open: {'IAF':>5s} {'exponent':>9s} {'offset':>8s}")
for s in SUBJECTS:
    print(f"{s:8s} {ec.loc[s, 'iaf']:9.2f} {ec.loc[s, 'exponent']:9.2f} {ec.loc[s, 'offset']:8.2f} {ec.loc[s, 'r_squared']:6.3f}   | "
          f"{eo.loc[s, 'iaf']:16.2f} {eo.loc[s, 'exponent']:9.2f} {eo.loc[s, 'offset']:8.2f}")
print(f"group medians, eyes closed: IAF {ec.iaf.median():.2f} Hz (n = {int(ec.iaf.notna().sum())} with a peak), exponent {ec.exponent.median():.2f}, offset {ec.offset.median():.2f}")
print(f"group medians, eyes open  : IAF {eo.iaf.median():.2f} Hz (n = {int(eo.iaf.notna().sum())} with a peak), exponent {eo.exponent.median():.2f}, offset {eo.offset.median():.2f}")
print(f"paired EC - EO medians: exponent {(ec.exponent - eo.exponent).median():+.2f}, offset {(ec.offset - eo.offset).median():+.2f}, "
      f"IAF {(ec.iaf - eo.iaf).dropna().median():+.2f} Hz (n = {int((ec.iaf - eo.iaf).notna().sum())})")
print("flags:", "; ".join(flags) if flags else "none")