# Setup: dependencies, the shared helpers, non-interactive plotting.
import importlib.util
import io
import re
import subprocess
import sys
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
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 kernel installs it once.  On Colab,
#    run from a clone of the repository so that notebooks/_shared/ is importable
#    (repository URL: TODO(confirm), spec section 13 item 3).
_needed = ("mne", "scipy", "matplotlib", "pandas", "pooch", "moabb")
_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", "moabb==1.7.2"]
    subprocess.check_call(_cmd)

# 2. Shared helpers, located relative to the working directory -- notebooks/L7/ or notebooks/ --
#    never through an absolute path.  helpers_l7.py belongs to track notebooks-L7a; this notebook
#    imports it when it is present and works without it (nothing below depends on it).
_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/L7/ (or notebooks/) so that "
                            "_shared/helpers_l4.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l1 as L1
import helpers_l4 as L4
import helpers_l5 as L5

HAVE_L7 = importlib.util.find_spec("helpers_l7") is not None
if HAVE_L7:
    import helpers_l7 as L7

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import mne
from scipy import signal as sps

mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72

# 3. Quiet the downloader.  pooch, which MNE and moabb use 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.  A stored notebook may not contain
#    one (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:
#    every cell below prints the file NAMES it fetched and the free disk before and after.
try:
    import pooch

    pooch.get_logger().setLevel("WARNING")
except Exception:  # pooch absent or its API moved: the scrub script is the backstop
    pass

SEED = 20260918
rng = np.random.default_rng(SEED)

# How much of each dataset is read.  FULL_RUN is the thorough local run; the stored run is the
# documented subset.  None of the conclusions below depends on the subset size, and section 3
# prints the spread across subjects so a reader can see how stable each number is.
FULL_RUN = False
N_EEGBCI = 5 if not FULL_RUN else 10          # ds-eegbci subjects for the resting spectra
N_AREEG = 3 if not FULL_RUN else 10           # ds-areeg participants
AREEG_TRIALS = None                            # None = every recording that participant has

print(f"MNE {mne.__version__}; helpers imported from notebooks/_shared "
      f"(helpers_l7 present: {HAVE_L7})")
print(f"FULL_RUN = {FULL_RUN}: {N_EEGBCI} ds-eegbci subjects, {N_AREEG} ds-areeg participants, "
      f"ds-brain-invaders subject 1")

# data/directory.yaml, found by walking up from the working directory (no absolute path).
# helpers_l5 reads the same file for its licence lines; this is the acquisition side of it, which
# no shared helper exposes yet (recorded as a request to helpers_l7 in site/notes/notebooks-L7b.md).
import yaml

_DIR_YAML = next((d / "data" / "directory.yaml" for d in (Path.cwd(), *Path.cwd().parents)
                  if (d / "data" / "directory.yaml").exists()), None)
if _DIR_YAML is None:
    raise FileNotFoundError("data/directory.yaml not found above the working directory")
DIRECTORY = {e["id"]: e for e in yaml.safe_load(_DIR_YAML.read_text())["datasets"]}
print(f"data/directory.yaml: {len(DIRECTORY)} dataset entries")

CLASSES = {
    "ds-eegbci": "64-ch gel research cap",
    "ds-brain-invaders": "16-ch dry research headset",
    "ds-areeg": "14-ch consumer saline headset",
    "ds-mpeng": "14-ch consumer saline headset (named only -- CC BY-NC)",
}
FIELDS = ("device_class", "device", "channels", "sfreq_hz", "reference", "online_filters",
          "mains_hz", "paradigms", "population")
rows = []
for ds, label in CLASSES.items():
    e = DIRECTORY[ds]
    row = {"dataset": ds, "class in this notebook": label}
    for f in FIELDS:
        v = e.get(f, e.get(f + "_note", "TODO(confirm)"))
        row[f] = ", ".join(map(str, v)) if isinstance(v, list) else str(v)
    rows.append(row)
facts = pd.DataFrame(rows).set_index("dataset")
for col in facts.columns:
    print(f"\n--- {col} ---")
    for ds in facts.index:
        print(f"   {ds:20s} {facts.loc[ds, col]}")

# Licences, from the same file, through the helper that formats them.
L5.print_licences("ds-eegbci", "ds-brain-invaders", "ds-areeg", "ds-mpeng", notes=True)
print()
for ds in ("ds-eegbci", "ds-brain-invaders", "ds-areeg", "ds-mpeng"):
    cav = DIRECTORY[ds].get("caveats", [])
    print(f"{ds} -- caveats the catalog records:")
    for c in cav:
        print(f"   - {c}")
    print()

ANALYSES = {
    "alpha peak / IAF": {
        "needs": "a posterior electrode, >= 30 s of a resting or low-demand condition, resolution well "
                 "below 1 Hz, usable bandwidth to ~15 Hz",
    },
    "aperiodic slope (1-40 Hz)": {
        "needs": "usable bandwidth across the whole fit range; a hardware roll-off inside the range "
                 "becomes part of the fitted slope",
    },
    "mu ERD lateralisation": {
        "needs": "C3 and C4 (or near), a motor execution/imagery paradigm with cues, and a reference "
                 "that does not itself carry the sensorimotor signal",
    },
    "P300 amplitude": {
        "needs": "a centro-parietal electrode (Pz/CPz), an oddball or comparable rare-event paradigm, "
                 "and enough trials for an average",
    },
    "scalp topography / source": {
        "needs": "enough electrodes with known positions to interpolate a field -- and an average "
                 "reference that approximates an inactive one",
    },
}
DEVICE_CHANNELS = {
    "ds-eegbci": None,           # filled from the loaded recordings below
    "ds-brain-invaders": None,
    "ds-areeg": ("AF3", "F7", "F3", "FC5", "T7", "P7", "O1", "O2", "P8", "T8", "FC6", "F4", "F8", "AF4"),
}
PARADIGM = {
    "ds-eegbci": "eyes-open/eyes-closed baselines + motor execution and imagery (cued)",
    "ds-brain-invaders": "visual P300 oddball (Brain Invaders game), target/non-target",
    "ds-areeg": "imagined speech of 16 Arabic words, 10 s per word with eyes closed, no cue-locked "
                "sensory or motor event",
}
for ds, p in PARADIGM.items():
    print(f"{ds:20s} paradigm: {p}")
print()
for a, d in ANALYSES.items():
    print(f"{a}\n   needs: {d['needs']}")
print()
print("The matrix is filled in section 9, after every entry in it has a measured number behind it.")

FIT_RANGE = (2.0, 40.0)
ALPHA_BAND = (7.0, 13.0)
MIN_DB = 3.0


def peak_above_trend(freqs, psd_uv2_hz, *, band=ALPHA_BAND, fit_range=FIT_RANGE, min_db=MIN_DB):
    """helpers.alpha_peak's algorithm on a spectrum instead of a Raw.

    Straight-line fit of 10*log10(PSD) against log10(f) over ``fit_range`` excluding 5-15 Hz, then the
    largest residual inside ``band``.  Returns frequency (Hz), height above the trend (dB), and the
    fitted slope, which is reported because a device's roll-off shows up in it.
    """
    freqs = np.asarray(freqs, float)
    psd = np.asarray(psd_uv2_hz, float)
    if psd.ndim > 1:
        psd = psd.mean(axis=0)
    keep_fit = (freqs >= fit_range[0]) & (freqs <= fit_range[1]) & (psd > 0)
    f, p = freqs[keep_fit], 10 * np.log10(psd[keep_fit])
    keep = (f < 5) | (f > 15)
    slope, icpt = np.polyfit(np.log10(f[keep]), p[keep], 1)
    resid = p - (slope * np.log10(f) + icpt)
    m = (f >= band[0]) & (f <= band[1])
    i = int(np.argmax(np.where(m, resid, -np.inf)))
    # A maximum sitting ON the search limit is the band's choice, not the data's: the residual is still
    # rising where the band stops.  A recording with no posterior rhythm does exactly that, and calling
    # the edge an "alpha peak" manufactures a value.  Reported as absent instead.  (nb-7-6 uses the same
    # guard; without it four eyes-open ds-iowapd recordings there all landed on 7.00 Hz and produced a
    # group difference built entirely on non-peaks.)
    idx = np.where(m)[0]
    at_edge = bool(i == idx[0] or i == idx[-1])
    ok = (resid[i] >= min_db) and not at_edge
    return {"freq_hz": float(f[i]) if ok else None, "at_band_edge": at_edge,
            "db_above_trend": float(resid[i]), "slope_db_per_decade": float(slope)}


def welch_uv2(x_uv, sfreq, *, seconds=2.0):
    """Welch PSD in uV^2/Hz with a Hann window of ``seconds`` and 50 % overlap."""
    nperseg = int(round(seconds * sfreq))
    nperseg = min(nperseg, np.asarray(x_uv).shape[-1])
    f, p = sps.welch(np.asarray(x_uv, float), fs=sfreq, window="hann", nperseg=nperseg,
                     noverlap=nperseg // 2, detrend="constant", scaling="density")
    return f, p


print(f"estimator: line fit over {FIT_RANGE[0]:g}-{FIT_RANGE[1]:g} Hz excluding 5-15 Hz; peak = largest "
      f"residual in {ALPHA_BAND[0]:g}-{ALPHA_BAND[1]:g} Hz, reported only above {MIN_DB:g} dB and only "
      f"when it is not on the band's own edge")

dl = L4.Downloads("nb-7-5").start()

POSTERIOR = ("O1", "Oz", "O2", "PO3", "POz", "PO4", "PO7", "PO8", "P3", "Pz", "P4", "P7", "P8")


def existing_then_add(paths, downloads):
    """Register for deletion ONLY the files this run created.

    Two Phase 3 tracks sharing one cache deleted each other's subjects (notebooks/README.md).  A file
    that was already on disk when this notebook started was not downloaded by it and is left alone.
    """
    new = [p for p in paths if Path(p).is_file() and Path(p).stat().st_mtime >= _T_START]
    downloads.add(*new)
    return new


_T_START = time.time()

# --- ds-eegbci: R02, the one-minute eyes-CLOSED baseline ------------------------------------------
t0 = time.time()
eegbci_psd, eegbci_rows = {}, []
subjects = [f"S{i:03d}" for i in range(1, N_EEGBCI + 1)]
for s in subjects:
    paths = L4.fetch_eegbci(s, (2,), verbose=False)
    existing_then_add(paths, dl)
    raw = helpers.load_spine("ds-eegbci", s, "R02")
    chs = [c for c in POSTERIOR if c in raw.ch_names]
    x = raw.get_data(picks=chs) * 1e6
    f, P = welch_uv2(x, raw.info["sfreq"])
    eegbci_psd[s] = (f, P.mean(0))
    pk = peak_above_trend(f, P.mean(0))
    ref = helpers.alpha_peak(raw, picks=POSTERIOR, band=ALPHA_BAND, fit_range=FIT_RANGE, min_db=MIN_DB)
    eegbci_rows.append({"subject": s, "n_eeg": len(mne.pick_types(raw.info, eeg=True)),
                        "duration_s": round(float(raw.times[-1]), 1),
                        "iaf_hz": pk["freq_hz"], "db": round(pk["db_above_trend"], 2),
                        "helpers.alpha_peak": ref["freq_hz"],
                        "agree": (pk["freq_hz"] == ref["freq_hz"])})
    DEVICE_CHANNELS["ds-eegbci"] = tuple(raw.copy().pick("eeg").ch_names)
    del raw
eegbci_tab = pd.DataFrame(eegbci_rows)
print(f"ds-eegbci: {len(eegbci_tab)} subjects, run R02 (eyes closed), {time.time() - t0:.0f} s")
print(eegbci_tab.to_string(index=False))
print(f"\nCROSS-CHECK: the notebook's spectrum-domain estimator agrees with helpers.alpha_peak on "
      f"{int(eegbci_tab.agree.sum())} of {len(eegbci_tab)} subjects "
      f"({'identical' if eegbci_tab.agree.all() else 'DISAGREEMENT -- see the rows above'}).")

# --- ds-brain-invaders: the 16-channel dry headset -------------------------------------------------
t0 = time.time()
raw_bi = helpers.load_spine("ds-brain-invaders", 1)
bi_channels = tuple(raw_bi.copy().pick("eeg").ch_names)
DEVICE_CHANNELS["ds-brain-invaders"] = bi_channels
chs = [c for c in POSTERIOR if c in raw_bi.ch_names]
x = raw_bi.get_data(picks=chs) * 1e6
f_bi, P_bi = welch_uv2(x, raw_bi.info["sfreq"])
P_bi = P_bi.mean(0)
pk_bi = peak_above_trend(f_bi, P_bi)
print(f"ds-brain-invaders subject 1: {len(bi_channels)} EEG channels at "
      f"{raw_bi.info['sfreq']:.0f} Hz, {raw_bi.times[-1]:.0f} s, loaded in {time.time() - t0:.0f} s")
print(f"   channels: {', '.join(bi_channels)}")
print(f"   posterior picks used: {', '.join(chs)}")
print(f"   alpha peak {pk_bi['freq_hz']} Hz at {pk_bi['db_above_trend']:.2f} dB above trend")
print()
print("CAVEAT, and it is not small: this recording is a P300 game, eyes open and visually engaged, not a")
print("resting baseline.  Alpha is suppressed by exactly that.  The number is comparable across the")
print("SPECTRAL ESTIMATOR, not across the STATE -- which is the confound section 8's degradation removes.")

# --- ds-areeg: the 14-channel consumer headset ----------------------------------------------------
# One 32.2 MB zip on Mendeley Data is the whole dataset.  Rather than fetch it, the central directory
# is read over two HTTP range requests (about 125 kB) and only the members this notebook needs are
# pulled -- the same technique helpers_l4.aszed_index uses for ds-aszed, pointed at a different archive.
AREEG_URL = ("https://data.mendeley.com/public-files/datasets/7m472ykkx7/files/"
             "8b06aa16-8df4-436a-b95c-fc0f1d821be7/file_downloaded")
AREEG_ZIP_BYTES = 32_207_245           # Mendeley public API, record 7m472ykkx7 v1, read 2026-09-18
AREEG_CHANNELS = DEVICE_CHANNELS["ds-areeg"]
AREEG_SFREQ = 128.0

t0 = time.time()
areeg_ix = L4.aszed_index(AREEG_URL, size=AREEG_ZIP_BYTES, verbose=False)
members = [e for e in areeg_ix["entries"] if e["name"].lower().endswith(".csv")]
by_par = {}
for e in members:
    base = e["name"].split("/")[-1]
    m = re.match(r"par\.\s*(\d+)\s", base)
    by_par.setdefault(int(m.group(1)) if m else -1, []).append(e)
print(f"ds-areeg archive index: {len(areeg_ix['entries'])} members, {len(members)} CSV recordings, "
      f"read in {time.time() - t0:.1f} s over 2 range requests")
print(f"   participant ids parsed from filenames: "
      f"{sorted(k for k in by_par if k > 0)}")
print(f"   recordings per participant: "
      f"{ {k: len(v) for k, v in sorted(by_par.items()) if k > 0} }")
print(f"   filenames that do not match 'par.<N> ': {len(by_par.get(-1, []))}")
print()
print("TODO(confirm) -- a discrepancy with the descriptor, recorded rather than reconciled: the data")
print(f"   descriptor documents 22 participants and 352 recordings; this archive holds {len(members)} CSVs")
print(f"   whose filenames parse to participants 1-{max(k for k in by_par if k > 0)}, plus "
      f"{len(by_par.get(-1, []))} filenames that do not follow the pattern.  Filename-derived ids are not")
print("   a participant table, and this dataset ships no participants file (catalog: 'No formal")
print("   demographic/clinical metadata file beyond the in-paper participant table').")

def areeg_load(entry):
    """One ArEEG CSV, inflated from the archive in memory.  Returns (metadata, DataFrame)."""
    txt = L4.aszed_extract(areeg_ix, entry["name"]).decode("utf-8", "replace")
    head, rest = txt.split("\n", 1)
    meta = {k.strip(): v.strip() for k, v in
            (kv.split(":", 1) for kv in head.split(",") if ":" in kv)}
    return meta, pd.read_csv(io.StringIO(rest))


def areeg_participant(par, *, max_trials=None):
    """Every recording of one ArEEG participant, fetched in parallel (the archive host's latency,
    not its bandwidth, is the cost: ~2 s per request, ~90 kB each)."""
    ents = sorted(by_par[par], key=lambda e: e["name"])[:max_trials]
    with ThreadPoolExecutor(max_workers=6) as ex:
        out = list(ex.map(areeg_load, ents))
    return ents, out


t0 = time.time()
areeg_par = sorted(k for k in by_par if k > 0)[:N_AREEG]
areeg_data, areeg_rows, areeg_bytes = {}, [], 0
for par in areeg_par:
    ents, loaded = areeg_participant(par, max_trials=AREEG_TRIALS)
    areeg_bytes += sum(e["compressed_bytes"] for e in ents)
    segs, quality = [], []
    for meta, df in loaded:
        X = df[[f"EEG.{c}" for c in AREEG_CHANNELS]].to_numpy(float)
        segs.append(X.T)                                    # (14, n_times), microvolts
        quality.append(df)
    areeg_data[par] = {"segments": segs, "frames": quality, "meta": loaded[0][0]}
    # Per-segment Welch, then average: the exports are separate recordings, so concatenating them
    # would put 15 discontinuities inside the windows and spread broadband power everywhere.
    Ps = [welch_uv2(s - s.mean(1, keepdims=True), AREEG_SFREQ)[1] for s in segs]
    f_ar = welch_uv2(segs[0], AREEG_SFREQ)[0]
    P = np.mean(Ps, axis=0)
    posterior = [AREEG_CHANNELS.index(c) for c in ("O1", "O2", "P7", "P8")]
    pk = peak_above_trend(f_ar, P[posterior])
    areeg_data[par]["freqs"], areeg_data[par]["psd"] = f_ar, P
    areeg_rows.append({"participant": f"par.{par}", "n_recordings": len(segs),
                       "total_s": round(sum(s.shape[1] for s in segs) / AREEG_SFREQ, 1),
                       "iaf_hz": pk["freq_hz"], "db": round(pk["db_above_trend"], 2),
                       "slope_dB_per_decade": round(pk["slope_db_per_decade"], 2)})
areeg_tab = pd.DataFrame(areeg_rows)
print(f"ds-areeg: {len(areeg_tab)} participants, {int(areeg_tab.n_recordings.sum())} recordings, "
      f"{areeg_bytes / 1e6:.1f} MB fetched of a {AREEG_ZIP_BYTES / 1e6:.1f} MB archive, "
      f"{time.time() - t0:.0f} s")
print(areeg_tab.to_string(index=False))
print()
print(f"header of the first recording: {areeg_data[areeg_par[0]]['meta']}")
print()
print("What the 10-s imagined-speech period is, for this purpose: eyes closed, no visual input, a")
print("light cognitive task.  It is not a resting baseline and this notebook does not call it one.")
print("TODO(confirm): each export runs ~11.5 s against a protocol whose imagine period is 10 s; which")
print("part of the trial the export covers is not stated in the file, so the whole export is used.")

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.4))

ax = axes[0]
f_e, P_e = eegbci_psd[subjects[0]]
ax.semilogy(f_e, P_e, lw=1.4, label=f"ds-eegbci {subjects[0]} R02, 64 ch @ 160 Hz (eyes closed)")
ax.semilogy(f_bi, P_bi, lw=1.4, label="ds-brain-invaders sub 1, 16 ch dry @ 512 Hz (P300 game)")
par0 = areeg_par[0]
posterior = [AREEG_CHANNELS.index(c) for c in ("O1", "O2", "P7", "P8")]
ax.semilogy(areeg_data[par0]["freqs"], areeg_data[par0]["psd"][posterior].mean(0), lw=1.4,
            label=f"ds-areeg par.{par0}, 14 ch consumer @ 128 Hz (imagined speech, eyes closed)")
ax.axvspan(43, 64, color="0.85", zorder=0)
ax.text(45, ax.get_ylim()[1] * 0.3, "above the consumer\nheadset's documented\n~43 Hz ceiling",
        fontsize=7.5, va="top")
ax.set(xlim=(1, 70), xlabel="Frequency (Hz)", ylabel="Power spectral density (µV²/Hz)",
       title="Posterior spectra, three device classes (µV²/Hz vs Hz)")
ax.legend(fontsize=7.4, loc="lower left")
ax.grid(alpha=0.3, which="both")

ax = axes[1]
groups = [("ds-eegbci\n64 ch gel", [r for r in eegbci_tab.iaf_hz if r]),
          ("ds-brain-invaders\n16 ch dry", [pk_bi["freq_hz"]] if pk_bi["freq_hz"] else []),
          ("ds-areeg\n14 ch consumer", [r for r in areeg_tab.iaf_hz if r])]
for i, (name, vals) in enumerate(groups):
    if vals:
        ax.scatter(np.full(len(vals), i) + rng.normal(0, 0.05, len(vals)), vals, s=44, alpha=0.85)
        ax.hlines(np.median(vals), i - 0.22, i + 0.22, color="k", lw=2.2)
    else:
        ax.text(i, 10, "no peak\nabove 3 dB", ha="center", fontsize=8)
ax.set_xticks(range(3), [g[0] for g in groups], fontsize=8)
ax.set(ylabel="Individual alpha frequency (Hz)", ylim=(6.5, 13.5),
       title="Alpha peak, one estimator, three classes (Hz)\nblack bar = median")
ax.grid(alpha=0.3, axis="y")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

CEILING_SLOPE = -40.0        # dB per decade; steeper than any plausible aperiodic EEG background


def local_slopes(freqs, psd, *, width_oct=1.0, fmin=12.0, step=0.5):
    """Slope of log10 power against log10 frequency in sliding one-octave windows (dB/decade)."""
    freqs = np.asarray(freqs, float)
    psd = np.asarray(psd, float)
    out = []
    for fc in np.arange(fmin, freqs.max() / 2 ** (width_oct / 2) + 1e-9, step):
        lo, hi = fc / 2 ** (width_oct / 2), fc * 2 ** (width_oct / 2)
        m = (freqs >= lo) & (freqs <= hi) & (psd > 0)
        if m.sum() < 5:
            continue
        out.append((fc, float(np.polyfit(np.log10(freqs[m]), 10 * np.log10(psd[m]), 1)[0])))
    return np.asarray(out)


def sustained_ceiling(freqs, psd, *, threshold=CEILING_SLOPE, **kw):
    """Lowest frequency above which EVERY window is steeper than ``threshold``.

    One steep window is a notch or a noisy estimate; a ceiling is a slope that never comes back.
    Returns (frequency or None, the slope curve).
    """
    ls = local_slopes(freqs, psd, **kw)
    steep = (ls[:, 1] <= threshold).astype(int)
    sustained = np.flip(np.minimum.accumulate(np.flip(steep)))
    idx = np.where(sustained == 1)[0]
    return (float(ls[idx[0], 0]) if idx.size else None), ls


SPECTRA = {
    f"ds-eegbci {subjects[0]} R02 (64 ch, 160 Hz out)": eegbci_psd[subjects[0]],
    "ds-brain-invaders sub 1 (16 ch dry, 512 Hz out)": (f_bi, P_bi),
    f"ds-areeg par.{areeg_par[0]} (14 ch consumer, 128 Hz out)":
        (areeg_data[areeg_par[0]]["freqs"], areeg_data[areeg_par[0]]["psd"][posterior].mean(0)),
}
rows, slope_curves = [], {}
for label, (f, P) in SPECTRA.items():
    ceil, ls = sustained_ceiling(f, P)
    slope_curves[label] = ls
    at = {hz: (float(ls[np.argmin(np.abs(ls[:, 0] - hz)), 1]) if ls[:, 0].max() >= hz else np.nan)
          for hz in (20, 30, 40, 45)}
    rows.append({"recording": label, "Nyquist (Hz)": f.max(),
                 "sustained roll-off from (Hz)": ceil,
                 "slope @20 Hz": round(at[20], 1), "slope @30 Hz": round(at[30], 1),
                 "slope @40 Hz": round(at[40], 1), "slope @45 Hz": round(at[45], 1)})
roll = pd.DataFrame(rows)
print(f"Local spectral slope in sliding one-octave windows (dB per decade); a sustained slope past "
      f"{CEILING_SLOPE:g} dB/decade is the front end, not the brain:")
print(roll.to_string(index=False))
print()
print("ds-eegbci and ds-brain-invaders have NO sustained ceiling below their Nyquist frequencies, which is")
print("what the catalog says of both: no hardware filters on the first, no online filter on the second.")
print("Their slopes wander between about -10 and -50 dB/decade and always come back -- that is EEG plus")
print("muscle plus a mains line, not a filter.")
print()
ceil_ar = roll.iloc[2]["sustained roll-off from (Hz)"]
print(f"ds-areeg's slope leaves that range at {ceil_ar:g} Hz and never returns, reaching "
      f"{roll.iloc[2]['slope @45 Hz']:.0f} dB/decade by 45 Hz.")
print()
print("A DISAGREEMENT WITH A RECORDED FACT, reported rather than reconciled: data/directory.yaml records")
print(f"   'nominal 128 Hz output but ~= 43 Hz hardware bandwidth' for this device class, and the")
print(f"   sustained roll-off measured here begins at {ceil_ar:g} Hz -- lower than 43.  The two numbers are")
print("   not the same quantity: a manufacturer's 'bandwidth' figure is a single point on a response")
print("   (a -3 dB corner, usually) and this is the frequency past which the response is steepening")
print("   everywhere.  A corner at 43 Hz and a roll-off that has begun by 31 Hz are compatible.")
print("   TODO(confirm): which convention the 43 Hz figure uses.  What both agree on, and what the lesson")
print(f"   needs, is that the usable ceiling is FAR below the "
      f"{roll.iloc[2]['Nyquist (Hz)']:.0f} Hz Nyquist frequency the export's sampling rate advertises.")

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.3))
ax = axes[0]
for label, ls in slope_curves.items():
    ax.plot(ls[:, 0], ls[:, 1], lw=1.5, label=label)
ax.axhline(CEILING_SLOPE, color="0.4", lw=1.0, ls="--")
ax.text(13, CEILING_SLOPE + 4, f"{CEILING_SLOPE:g} dB/decade: steeper than any\naperiodic EEG background",
        fontsize=7.4, color="0.3")
ax.axvline(43, color="0.6", lw=1.0, ls=":")
ax.text(43.8, -175, "documented\n~43 Hz", fontsize=7.2, color="0.4")
ax.set(xlim=(12, 70), ylim=(-200, 60), xlabel="Frequency (Hz)",
       ylabel="Local spectral slope (dB per decade)",
       title="Where the amplifier takes over (dB/decade vs Hz)")
ax.legend(fontsize=7.2, loc="lower left")
ax.grid(alpha=0.3)

ax = axes[1]
f_ar = areeg_data[areeg_par[0]]["freqs"]
for par in areeg_par:
    ax.semilogy(f_ar, areeg_data[par]["psd"][posterior].mean(0), lw=1.3, label=f"par.{par}")
for hz in (50, 60):
    ax.axvline(hz, color="0.6", lw=0.9, ls=":")
ax.axvline(64, color="k", lw=1.0)
ax.text(63, 3e-3, "Nyquist of the 128 Hz export", fontsize=7.2, ha="right", rotation=90)
ax.set(xlim=(1, 64), xlabel="Frequency (Hz)", ylabel="Power spectral density (µV²/Hz)",
       title="ds-areeg posterior spectra to Nyquist (µV²/Hz vs Hz)")
ax.legend(fontsize=8)
ax.grid(alpha=0.3, which="both")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

# Is a mains notch visible, or has the roll-off already removed everything at 50 and 60 Hz?
print(f"ds-areeg, mains region.  Recording location: Syria (catalog); "
      f"data/directory.yaml records mains_hz: {DIRECTORY['ds-areeg'].get('mains_hz')}.")
print()
P_ar = np.mean([areeg_data[p]["psd"][posterior].mean(0) for p in areeg_par], axis=0)
tot = np.trapezoid(P_ar[(f_ar >= 1) & (f_ar <= 64)], f_ar[(f_ar >= 1) & (f_ar <= 64)])
for lo, hi in ((8, 12), (20, 25), (30, 35), (38, 43), (45, 50), (48, 52), (55, 62)):
    m = (f_ar >= lo) & (f_ar <= hi)
    band = np.trapezoid(P_ar[m], f_ar[m])
    print(f"   {lo:2d}-{hi:2d} Hz : {band:9.4f} uV^2 ({100 * band / tot:6.3f} % of 1-64 Hz)")
print()
print("A notch cannot be told apart from the roll-off here, and that is the honest reading: by 50 Hz the")
print("band-limited front end has already removed so much that there is nothing left for a notch to")
print("remove.  The catalog states the notches as a device specification; this recording neither")
print("confirms nor contradicts them, and TODO(confirm) is the right answer to 'can you see the notch'.")

frames = [df for par in areeg_par for df in areeg_data[par]["frames"]]
cols = frames[0].columns
fam = {"EEG (signal)": [c for c in cols if c.startswith("EEG.") and c.split(".", 1)[1] in AREEG_CHANNELS],
       "EEG (housekeeping)": [c for c in cols if c.startswith("EEG.")
                              and c.split(".", 1)[1] not in AREEG_CHANNELS],
       "CQ (contact quality)": [c for c in cols if c.startswith("CQ.")],
       "EQ (signal quality)": [c for c in cols if c.startswith("EQ.")],
       "MOT (motion)": [c for c in cols if c.startswith("MOT.")],
       "markers / timestamps": [c for c in cols if c.startswith(("Timestamp", "Original", "Marker"))]}
print(f"{len(cols)} columns in an ArEEG export:")
for k, v in fam.items():
    print(f"   {len(v):3d}  {k:22s} {', '.join(v[:6])}{' ...' if len(v) > 6 else ''}")
print(f"\n   {len(fam['EEG (signal)'])} of {len(cols)} columns "
      f"({100 * len(fam['EEG (signal)']) / len(cols):.0f} %) are electrode signal.")
print()
# The vendor criterion, computed.  EQ.OVERALL is written only when the firmware updates it.
eq = pd.concat([df["EQ.OVERALL"] for df in frames])
updates = eq.dropna()
print(f"EQ.OVERALL is present on {len(updates):,} of {len(eq):,} rows "
      f"({100 * len(updates) / len(eq):.1f} %) -- the firmware updates it about "
      f"{len(updates) / (len(eq) / 128.0):.1f} times a second, so it is a periodic report, not a "
      f"per-sample flag.")
print(f"   values seen: {sorted(updates.unique().tolist())}")
frac75 = float((updates >= 75).mean())
print(f"   fraction of quality updates meeting the ds-mpeng criterion (EQ.OVERALL >= 75 %): "
      f"{100 * frac75:.1f} %")
print()
print("Beside the recorded fact for ds-mpeng -- about 50 % of samples during intense video-game play --")
print(f"this imagined-speech task at {100 * frac75:.0f} % is the contrast the lesson wants: the same class")
print("of hardware and the same vendor criterion, with and without vigorous movement.  The two numbers")
print("are not measured the same way (samples there, firmware updates here) and are not interchangeable.")

# Per-electrode quality, and the one that matters most for this analysis.
per_ch = pd.DataFrame({c: pd.concat([df[f"EQ.{c}"] for df in frames]).dropna()
                       for c in AREEG_CHANNELS}).mean().sort_values()
print("Mean EQ.* per electrode over every loaded recording (vendor scale, higher is better):")
for c, v in per_ch.items():
    flag = "  <-- posterior, and the worst-scored electrode in the set" if c == per_ch.index[0] and c in ("O1", "O2", "P7", "P8") else ""
    print(f"   {c:4s} {v:5.2f}{flag}")
print()
worst = per_ch.index[0]
print(f"The vendor's own column flags {worst} as the least trustworthy electrode of the fourteen.")
if worst in ("O1", "O2", "P7", "P8"):
    print("It is also one of the four posterior electrodes this notebook measures alpha on.  That is the")
    print("whole argument for reading the quality columns BEFORE choosing a channel, rather than after a")
    print("result disappoints -- and it is a decision the 64-channel cap does not have to make, because")
    print("it has eleven posterior electrodes and can drop one.")
print()
mot = [c for c in cols if c.startswith("MOT.Acc")]
m0 = frames[0]
print(f"Motion block: {', '.join(mot)} is present on {100 * (1 - m0[mot[0]].isna().mean()):.0f} % of rows "
      f"-- the header declares 'sampling rate: {areeg_data[areeg_par[0]]['meta'].get('sampling rate')}', so "
      f"the motion channels run at a quarter of the EEG rate and the column is padded with blanks.")
print("Reading this CSV with a naive 'drop rows with missing values' throws away three quarters of the EEG.")

MI_SUBJECT = "S001"
t0 = time.time()
paths = L4.fetch_eegbci(MI_SUBJECT, L4.MI_RUNS, verbose=False)
existing_then_add(paths, dl)
epochs_mi, mi_info = L4.load_imagery_epochs(MI_SUBJECT, L4.MI_RUNS, reference="average")
print(f"ds-eegbci {MI_SUBJECT} runs {'+'.join(mi_info['runs'])}: {mi_info['n_eeg']} EEG channels at "
      f"{mi_info['sfreq']:.0f} Hz, epochs {mi_info['n_epochs']}, {time.time() - t0:.0f} s")
print(f"   conditions: T1 = left-fist imagery, T2 = right-fist imagery (the EDF's own annotations)")
print(f"   reference: {mi_info['reference']} over all {mi_info['n_eeg']} channels")

CH_MI = ("C3", "C4")
t0 = time.time()
tf = L4.morlet_power(epochs_mi["T2"], picks=CH_MI)
erd64 = {ch: L4.band_erd(tf["power"][:, i], tf["freqs"], tf["times"], L4.MU_BAND,
                         baseline=L4.TF_BASELINE, active=L4.TF_ACTIVE, order="band-first")
         for i, ch in enumerate(CH_MI)}
KEY_44 = {"C3": -47.87, "C4": -10.46}
print(f"\nmu ERD (8-13 Hz, band-first, baseline {L4.TF_BASELINE}, active {L4.TF_ACTIVE}), T2, "
      f"{mi_info['n_eeg']}-channel average reference, {time.time() - t0:.0f} s:")
for ch in CH_MI:
    d = erd64[ch] - KEY_44[ch]
    print(f"   {ch}: {erd64[ch]:+7.2f} %   nb-4-4 key {KEY_44[ch]:+7.2f} %   difference {d:+.3f} pp "
          f"{'-- reproduced' if abs(d) < 0.05 else '-- DISAGREEMENT, see the notes'}")
LAT64 = erd64["C3"] - erd64["C4"]
print(f"   lateralisation (C3 - C4) for right-fist imagery: {LAT64:+.2f} pp "
      f"(negative = more desynchronisation contralateral to the imagined hand, which is the expected sign)")

import helpers_l3 as L3

P3_CH, P3_WINDOW = L3.P3_CHANNEL, L3.P3_WINDOW
raw_bi_f = raw_bi.copy()
raw_bi_f.filter(L3.HP_HZ, L3.LP_HZ, picks="eeg", method="fir", fir_design="firwin", phase="zero",
                verbose=False)
events_bi = mne.find_events(raw_bi_f, stim_channel="STI 014", shortest_event=1, verbose=False)
ep_all = mne.Epochs(raw_bi_f, events_bi, helpers.BI2014A_EVENT_ID, tmin=L3.EPOCH_TMIN, tmax=L3.EPOCH_TMAX,
                    baseline=L3.EPOCH_BASELINE, picks="eeg", preload=True, verbose=False)
d_all = ep_all.get_data() * 1e6
ptp_all = (d_all.max(2) - d_all.min(2)).max(1)
keep_all = ptp_all <= L3.REJECT_PTP_UV
ep_clean = ep_all[np.where(keep_all)[0]]


def p300(epochs, *, ch=None, window=None):
    """Target minus non-target mean amplitude (uV) at one electrode, with its SME and trial counts."""
    ch = P3_CH if ch is None else ch
    window = P3_WINDOW if window is None else window
    if ch not in epochs.ch_names:
        return {"amp_uv": np.nan, "sme_uv": np.nan, "evoked": None,
                "n": {c: len(epochs[c]) for c in ("Target", "NonTarget")}}
    ev = {c: epochs[c].average() for c in ("Target", "NonTarget")}
    diff = mne.combine_evoked([ev["Target"], ev["NonTarget"]], weights=[1, -1])
    m = (diff.times >= window[0]) & (diff.times <= window[1])
    per = {c: epochs[c].get_data(picks=[ch])[:, 0, :][:, m].mean(1) * 1e6 for c in ev}
    return {"amp_uv": float(L3.mean_amplitude(diff, ch, window)),
            "sme_uv": float(np.sqrt(per["Target"].var(ddof=1) / len(per["Target"])
                                    + per["NonTarget"].var(ddof=1) / len(per["NonTarget"]))),
            "evoked": diff, "n": {c: len(per[c]) for c in per}}


p3_ref = p300(ep_clean)
diff_clean = mne.combine_evoked([ep_clean["Target"].average(), ep_clean["NonTarget"].average()],
                                weights=[1, -1])
KEY_32 = 1.43
print(f"ds-brain-invaders subject 1, the nb-3-2 pipeline: {L3.HP_HZ:g}-{L3.LP_HZ:g} Hz, epochs "
      f"{L3.EPOCH_TMIN:g}..{L3.EPOCH_TMAX:g} s, baseline {L3.EPOCH_BASELINE}, "
      f"{L3.REJECT_PTP_UV:g} uV peak-to-peak over all 16 channels, RECORDING REFERENCE KEPT")
print(f"   rejected {int((~keep_all).sum())} of {len(keep_all)} epochs ({100 * (~keep_all).mean():.0f} %) "
      f"-- kept {p3_ref['n']['Target']} Target / {p3_ref['n']['NonTarget']} NonTarget")
print(f"   P300 at {P3_CH}, {P3_WINDOW[0] * 1000:.0f}-{P3_WINDOW[1] * 1000:.0f} ms: "
      f"{p3_ref['amp_uv']:+.3f} uV   nb-3-2 key {KEY_32:+.2f} uV   "
      f"difference {p3_ref['amp_uv'] - KEY_32:+.3f} uV "
      f"{'-- reproduced' if abs(p3_ref['amp_uv'] - KEY_32) < 0.006 else '-- DISAGREEMENT, see the notes'}")
print(f"   nb-3-2 also records that Pz is not this subject's best site, and it is not: "
      f"P4 {L3.mean_amplitude(diff_clean, 'P4', P3_WINDOW):+.2f} uV, "
      f"Cz {L3.mean_amplitude(diff_clean, 'Cz', P3_WINDOW):+.2f} uV, "
      f"P3 {L3.mean_amplitude(diff_clean, 'P3', P3_WINDOW):+.2f} uV")

# Two decisions, each of which flips the sign of this component on this recording.
def rereference(epochs):
    """Average over the channels that remain, applied to the data (not as a projection)."""
    d = epochs.get_data(copy=True)
    d = d - d.mean(axis=1, keepdims=True)
    return mne.EpochsArray(d, epochs.info.copy(), events=epochs.events, event_id=epochs.event_id,
                           tmin=epochs.tmin, verbose=False)


variants = [
    ("recording reference, 150 uV rejection  (nb-3-2)", p300(ep_clean)),
    ("recording reference, NO rejection", p300(ep_all)),
    ("16-channel average reference, 150 uV rejection", p300(rereference(ep_clean))),
    ("16-channel average reference, NO rejection", p300(rereference(ep_all))),
]
print(f"The same flashes, the same electrode, the same window -- {P3_CH}, "
      f"{P3_WINDOW[0] * 1000:.0f}-{P3_WINDOW[1] * 1000:.0f} ms, target minus non-target:\n")
print(f"{'pipeline':50s} {'P300 (uV)':>10s} {'SME (uV)':>9s} {'target n':>9s}")
for name, r in variants:
    print(f"{name:50s} {r['amp_uv']:+10.3f} {r['sme_uv']:9.3f} {r['n']['Target']:9d}")
print()
print("Two independent decisions each change the SIGN of a P300 on this recording, and neither is a")
print("mistake a reader would notice in a methods section:")
print("  - Artifact rejection.  57 % of these dry-electrode epochs exceed 150 uV against 6 % of ERP")
print("    CORE's gel epochs (nb-3-2).  Without rejection the average is dominated by what the dry")
print("    contacts did, not by what the brain did.")
print("  - The reference.  An average over 64 electrodes approximates an inactive reference; an average")
print("    over 16 that are mostly frontal and temporal does not, and a P300 is broad enough that much")
print("    of it lands in the reference and is subtracted from Pz.")
print()
print("This is the low-channel lesson in its sharpest form, and it is NOT a claim that dry electrodes")
print("are worse: it is that a pipeline transported unchanged from a 64-channel gel cap can invert a")
print("component on 16 dry ones, and the montage is the reason.  pf-reference-changes-everything.")

SUBSETS = {
    "64 (full cap)": None,
    "19 (10-20 clinical)": ("Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "T7", "C3", "Cz", "C4", "T8",
                            "P7", "P3", "Pz", "P4", "P8", "O1", "O2"),
    "8 (task-matched)": ("F3", "F4", "C3", "Cz", "C4", "P3", "O1", "O2"),
    "4 (task-matched)": ("C3", "C4", "O1", "O2"),
    "4 (consumer-shaped)": ("AF3", "AF4", "O1", "O2"),
}

# The alpha peak is measured on the SAME eyes-closed baseline section 3 used (R02), so that the only
# thing changing here is the montage -- not the state, the run or the window.
raw_ec = helpers.load_spine("ds-eegbci", MI_SUBJECT, "R02")

rows = []
for name, chs in SUBSETS.items():
    # --- alpha peak on the eyes-closed baseline, reference re-derived over the retained channels ---
    r = raw_ec.copy()
    if chs is not None:
        r.pick([c for c in chs if c in r.ch_names])
    x = r.get_data(picks="eeg") * 1e6
    x = x - x.mean(axis=0, keepdims=True)
    post = [i for i, c in enumerate(r.copy().pick("eeg").ch_names) if c in POSTERIOR]
    f_d, P_d = welch_uv2(x[post] if post else x, r.info["sfreq"])
    pk = peak_above_trend(f_d, P_d)
    # --- mu ERD on the imagery epochs, reference re-derived the same way ---
    ep = epochs_mi.copy()
    if chs is not None:
        ep.pick([c for c in chs if c in ep.ch_names])
    ep_r = rereference(ep)
    row = {"montage": name, "n_channels": len(ep_r.ch_names), "posterior ch": len(post),
           "alpha peak (Hz)": pk["freq_hz"], "peak height (dB)": round(pk["db_above_trend"], 2)}
    if all(c in ep_r.ch_names for c in CH_MI):
        t = L4.morlet_power(ep_r["T2"], picks=CH_MI)
        e = {ch: L4.band_erd(t["power"][:, i], t["freqs"], t["times"], L4.MU_BAND,
                             baseline=L4.TF_BASELINE, active=L4.TF_ACTIVE, order="band-first")
             for i, ch in enumerate(CH_MI)}
        row |= {"mu ERD C3 (%)": round(e["C3"], 2), "mu ERD C4 (%)": round(e["C4"], 2),
                "lateralisation C3-C4 (pp)": round(e["C3"] - e["C4"], 2)}
    else:
        row |= {"mu ERD C3 (%)": np.nan, "mu ERD C4 (%)": np.nan,
                "lateralisation C3-C4 (pp)": np.nan}
    rows.append(row)
deg = pd.DataFrame(rows)
print(f"Controlled degradation -- ds-eegbci {MI_SUBJECT}, identical recordings, identical pipeline, "
      f"only the montage changes.")
print(f"   alpha peak: run R02 (eyes closed), the posterior electrodes that survive the subset")
print(f"   mu ERD    : runs {'+'.join(mi_info['runs'])}, condition T2, 8-13 Hz, band-first\n")
print(deg.to_string(index=False))
edge = [r["montage"] for _, r in deg.iterrows() if r["alpha peak (Hz)"] is None]
print(f"\nband-edge / threshold check: "
      f"{'every montage produced a real peak inside the band' if not edge else 'NO PEAK REPORTED for ' + ', '.join(edge) + ' -- either below ' + str(MIN_DB) + ' dB or on the band limit, and the estimator says absent rather than inventing a value'}")

base = deg.iloc[0]
print("What moved, and by how much, against the full cap:\n")
for _, r in deg.iloc[1:].iterrows():
    print(f"{r['montage']}:")
    if not np.isnan(r["mu ERD C3 (%)"]):
        print(f"   mu ERD C3   {r['mu ERD C3 (%)']:+7.2f} % vs {base['mu ERD C3 (%)']:+7.2f} %  "
              f"({r['mu ERD C3 (%)'] - base['mu ERD C3 (%)']:+.2f} pp, "
              f"{100 * abs(r['mu ERD C3 (%)'] - base['mu ERD C3 (%)']) / abs(base['mu ERD C3 (%)']):.0f} % "
              f"of the full-cap value)")
        print(f"   mu ERD C4   {r['mu ERD C4 (%)']:+7.2f} % vs {base['mu ERD C4 (%)']:+7.2f} %  "
              f"({r['mu ERD C4 (%)'] - base['mu ERD C4 (%)']:+.2f} pp)")
        print(f"   C3 - C4     {r['lateralisation C3-C4 (pp)']:+7.2f} pp vs "
              f"{base['lateralisation C3-C4 (pp)']:+7.2f} pp   "
              f"(sign {'PRESERVED' if np.sign(r['lateralisation C3-C4 (pp)']) == np.sign(base['lateralisation C3-C4 (pp)']) else 'FLIPPED'})")
    else:
        print("   mu ERD      NOT COMPUTABLE -- this montage has no C3/C4")
    if r["alpha peak (Hz)"] is not None and base["alpha peak (Hz)"] is not None:
        print(f"   alpha peak  {r['alpha peak (Hz)']:.2f} Hz vs {base['alpha peak (Hz)']:.2f} Hz  "
              f"({r['alpha peak (Hz)'] - base['alpha peak (Hz)']:+.2f} Hz) from "
              f"{r['posterior ch']} posterior electrode(s); height "
              f"{r['peak height (dB)']:.2f} dB vs {base['peak height (dB)']:.2f} dB")
    print()
print("The peak FREQUENCY is the robust one, and the reason is worth stating: a frequency does not care")
print("how many electrodes were averaged into the reference, as long as the rhythm is on one of them.")
print("Its HEIGHT above the aperiodic background does move, because that is an amplitude.")
print("The ERD percentages move too -- most at four channels, where C3 and C4 are half the reference and")
print("a quarter of each channel's own signal is subtracted from itself.  The lateralisation keeps its")
print("sign at every montage that has both electrodes, which is the useful part: a low-channel system")
print("can say WHICH side desynchronised and cannot say BY HOW MUCH on a research cap's scale.")