# 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')
_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_l4.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L4/ (or notebooks/) so that "
                            "_shared/helpers_l4.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l1
import helpers_l4 as L4

# 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

# 4. Quiet the downloader.  pooch, which MNE uses to fetch datasets, logs
#    "Downloading file '...' from '...' to '<cache directory>'" at INFO, and that last field is an
#    ABSOLUTE PATH from whichever machine executed the notebook.  Absolute paths are not allowed in a
#    stored notebook (scripts/scrub-notebooks.py is a CI gate) and re-executing would put them straight
#    back, so the message is suppressed at the source rather than cleaned up afterwards.  Nothing is
#    hidden by this: every cell below prints the file NAMES it fetched and helpers_l4.Downloads prints
#    free disk before and after.  Please do not delete this as noise.
try:
    import pooch

    pooch.get_logger().setLevel("WARNING")
except Exception:  # pooch absent or its API moved: the scrub script is the backstop
    pass
print(f"MNE {mne.__version__}; helpers_l4 imported from notebooks/_shared")
print(f"downloads go to {helpers_l1.download_dir().name}/ (resolved relative to the working directory, "
      "or $EEG_COURSE_DOWNLOADS) and are deleted at the end of this notebook")

SUBJECT = "S001"
CHANNELS = ["C3", "C4"]
dl = L4.Downloads("nb-4-4").start()
epochs, info = L4.load_imagery_epochs(SUBJECT, L4.MI_RUNS, dl)
print(f"{SUBJECT}: {info['n_epochs']} epochs, {info['n_eeg']} EEG channels at {info['sfreq']:g} Hz, "
      f"{info['epoch_window_s'][0]:g}..{info['epoch_window_s'][1]:g} s, average reference, no filtering")
print(f"task duration from the EDF's own annotations: {info['task_duration_s']:g} s")
print("downloaded for this notebook: " + ", ".join(p.name for p in dl.paths))

TF = {}
for cond in ("T1", "T2"):
    X = epochs[cond].get_data(picks=CHANNELS) * 1e6
    r = L4.morlet_power(X, sfreq=float(epochs.info["sfreq"]), times=epochs.times)
    TF[cond] = r
    print(f"  {cond} ({L4.MI_CONDITIONS[cond]}): single-trial power {r['power'].shape} "
          f"(trials x channels x freqs x times)")
freqs, times = TF["T2"]["freqs"], TF["T2"]["times"]
edge = L4.edge_seconds(freqs, L4.TF_CYCLES)
print(f"  call: {TF['T2']['call']}")

i_c3 = CHANNELS.index("C3")
P_raw = TF["T2"]["power"].mean(0)[i_c3]              # trial-averaged raw power, C3, T2

fig, axes = plt.subplots(1, 2, figsize=(14, 4.2))
im = axes[0].pcolormesh(times, freqs, P_raw, cmap="viridis", shading="nearest")
axes[0].axvline(0, color="w", lw=1.0, ls="--")
axes[0].set(xlabel="Time from cue (s)", ylabel="Frequency (Hz)",
            title=f"RAW trial-averaged power, {SUBJECT} C3, T2\n(library units, linear scale) -- this is a "
                  f"picture of 1/f")
fig.colorbar(im, ax=axes[0]).set_label("Power (library units)")
L4.plot_tfr(L4.baseline_normalise(P_raw, times, L4.TF_BASELINE, "percent"), freqs, times, ax=axes[1],
            vlim=(-100, 100), edge_s=edge, baseline=L4.TF_BASELINE,
            cbar_label="Power change from baseline (%)",
            title=f"The SAME array, per-frequency percent change\n{SUBJECT} C3, T2, baseline "
                  f"{L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

rows = [{"frequency (Hz)": f, "mean raw power over the epoch (library units)": P_raw[i].mean(),
         "as a multiple of the 40 Hz row": P_raw[i].mean() / P_raw[-1].mean()}
        for i, f in enumerate(freqs) if f in (4, 10, 20, 30, 40)]
print(L4.fmt_table(rows, floatfmt="{:.4g}"))
print()
print("That last column is the reason a raw map is useless as a picture: the 4 Hz row is "
      f"{P_raw[0].mean() / P_raw[-1].mean():.0f} times the 40 Hz row before anything has happened. Any "
      "colour scale that shows one hides the other.")

fig, axes = plt.subplots(1, 4, figsize=(21.5, 4.0))
maps = {}
for ax, mode, lim, lab in ((axes[0], "percent", 100, "Power change from baseline (%)"),
                           (axes[1], "db", 6, "Power change from baseline (dB)"),
                           (axes[2], "z", 6, "z against the baseline window's own SD over time")):
    M = L4.baseline_normalise(P_raw, times, L4.TF_BASELINE, mode)
    maps[mode] = M
    L4.plot_tfr(M, freqs, times, ax=ax, vlim=(-lim, lim), edge_s=edge, baseline=L4.TF_BASELINE,
                cbar_label=lab, title=f"{mode} \n{SUBJECT} C3, T2, baseline "
                                      f"{L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s")
# z across trials: each trial's own baseline mean, then the SD of those per-trial values
tr = TF["T2"]["power"][:, i_c3]
bm = (times >= L4.TF_BASELINE[0]) & (times <= L4.TF_BASELINE[1])
base_per_trial = tr[:, :, bm].mean(-1)                       # (n_trials, n_freqs)
z_trials = (tr.mean(0) - base_per_trial.mean(0)[:, None]) / base_per_trial.std(0, ddof=1)[:, None]
maps["z across trials"] = z_trials
L4.plot_tfr(z_trials, freqs, times, ax=axes[3], vlim=(-6, 6), edge_s=edge, baseline=L4.TF_BASELINE,
            cbar_label="z against the ACROSS-TRIAL SD of the baseline",
            title=f"z across trials\n{SUBJECT} C3, T2 -- a different quantity")
fig.suptitle(f"One array, four normalisations ({SUBJECT} C3, T2, {tr.shape[0]} trials)", y=1.04)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

mu = (freqs >= L4.MU_BAND[0]) & (freqs <= L4.MU_BAND[1])
act = (times >= L4.TF_ACTIVE[0]) & (times <= L4.TF_ACTIVE[1])
print(f"Mu band over {L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s at C3, the same cells read four ways:")
for k, M in maps.items():
    print(f"    {k:18s}: {M[np.ix_(mu, act)].mean():+9.3f}")
print(f"    percent -> dB check: 10 log10(1 + {maps['percent'][np.ix_(mu, act)].mean():.3f}/100) = "
      f"{10 * np.log10(1 + maps['percent'][np.ix_(mu, act)].mean() / 100):+.3f} dB, which is NOT the mean dB "
      f"({maps['db'][np.ix_(mu, act)].mean():+.3f}) -- the log of a mean is not the mean of a log, and that is "
      "a third place where an averaging order changes an answer.")

WINDOWS = {
    f"default {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s": L4.TF_BASELINE,
    "-2.5..-2.0 s (inside the edge region)": (-2.5, -2.0),
    "-0.5..0.0 s (right up to the cue)": (-0.5, 0.0),
    "-2.5..0.0 s (everything before the cue)": (-2.5, 0.0),
    "1.0..2.0 s (inside the task -- wrong on purpose)": (1.0, 2.0),
}
rows = []
for name, win in WINDOWS.items():
    r = {"baseline window": name,
         "inside the edge region?": "yes" if win[0] < times[0] + edge[0] else "no"}
    for ch in CHANNELS:
        i = CHANNELS.index(ch)
        r[f"mu ERD at {ch} (%)"] = L4.band_erd(TF["T2"]["power"][:, i], freqs, times, L4.MU_BAND,
                                               baseline=win, active=L4.TF_ACTIVE, order="band-first")
    rows.append(r)
print(f"Right-fist imagery (T2), mu {L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz over "
      f"{L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s, band-first order:")
print(L4.fmt_table(rows, floatfmt="{:.2f}"))
print()
bc3 = TF["T2"]["power"].mean(0)[i_c3][mu].mean(0)
for name, win in WINDOWS.items():
    m = (times >= win[0]) & (times <= win[1])
    print(f"    baseline power in {name:48s}: {bc3[m].mean():9.2f} library units")
print()
print("The edge-region baseline reads low because part of it is convolved with zeros, so every later value is "
      "inflated against it; the up-to-the-cue baseline reads low because the desynchronisation has already "
      "begun; the in-task baseline is the effect itself, so the ERD it reports is close to nothing.")

fig, axes = plt.subplots(1, 3, figsize=(17.5, 4.0))
for ax, (name, win) in zip(axes, list(WINDOWS.items())[:3]):
    L4.plot_tfr(L4.baseline_normalise(P_raw, times, win, "percent"), freqs, times, ax=ax, vlim=(-100, 100),
                edge_s=edge, baseline=win, cbar_label="Power change from baseline (%)",
                title=f"baseline {name}\n{SUBJECT} C3, T2 (%)")
fig.suptitle("Moving the baseline moves the map (the black bar at the bottom of each panel is the window)",
             y=1.04)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

table = []
for cond in ("T2", "T1"):
    for band_name, band in (("mu", L4.MU_BAND), ("beta", L4.BETA_BAND)):
        for ch in CHANNELS:
            i = CHANNELS.index(ch)
            row = {"condition": f"{cond} ({L4.MI_CONDITIONS[cond]})", "band": band_name, "channel": ch}
            for order in L4.ERD_ORDERS:
                row[order] = L4.band_erd(TF[cond]["power"][:, i], freqs, times, band,
                                         baseline=L4.TF_BASELINE, active=L4.TF_ACTIVE, order=order)
            table.append(row)
print(f"ERD in percent, baseline {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s, active window "
      f"{L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s, {SUBJECT}")
print(L4.fmt_table(table, floatfmt="{:.2f}"))
print()
for k, v in L4.ERD_ORDERS.items():
    print(f"  {k:18s}: {v}")

fig, axes = plt.subplots(2, 2, figsize=(14.5, 8.0))
peaks = {}
for r, cond in enumerate(("T2", "T1")):
    for c, ch in enumerate(CHANNELS):
        i = CHANNELS.index(ch)
        M = L4.baseline_normalise(TF[cond]["power"].mean(0)[i], times, L4.TF_BASELINE, "percent")
        L4.plot_tfr(M, freqs, times, ax=axes[r, c], vlim=(-80, 80), edge_s=edge, baseline=L4.TF_BASELINE,
                    cbar_label="Power change from baseline (%)",
                    title=f"{SUBJECT} {ch}, {cond} ({L4.MI_CONDITIONS[cond]})\n% change re "
                          f"{L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s")
        pb = TF[cond]["power"].mean(0)[i][mu].mean(0)
        curve = 100 * (pb - pb[(times >= L4.TF_BASELINE[0]) & (times <= L4.TF_BASELINE[1])].mean()) / \
            pb[(times >= L4.TF_BASELINE[0]) & (times <= L4.TF_BASELINE[1])].mean()
        safe = (times >= times[0] + edge[0]) & (times <= times[-1] - edge[0])
        j = int(np.argmin(np.where(safe, curve, np.inf)))
        peaks[(cond, ch)] = (float(curve[j]), float(times[j]))
fig.suptitle(f"Mu and beta desynchronisation is contralateral ({SUBJECT}, % change from baseline)", y=1.01)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"Peak mu ERD (most negative point of the {L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz mean over time, "
      f"searched only inside the edge-safe region {times[0] + edge[0]:.3f}..{times[-1] - edge[0]:.3f} s):")
for (cond, ch), (v, t) in peaks.items():
    print(f"    {cond} {ch}: {v:+6.1f} % at {t:+.3f} s")

ALL = [epochs.ch_names[i] for i in mne.pick_types(epochs.info, eeg=True)]
topo = {}
for cond in ("T1", "T2"):
    X = epochs[cond].get_data(picks=ALL) * 1e6
    r = L4.morlet_power(X, sfreq=float(epochs.info["sfreq"]), freqs=np.arange(8.0, 13.5, 1.0),
                        n_cycles=np.full(6, 5.0), times=epochs.times)
    P = r["power"].mean(0)                                   # (n_ch, n_freqs, n_times)
    b = (r["times"] >= L4.TF_BASELINE[0]) & (r["times"] <= L4.TF_BASELINE[1])
    a = (r["times"] >= L4.TF_ACTIVE[0]) & (r["times"] <= L4.TF_ACTIVE[1])
    pb = P.mean(1)                                           # band-first
    topo[cond] = 100 * (pb[:, a].mean(1) - pb[:, b].mean(1)) / pb[:, b].mean(1)
info_eeg = mne.pick_info(epochs.info, mne.pick_types(epochs.info, eeg=True))

fig, axes = plt.subplots(1, 3, figsize=(14, 4.2))
for ax, (name, v, lim) in zip(axes, (("T1, left-fist imagery", topo["T1"], 60),
                                     ("T2, right-fist imagery", topo["T2"], 60),
                                     ("T2 minus T1", topo["T2"] - topo["T1"], 40))):
    im, _ = mne.viz.plot_topomap(v, info_eeg, axes=ax, show=False, cmap="RdBu_r", vlim=(-lim, lim),
                                 contours=4, sensors=True)
    ax.set_title(f"{name}\nmu {L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz, % change re baseline", fontsize=9)
    fig.colorbar(im, ax=ax, shrink=0.8).set_label("Power change from baseline (%)")
fig.suptitle(f"{SUBJECT}: mu-band percent change over the scalp, band-first order "
             f"({L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s vs {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s)",
             y=1.03)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

for cond in ("T1", "T2"):
    order = np.argsort(topo[cond])
    print(f"{cond} ({L4.MI_CONDITIONS[cond]}): strongest mu decrease at "
          + ", ".join(f"{ALL[i]} {topo[cond][i]:+.1f} %" for i in order[:5]))
lat = {ch: (topo["T2"][ALL.index(ch)], topo["T1"][ALL.index(ch)]) for ch in CHANNELS}
print()
print(f"At C3 and C4: T2 gives {lat['C3'][0]:+.1f} % / {lat['C4'][0]:+.1f} %, "
      f"T1 gives {lat['C3'][1]:+.1f} % / {lat['C4'][1]:+.1f} %.")
print(f"Contralateral dominance for T2 (C3 more negative than C4): {lat['C3'][0] < lat['C4'][0]}. "
      f"For T1 (C4 more negative than C3): {lat['C4'][1] < lat['C3'][1]}.")
print("These topography numbers use a 6-row 8-13 Hz axis with 5 cycles everywhere, not the 37-row "
      "max(3, f/2) axis of the tables above, so they are close to but not identical with them; the axis is "
      "stated with them for that reason.")

try:
    KEY = {"C3": -47.9, "C4": -10.5}            # site/notes/integration-phase3.md, data-p3a and widgets-J
    KEY_OTHER = {"C3": -44.0, "C4": -9.0}       # the same question, cell-first order
    got = {ch: {o: L4.band_erd(TF["T2"]["power"][:, CHANNELS.index(ch)], freqs, times, L4.MU_BAND,
                               baseline=L4.TF_BASELINE, active=L4.TF_ACTIVE, order=o)
                for o in L4.ERD_ORDERS} for ch in CHANNELS}
    print("nb-4-4-erd -- L4.4 exercise numbers (draft; TODO(confirm) at author review)")
    print(f"Data: ds-eegbci {SUBJECT}, runs {'+'.join(info['runs'])} "
          f"(PhysioNet DOI {L4.DATASETS_L4['ds-eegbci']['dataset_doi']}; ODC-By 1.0); "
          f"{info['n_eeg']} EEG channels at {info['sfreq']:g} Hz, average reference, NO filtering, "
          f"no epoch rejection.")
    print(f"Trials: T2 (right-fist imagery) {info['n_epochs']['T2']}, T1 (left-fist imagery) "
          f"{info['n_epochs']['T1']}; largest absolute sample {info['max_abs_uv']['T2']:.0f} uV (T2).")
    print(f"Time-frequency: {TF['T2']['call']}")
    print(f"Windows: baseline {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s, active "
          f"{L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s, mu {L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz, "
          f"beta {L4.BETA_BAND[0]:g}-{L4.BETA_BAND[1]:g} Hz.")
    print(f"Subject selection: S001 was the top-ranked of ten candidates by mu ERD at C3 "
          f"(w-tf-baseline-explorer/tfpower.json records the ranking). A subject chosen for the size of its "
          f"effect is not a sample of anything.")
    print()
    print("ANSWER KEY -- ex-4-4 (numeric pair): peak mu ERD at C3 and C4 for RIGHT-hand imagery")
    print(f"    band-first order:  C3 {got['C3']['band-first']:+.1f} %   C4 {got['C4']['band-first']:+.1f} %"
          f"    <-- the key")
    print(f"    THE PROMPT MUST NAME THE ORDER. The same question under the other three:")
    for o in ("cell-first", "trial-band-first", "trial-cell-first"):
        print(f"        {o:18s}: C3 {got['C3'][o]:+7.1f} %   C4 {got['C4'][o]:+7.1f} %")
    print(f"    Cross-check: site/notes/integration-phase3.md records C3 {KEY['C3']:g} % and "
          f"C4 {KEY['C4']:g} % (band-first) from data-p3a and widgets-J, with "
          f"C3 {KEY_OTHER['C3']:g} % / C4 {KEY_OTHER['C4']:g} % for the cell-first order.")
    ok_bf = all(abs(got[ch]['band-first'] - KEY[ch]) < 0.06 for ch in CHANNELS)
    ok_cf = all(abs(got[ch]['cell-first'] - KEY_OTHER[ch]) < 0.06 for ch in CHANNELS)
    print(f"    This notebook {'AGREES' if ok_bf else 'DISAGREES'} on the band-first pair "
          f"(differences {', '.join(f'{ch} {got[ch]['band-first'] - KEY[ch]:+.3f}' for ch in CHANNELS)}) "
          f"and {'AGREES' if ok_cf else 'DISAGREES'} on the cell-first pair "
          f"(differences {', '.join(f'{ch} {got[ch]['cell-first'] - KEY_OTHER[ch]:+.3f}' for ch in CHANNELS)}).")
    print(f"    NEW HERE: the two PER-TRIAL orders were not reported by the widget tracks and are far apart "
          f"from both: C3 {got['C3']['trial-band-first']:+.1f} % and {got['C3']['trial-cell-first']:+.1f} %, "
          f"and on C4 they change SIGN ({got['C4']['trial-band-first']:+.1f} % and "
          f"{got['C4']['trial-cell-first']:+.1f} %). A per-trial ratio is heavy-tailed, so its mean is pulled "
          f"upward by trials with a quiet baseline. A learner who normalises single trials and then averages "
          f"will not reproduce the key, and the prompt should keep them from having to guess.")
    print()
    print("Supporting -- peak (not window-mean) mu ERD, searched inside the edge-safe region:")
    for (cond, ch), (v, t) in peaks.items():
        print(f"    {cond} {ch}: {v:+6.1f} % at {t:+.3f} s")
    print("    (site/notes/integration-phase3.md records C3 -64.3 % at 1.15 s and C4 -36.7 % at 0.925 s "
          "for T2.)")
    print()
    print("Supporting -- beta, same windows, band-first:")
    for ch in CHANNELS:
        v = L4.band_erd(TF["T2"]["power"][:, CHANNELS.index(ch)], freqs, times, L4.BETA_BAND,
                        baseline=L4.TF_BASELINE, active=L4.TF_ACTIVE, order="band-first")
        print(f"    T2 {ch}: {v:+.1f} %   (the notes record C3 -35.8 % and C4 -12.0 %)")
    print()
    print("Supporting -- how much the baseline window alone moves the answer (band-first, T2, C3):")
    for r in rows:
        print(f"    {r['baseline window']:48s} {r['mu ERD at C3 (%)']:+7.2f} %")
    print()
    print("Pitfalls: pf-tf-edge-effects (every map here shades the edge region and one baseline is put "
          "inside it on purpose), pf-band-power-slope. Widget: w-tf-baseline-explorer.")
    print("TODO(confirm): the catalog carries no published ERD value for this dataset, so none of the above "
          "is compared with a literature number.")
finally:
    dl.finish()