# 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', 'pyprep')
_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 "pyprep" in _missing:
        _cmd += ["pyprep>=0.9"]
    subprocess.check_call(_cmd)

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

# 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

# Warnings are worth reading, so they are not silenced -- but their default format prints the
# absolute path of the file that raised them, which is nobody else's business and would put this
# machine's directory layout into the saved outputs.  Only the class and the message are shown.
warnings.formatwarning = lambda message, category, *a, **k: f"{category.__name__}: {message}\n"

mne.viz.set_browser_backend("matplotlib", verbose=False)
mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers imported from notebooks/_shared")
print("ERP CORE downloads (~56 MB per subject) are cached under data/downloads/erpcore/ in a repository "
      "clone, otherwise under MNE's data directory; nothing is re-fetched.")

SUBJECTS = ["sub-001", "sub-002", "sub-003"]
RESAMPLE_HZ = 256.0

raws, facts = {}, {}
for sid in SUBJECTS:
    raws[sid], facts[sid] = l2.load_erpcore("P3", sid, resample_hz=RESAMPLE_HZ, verbose=True)

f0 = facts[SUBJECTS[0]]
print(f"\nfrom the dataset's own sidecar: {f0['eeg_json']}")
print(f"EEG channels ({len(f0['eeg_channels'])}): {f0['eeg_channels']}")
print(f"EOG channels typed: {f0['eog_channels']}; units column: {f0['units']}")

raw0 = raws[SUBJECTS[0]]
eeg_names = [raw0.ch_names[i] for i in mne.pick_types(raw0.info, eeg=True)]
pos = raw0.get_montage().get_positions()["ch_pos"]
xyz = np.array([pos[c] for c in eeg_names])

no_position = [c for c in eeg_names if c not in pos or not np.isfinite(pos[c]).all()]
d = np.linalg.norm(xyz[:, None, :] - xyz[None, :, :], axis=-1)
np.fill_diagonal(d, np.inf)
duplicates = [(eeg_names[i], eeg_names[j]) for i, j in zip(*np.where(d < 1e-6)) if i < j]
radius = np.linalg.norm(xyz, axis=1)

print(f"montage checks for {SUBJECTS[0]} ({len(eeg_names)} EEG channels, {facts[SUBJECTS[0]]['montage']}):")
print(f"  channels without a position:            {no_position or 'none'}")
print(f"  channel pairs sharing a position:       {duplicates or 'none'}")
print(f"  radius from the head centre (m):        min {radius.min():.3f}, max {radius.max():.3f} "
      f"(a template head is a sphere of about 0.095 m)")
print(f"  closest neighbour distance (m):         min {d.min():.3f} ({eeg_names[int(np.unravel_index(d.argmin(), d.shape)[0])]}"
      f"-{eeg_names[int(np.unravel_index(d.argmin(), d.shape)[1])]}), "
      f"median of each channel's nearest neighbour {np.median(d.min(axis=1)):.3f}")
print(f"  channels the montage did not know:      {facts[SUBJECTS[0]]['montage_missing'] or 'none (the 3 EOG channels are typed eog, not eeg)'}")
print("\nThat median nearest-neighbour distance is the fact behind everything below: with 30 electrodes over "
      "the whole head, a channel's nearest neighbour is several centimetres away, so neighbouring channels "
      "are genuinely less correlated than they would be on a 64- or 128-channel cap.")

fig, axes = plt.subplots(1, 2, figsize=(11, 4.4), gridspec_kw=dict(width_ratios=[1, 1.5]))
raw0.plot_sensors(show_names=True, axes=axes[0], show=False)
axes[0].set_title(f"{SUBJECTS[0]}: 30 EEG positions, nose up (standard_1005)")
helpers.plot_traces(raw0, ["Fp1", "Fz", "Cz", "Pz", "P9", "P10", "O1", "Oz"], t0=60, duration=10,
                    spacing_uV=150, ax=axes[1], title=f"{SUBJECTS[0]}, 10 s from 60 s (uV)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

MANUAL = dict(z_sd=3.5, excursion_uv=250.0, flat_uv=0.5, highpass_hz=1.0)


def manual_flags(raw, *, z_sd=MANUAL["z_sd"], excursion_uv=MANUAL["excursion_uv"],
                 flat_uv=MANUAL["flat_uv"], highpass_hz=MANUAL["highpass_hz"]):
    """The stated by-eye rule, as code. Returns {criterion: [channels]} plus the per-channel numbers."""
    r = raw.copy().pick("eeg").filter(highpass_hz, None, verbose=False)
    x = r.get_data() * 1e6
    names = r.ch_names
    sd = x.std(axis=1)
    exc = np.percentile(np.abs(x), 99.9, axis=1)
    mad = 1.4826 * np.median(np.abs(sd - np.median(sd)))
    z = (sd - np.median(sd)) / mad if mad > 0 else np.zeros_like(sd)
    flags = {
        "manual_amplitude": [names[i] for i in np.where(np.abs(z) > z_sd)[0]],
        "manual_excursion": [names[i] for i in np.where(exc > excursion_uv)[0]],
        "manual_flat": [names[i] for i in np.where(sd < flat_uv)[0]],
    }
    metrics = {names[i]: dict(std_uv=float(sd[i]), robust_z=float(z[i]), p999_abs_uv=float(exc[i]))
               for i in range(len(names))}
    return flags, metrics


manual, metrics = {}, {}
for sid in SUBJECTS:
    manual[sid], metrics[sid] = manual_flags(raws[sid])
    flagged = sorted({c for v in manual[sid].values() for c in v})
    print(f"{sid}: manual flags {flagged or 'none'}  " +
          "  ".join(f"{k.replace('manual_', '')}={v}" for k, v in manual[sid].items() if v))

sid = SUBJECTS[0]
order = sorted(metrics[sid], key=lambda c: -metrics[sid][c]["std_uv"])
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
axes[0].bar(range(len(order)), [metrics[sid][c]["std_uv"] for c in order], color="tab:blue")
axes[0].axhline(np.median([m["std_uv"] for m in metrics[sid].values()]), color="k", lw=0.8, label="median")
axes[0].set(xticks=range(len(order)), ylabel="standard deviation (uV)",
            title=f"{sid}: per-channel amplitude after a 1 Hz high-pass (uV)")
axes[0].set_xticklabels(order, rotation=90, fontsize=7)
axes[0].legend(fontsize=8)
axes[0].grid(alpha=0.3, axis="y")
zs = np.array([metrics[sid][c]["robust_z"] for c in order])
axes[1].bar(range(len(order)), zs, color=["tab:orange" if abs(v) > MANUAL["z_sd"] else "0.6" for v in zs])
axes[1].axhline(MANUAL["z_sd"], color="tab:red", lw=0.8)
axes[1].axhline(-MANUAL["z_sd"], color="tab:red", lw=0.8)
axes[1].set(xticks=range(len(order)), ylabel="robust z of the standard deviation (dimensionless)",
            title=f"{sid}: robust z, flagged above |z| = {MANUAL['z_sd']:g}")
axes[1].set_xticklabels(order, rotation=90, fontsize=7)
axes[1].grid(alpha=0.3, axis="y")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

import time

detection, detection_after_lowpass = {}, {}
t0 = time.time()
for sid in SUBJECTS:
    detection[sid] = l2.detect_bad_channels(raws[sid], highpass_hz=1.0, ransac=True, seed=l2.SEED)
    lp = raws[sid].copy().filter(0.1, 30.0, picks=["eeg", "eog"], verbose=False)
    detection_after_lowpass[sid] = l2.detect_bad_channels(lp, highpass_hz=None, ransac=True, seed=l2.SEED)
print(f"pyprep NoisyChannels on {len(SUBJECTS)} subjects, twice each: {time.time() - t0:.0f} s\n")

for sid in SUBJECTS:
    print(f"{sid} -- detection on the 1 Hz high-passed copy (the pipeline's setting)")
    for crit, chs in detection[sid]["flags"].items():
        print(f"    {crit:22s} {l2.BAD_CRITERIA.get(crit, ''):.58s}\n        -> {chs}")
    print(f"    channels flagged by more than one criterion: "
          f"{[c for c, cr in detection[sid]['by_channel'].items() if len(cr) > 1] or 'none'}")
    print(f"{sid} -- detection after a 0.1-30 Hz band-pass (the wrong order)")
    print(f"    flagged: {detection_after_lowpass[sid]['all_flagged'] or 'none'}  "
          f"(criteria that fired: {sorted(detection_after_lowpass[sid]['flags'])})")
    print()
print("The high-frequency-noise criterion cannot fire once a 30 Hz low-pass has removed the evidence, and the "
      "neighbour-correlation criterion looks much better on smoothed data. Detect first, then filter (L2.8).")

rows = []
for sid in SUBJECTS:
    man = {c for v in manual[sid].values() for c in v}
    auto = set(detection[sid]["all_flagged"])
    names = set(m for m in metrics[sid])
    agree = len(names - (man ^ auto))
    rows.append({"subject": sid, "manual": sorted(man) or ["-"], "pyprep": sorted(auto) or ["-"],
                 "both": sorted(man & auto) or ["-"], "pyprep only": sorted(auto - man) or ["-"],
                 "manual only": sorted(man - auto) or ["-"],
                 "per-channel agreement": f"{100 * agree / len(names):.0f} %"})
print(l2.fmt_table(rows, ["subject", "manual", "pyprep", "both", "pyprep only", "manual only",
                          "per-channel agreement"]))
print()
print("The disagreements are informative, not embarrassing, and they run in both directions.")
print("  - The manual rule flags Fp1/Fp2 on two subjects. Those channels are not bad: they are frontal, and "
      "blinks are enormous there, so an amplitude rule measures the participant rather than the electrode. "
      "That is a false alarm no spatial criterion makes, and it is a reason not to interpolate on amplitude alone.")
print("  - pyprep flags posterior channels (P3, PO3, PO7, O1) that the amplitude rule finds unremarkable. Some "
      "of that is real -- an electrode recording something the rest of the head is not -- and some of it is the "
      "montage: section 1 measured a median nearest-neighbour distance of several centimetres, so a correlation "
      "cutoff tuned on denser caps fires more often here.")
print("  - Neither pass can see a channel that is clean but wrong (bridging, section 8).")

decisions = {}
for sid in SUBJECTS:
    n_eeg = len(mne.pick_types(raws[sid].info, eeg=True))
    dec = l2.select_bads(detection[sid], n_eeg, min_criteria=2, max_fraction=0.10)
    # The spatial clause: how close are the flagged channels to one another?
    p = raws[sid].get_montage().get_positions()["ch_pos"]
    fl = detection[sid]["all_flagged"]
    contiguous = 0
    if len(fl) > 1:
        dd = np.array([[np.linalg.norm(p[a] - p[b]) for b in fl] for a in fl])
        np.fill_diagonal(dd, np.inf)
        contiguous = int((dd.min(axis=1) < 0.045).sum())     # within ~4.5 cm: adjacent on this montage
    dec["contiguous_flagged"] = contiguous
    dec["measurement_channel_flagged"] = l2.P3_CHANNEL in fl
    decisions[sid] = dec
    print(f"{sid}: flagged {len(fl):2d} -> interpolate {dec['bads'] or 'none'} "
          f"({len(dec['bads'])}/{n_eeg} = {len(dec['bads']) / n_eeg:.0%}, cap {dec['cap']}); "
          f"kept despite a flag: {dec['flagged_not_interpolated'] or 'none'}")
    print(f"    spatial clause: {contiguous} of the flagged channels have another flagged channel within 4.5 cm")
    print(f"    ROI clause:     the measurement channel {l2.P3_CHANNEL} "
          f"{'IS flagged -- this needs a stated decision' if dec['measurement_channel_flagged'] else 'is not flagged'}")
    print(f"    over the cap:   {dec['over_cap']}")
print(f"\nRule in force: {decisions[SUBJECTS[0]]['rule']}")

# Use the subject with the most convincing flag; if none has one, use the first.
SID_DEMO = next((s for s in SUBJECTS if decisions[s]["bads"]), SUBJECTS[0])
BAD = decisions[SID_DEMO]["bads"][0] if decisions[SID_DEMO]["bads"] else detection[SID_DEMO]["all_flagged"][0]
print(f"demonstration subject {SID_DEMO}, bad channel {BAD} "
      f"(criteria: {detection[SID_DEMO]['by_channel'][BAD]})")

base = raws[SID_DEMO].copy().filter(0.1, 30.0, picks=["eeg", "eog"], verbose=False)
eeg = [base.ch_names[i] for i in mne.pick_types(base.info, eeg=True)]
good = [c for c in eeg if c != BAD]

a = base.copy().set_eeg_reference("average", verbose=False)
b = base.copy().set_eeg_reference(good, verbose=False)
c = base.copy()
c.info["bads"] = [BAD]
c.interpolate_bads(reset_bads=True, verbose=False)
c.set_eeg_reference("average", verbose=False)

probe = l2.P3_CHANNEL if l2.P3_CHANNEL != BAD else "Cz"
xa, xb, xc = (r.get_data(picks=[probe])[0] * 1e6 for r in (a, b, c))
print(f"\nat the good channel {probe}, over the whole recording:")
print(f"  (a) bad channel in the average   vs (b) bad channel excluded: "
      f"RMS difference {np.sqrt(((xa - xb) ** 2).mean()):6.2f} uV, max |difference| {np.abs(xa - xb).max():7.2f} uV")
print(f"  (c) interpolated, then averaged  vs (b) bad channel excluded: "
      f"RMS difference {np.sqrt(((xc - xb) ** 2).mean()):6.2f} uV, max |difference| {np.abs(xc - xb).max():7.2f} uV")
print(f"  for scale, {probe}'s own RMS in (b) is {np.sqrt((xb ** 2).mean()):.2f} uV")

t0 = 60.0
sl = slice(int(t0 * base.info["sfreq"]), int((t0 + 8) * base.info["sfreq"]))
t = base.times[sl]
fig, axes = plt.subplots(2, 1, figsize=(12, 6.5), sharex=True)
axes[0].plot(t, base.get_data(picks=[BAD])[0][sl] * 1e6, "k", lw=0.7, label=f"{BAD} (flagged)")
axes[0].plot(t, base.get_data(picks=[probe])[0][sl] * 1e6, color="tab:blue", lw=0.7, label=f"{probe} (good)")
axes[0].set(ylabel="Amplitude (uV)", title=f"{SID_DEMO}: the flagged channel next to a good one, recorded reference (uV)")
axes[0].legend(fontsize=8); axes[0].grid(alpha=0.3)
axes[1].plot(t, xb[sl], color="0.55", lw=1.1, label="(b) average over the good channels")
axes[1].plot(t, xa[sl], color="tab:red", lw=0.9, label="(a) average including the flagged channel")
axes[1].plot(t, xc[sl], color="tab:green", lw=0.9, ls="--", label="(c) interpolate, then average")
axes[1].set(xlabel="Time (s)", ylabel="Amplitude (uV)",
            title=f"{probe} under three references (uV, positive up)")
axes[1].legend(fontsize=8); axes[1].grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

def rank_of(r, label):
    est = mne.compute_rank(r, rank="info", verbose=False)["eeg"]
    data_rank = int(np.linalg.matrix_rank(r.get_data(picks="eeg"), tol=1e-6 * float(np.abs(r.get_data(picks="eeg")).max())))
    return {"stage": label, "n_eeg_rows": len(mne.pick_types(r.info, eeg=True)),
            "rank (compute_rank, 'info')": est, "rank (numerical, from the data)": data_rank}


n_eeg = len(mne.pick_types(base.info, eeg=True))
stages = [rank_of(base, "as loaded (recorded CMS reference)")]
stages.append(rank_of(c.copy(), f"1 channel interpolated + average reference"))
print(l2.fmt_table(stages))
print()
for sid in SUBJECTS:
    k = len(decisions[sid]["bads"])
    before = l2.data_rank(n_eeg)
    after = l2.data_rank(n_eeg, n_interpolated=k, average_reference=True)
    print(f"{sid}: rank before {before['rank']}  ->  after {after['rank']}   ({after['arithmetic']})")
print()
print(f"The arithmetic ({l2.data_rank(n_eeg, n_interpolated=1, average_reference=True)['arithmetic']}) matches "
      f"the numerical rank of the data ({stages[1]['rank (numerical, from the data)']}) and not what "
      f"compute_rank(rank='info') reports ({stages[1][chr(114) + 'ank (compute_rank, ' + chr(39) + 'info' + chr(39) + ')']}): "
      "the Info object records channels and projectors, not the fact that one channel is now a linear "
      "combination of the others and that the rows sum to zero.")
print("The array keeps 30 rows throughout, so nothing warns you. A script that asks for 30 ICA components from "
      "this data is asking for dimensions that no longer exist (pf-interpolation-rank, L2.6). Pass the number: "
      "ICA(n_components=rank).")

def bridge_check(raw, *, corr_min=0.99, ed_max=0.05, n_report=3):
    r = raw.copy().pick("eeg").filter(1.0, 40.0, verbose=False)
    x = r.get_data() * 1e6
    names = r.ch_names
    corr = np.corrcoef(x)
    ed = np.var(x[:, None, :] - x[None, :, :], axis=-1)      # electrical distance (uV^2)
    med = np.median(ed[np.triu_indices_from(ed, 1)])
    ed_n = ed / med
    iu = np.triu_indices_from(corr, 1)
    order = np.argsort(ed_n[iu])
    closest = [{"pair": f"{names[iu[0][k]]}-{names[iu[1][k]]}", "r": float(corr[iu][k]),
                "ed_over_median": float(ed_n[iu][k]),
                "bridged": bool(corr[iu][k] > corr_min and ed_n[iu][k] < ed_max)}
               for k in order[:n_report]]
    n_bridged = int(((corr[iu] > corr_min) & (ed_n[iu] < ed_max)).sum())
    return closest, n_bridged, float(np.median(corr[iu])), float(np.max(corr[iu])), med


BRIDGE_RULE = "r > 0.99 AND electrical distance < 5 % of the median over all pairs"
total_bridged = 0
for sid in SUBJECTS:
    closest, n_bridged, med_r, max_r, med_ed = bridge_check(raws[sid])
    total_bridged += n_bridged
    print(f"{sid}: median pair correlation {med_r:+.2f}, largest {max_r:+.2f}, median electrical distance "
          f"{med_ed:.0f} uV^2; pairs meeting the rule: {n_bridged}")
    print("    the three electrically closest pairs: " +
          ", ".join(f"{p['pair']} (r {p['r']:+.2f}, ED {100 * p['ed_over_median']:.0f} % of median"
                    f"{', BRIDGED' if p['bridged'] else ''})" for p in closest))
print(f"\nRule used here: {BRIDGE_RULE}. It is a screening rule chosen for this notebook, not a published "
      "threshold -- TODO(confirm) against a bridging-detection reference before quoting it.")
print(f"Pairs meeting it across the three recordings: {total_bridged}. " +
      ("None, which is the expected result for a well-run lab recording."
       if total_bridged == 0 else
       "Each is a neighbouring pair and needs a person to look before anything is done about it."))
print("Note which way the numbers point: the closest pairs are all physically adjacent electrodes, where a high "
      "correlation is exactly what should happen. That is why bridging cannot be caught by a correlation "
      "threshold alone, why no 'bad channel' criterion will ever fire on it, and why the real fix is at the cap "
      "before recording (pf-bridged-electrodes).")

print("nb-2-2-bad-channels -- L2.2 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-erpcore P3, {', '.join(SUBJECTS)} (CC BY-SA 4.0 per data/directory.yaml, contested at source; open; per-subject "
      f"downloadable). 30 EEG + 3 EOG, 1024 Hz resampled to {RESAMPLE_HZ:g} Hz, CMS reference, no software filters.")
print(f"Detector: pyprep NoisyChannels.find_all_bads(ransac=True) on a 1 Hz high-passed, not low-passed copy, "
      f"random_state {l2.SEED}. Manual rule: {MANUAL}.")
print(f"Interpolation rule: {decisions[SUBJECTS[0]]['rule']}.")
print()
print("Flagged channels per subject, per criterion:")
for sid in SUBJECTS:
    print(f"  {sid}")
    for crit, chs in detection[sid]["flags"].items():
        print(f"      {crit:22s} {chs}")
    print(f"      {'manual (stated rule)':22s} {sorted({c for v in manual[sid].values() for c in v}) or []}")
    print(f"      {'-> interpolated':22s} {decisions[sid]['bads'] or []}  "
          f"({len(decisions[sid]['bads'])} of {n_eeg} = {len(decisions[sid]['bads']) / n_eeg:.0%} of the montage)")
print()
print("Rank before and after (30 channels; the array keeps 30 rows throughout):")
for sid in SUBJECTS:
    k = len(decisions[sid]["bads"])
    print(f"  {sid}: rank {l2.data_rank(n_eeg)['rank']} as loaded  ->  "
          f"{l2.data_rank(n_eeg, n_interpolated=k, average_reference=True)['rank']} after interpolating {k} "
          f"channel(s) and taking an average reference   "
          f"[{l2.data_rank(n_eeg, n_interpolated=k, average_reference=True)['arithmetic']}]")
print()
print(f"Effect of interpolation on the average reference ({SID_DEMO}, bad channel {BAD}, probe {probe}):")
print(f"  including the flagged channel in the average shifts {probe} by "
      f"{np.sqrt(((xa - xb) ** 2).mean()):.2f} uV RMS (max {np.abs(xa - xb).max():.1f} uV), against "
      f"{probe}'s own {np.sqrt((xb ** 2).mean()):.2f} uV RMS")
print(f"  interpolating first leaves a residual of {np.sqrt(((xc - xb) ** 2).mean()):.2f} uV RMS -- "
      "interpolation does not restore information, it restores a usable channel set")
print()
print("ex-2-2-flag-bads-drill is scored in the w-bad-channel-detective widget against its own reference flags "
      "(label_source: algorithmic until the author reviews them); this notebook produces no numeric key.")