# 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", "moabb==1.7.2", "pooch>=1.8"]
    subprocess.check_call(_cmd)

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

# 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

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; "
      "downloads go to MNE's default data directory unless EEG_COURSE_DATA is set")

from scipy import signal

DATASET = "ds-eegbci"
RUNS = [("S001", "R01"), ("S001", "R02"), ("S002", "R01"), ("S003", "R02"),
        ("S004", "R01"), ("S005", "R02"), ("S011", "R01")]
raws = {(s, r): helpers.load_spine(DATASET, s, r) for s, r in RUNS}
SF = raws[("S001", "R01")].info["sfreq"]
for (s, r), raw in raws.items():
    rep = helpers.first_look_report(raw)
    print(f"{s} {r} ({helpers.EEGBCI_RUNS[r].split(': ')[1].split(' (')[0]}): {rep['sfreq_hz']:.0f} Hz, {rep['n_channels']} ch, "
          f"{rep['duration_s']:.1f} s, noisiest channel {rep['amplitude_uV']['noisiest_channel']} "
          f"(max |x| {rep['amplitude_uV']['max_abs']:.0f} uV)")

def band_power_db(raw, band, tmin=None, tmax=None):
    """Per-channel mean Welch PSD over `band` in the window (2-s segments), dB re 1 uV^2/Hz."""
    n_fft = int(2 * raw.info["sfreq"])
    spec = raw.compute_psd(method="welch", picks="eeg", fmin=band[0], fmax=band[1], tmin=tmin, tmax=tmax,
                           n_fft=n_fft, n_overlap=n_fft // 2, verbose=False)
    psd, _ = spec.get_data(return_freqs=True)
    return 10 * np.log10(psd.mean(axis=1) * 1e12)


def atlas_panel(raw, channels, window, *, focus, reference, title, band=None, values=None, unit=None, map_title=None,
                highlight=None, spacing_uV=150, psd_fmin=0.5, psd_fmax=None, n_fft_s=2.0):
    """Trace | spectrum | topography for one artifact.

    focus     : channel whose spectrum is shown over `window`
    reference : ("channel", name) -> another channel over the same window, or
                ("window", (t0, t1)) -> the focus channel over a clean window
    band      : topography of band power over the window (dB re 1 uV^2/Hz), or
    values    : a precomputed per-channel measure (with `unit` and `map_title`) to map instead
    """
    t0, t1 = window
    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2), gridspec_kw=dict(width_ratios=[2.4, 1.4, 1.0]))
    helpers.plot_traces(raw, channels, t0=t0, duration=t1 - t0, spacing_uV=spacing_uV,
                        title=f"{title}: traces", highlight=highlight, ax=axes[0])
    fmax = psd_fmax or raw.info["sfreq"] / 2
    span = f"{t0:.1f}-{t1:.1f} s"
    helpers.plot_psd(raw, [focus], fmin=psd_fmin, fmax=fmax, tmin=t0, tmax=t1, n_fft_s=n_fft_s, ax=axes[1],
                     label_prefix=f"{span}: ", title="")
    kind, ref = reference
    if kind == "channel":
        helpers.plot_psd(raw, [ref], fmin=psd_fmin, fmax=fmax, tmin=t0, tmax=t1, n_fft_s=n_fft_s, ax=axes[1],
                         label_prefix=f"{span}: ", title=f"PSD of {focus} vs {ref}, {span} (dB re 1 uV^2/Hz)")
    else:
        helpers.plot_psd(raw, [focus], fmin=psd_fmin, fmax=fmax, tmin=ref[0], tmax=ref[1], n_fft_s=n_fft_s, ax=axes[1],
                         label_prefix=f"{ref[0]:.1f}-{ref[1]:.1f} s: ",
                         title=f"PSD of {focus}, artifact vs clean window (dB re 1 uV^2/Hz)")
    axes[1].legend(fontsize=8)
    if band is not None:
        helpers.plot_topomap_values(raw, band_power_db(raw, band, t0, t1), unit="dB re 1 uV^2/Hz",
                                    title=f"{band[0]:g}-{band[1]:g} Hz power, {span}\n", ax=axes[2])
    else:
        helpers.plot_topomap_values(raw, values, unit=unit, title=f"{map_title or 'per-channel measure'}\n", ax=axes[2])
    for ax in axes[1:]:
        ax.title.set_fontsize(9)
    fig.tight_layout()
    return fig


def bandpass(x, lo, hi, sf=SF, order=2):
    """Zero-phase Butterworth band-pass used only to locate artifacts (never applied to plotted data)."""
    return signal.filtfilt(*signal.butter(order, [lo, hi], btype="band", fs=sf), x)

raw = raws[("S001", "R01")]
fp = bandpass(raw.get_data(picks=["Fp1", "Fp2"]) * 1e6, 0.5, 15)
o1 = bandpass(raw.get_data(picks=["O1"])[0] * 1e6, 0.5, 15)
same_sign = np.minimum(fp[0], fp[1])
i = int(same_sign.argmax())
t_blink = i / SF
above = same_sign > 0.5 * same_sign[i]
lo = i
while lo > 0 and above[lo - 1]:
    lo -= 1
hi = i
while hi < len(above) - 1 and above[hi + 1]:
    hi += 1
print(f"blink candidate at {t_blink:.2f} s: Fp1 {fp[0][i]:+.0f} uV, Fp2 {fp[1][i]:+.0f} uV, O1 {o1[i]:+.0f} uV (0.5-15 Hz); "
      f"width at half maximum {(hi - lo) / SF * 1000:.0f} ms")
fig = atlas_panel(raw, ["Fp1", "Fp2", "F3", "F4", "Fz", "Cz", "O1"], (t_blink - 3, t_blink + 3),
                  focus="Fp1", reference=("channel", "O1"), band=(0.5, 4), title="Blink -- S001 R01",
                  highlight=(t_blink - 0.3, t_blink + 0.3), spacing_uV=200, psd_fmax=40)

raw = raws[("S002", "R01")]
f78 = bandpass(raw.get_data(picks=["F7", "F8"]) * 1e6, 0.5, 15)
heog = f78[0] - f78[1]
k = int(0.3 * SF)
step = np.array([heog[j:j + k].mean() - heog[j - k:j].mean() for j in range(k, len(heog) - k)])
j = int(np.abs(step).argmax()) + k
t_sac = j / SF
print(f"largest opposite-polarity step between F7 and F8: {step[j - k]:+.0f} uV at {t_sac:.2f} s -- "
      f"F7 {f78[0][j - k:j].mean():+.0f} -> {f78[0][j:j + k].mean():+.0f} uV, "
      f"F8 {f78[1][j - k:j].mean():+.0f} -> {f78[1][j:j + k].mean():+.0f} uV (means over 0.3 s before/after)")
fig = atlas_panel(raw, ["F7", "F8", "Fp1", "Fp2", "C3", "O1"], (50.5, 54.5),
                  focus="F7", reference=("channel", "Cz"), band=(0.5, 4), title="Horizontal saccades -- S002 R01",
                  highlight=(t_sac - 0.15, t_sac + 0.45), spacing_uV=200, psd_fmax=40)

seg = raw.copy().crop(50.5, 54.5)
t = seg.times + 50.5
x = seg.get_data(picks=["F7", "F8"]) * 1e6
fig, ax = plt.subplots(figsize=(12, 3))
ax.plot(t, x[0], lw=0.8, label="F7")
ax.plot(t, x[1], lw=0.8, label="F8")
ax.plot(t, x[0] - x[1] - 300, lw=0.8, label="F7 - F8 (offset -300 uV)")
ax.set(title="S002 R01: F7 and F8 are mirror images during the saccade train (uV)", xlabel="Time (s)", ylabel="Amplitude (uV)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)
plt.show()   # render the static figure(s) of this cell inline

raw = raws[("S002", "R01")]
hf = signal.filtfilt(*signal.butter(4, [25, 72], btype="band", fs=SF), raw.get_data() * 1e6)
n = int(2 * SF)
n_win = hf.shape[1] // n
rms = np.sqrt((hf[:, :n_win * n].reshape(len(raw.ch_names), n_win, n) ** 2).mean(-1))     # channel x 2-s window
ratio = rms / np.median(rms, axis=0, keepdims=True)
ci, wi = np.unravel_index(int(ratio.argmax()), ratio.shape)
print(f"largest 25-72 Hz RMS relative to the median channel: {raw.ch_names[ci]} in the 2-s window starting at {wi * 2} s "
      f"({rms[ci, wi]:.0f} uV vs median {np.median(rms[:, wi]):.0f} uV, x{ratio[ci, wi]:.1f})")
fig = atlas_panel(raw, ["F7", "T7", "C3", "T8", "F8", "O1"], (44, 52),
                  focus="T7", reference=("channel", "C3"), band=(25, 72), title="EMG -- S002 R01, T7",
                  highlight=(44, 52), spacing_uV=150, psd_fmax=80)

raw = raws[("S003", "R02")]
T0, T1 = 38, 58
slow = signal.filtfilt(*signal.butter(2, 0.5, btype="low", fs=SF), raw.get_data() * 1e6)
pp_slow = np.ptp(slow[:, int(T0 * SF):int(T1 * SF)], axis=1)
ci = int(pp_slow.argmax())
print(f"largest peak-to-peak of the < 0.5 Hz component in {T0}-{T1} s: {pp_slow[ci]:.0f} uV on {raw.ch_names[ci]}; "
      f"median across channels {np.median(pp_slow):.0f} uV")
fig = atlas_panel(raw, ["Fp1", "Fp2", "F7", raw.ch_names[ci], "Cz", "O1"], (T0, T1),
                  focus="Fp1", reference=("channel", "O1"), values=pp_slow, unit="uV", map_title=f"peak-to-peak of the < 0.5 Hz component, {T0}-{T1} s",
                  title="Drift / slow potentials -- S003 R02", spacing_uV=300, psd_fmin=0.125, psd_fmax=20, n_fft_s=8.0)

raw = raws[("S011", "R01")]
d = raw.get_data() * 1e6
jump = np.abs(np.diff(d, axis=1))
ci, si = np.unravel_index(int(jump.argmax()), jump.shape)
t_pop = (si + 1) / SF
others = float(np.delete(jump[:, si], ci).max())
w = int(0.1 * SF)
step = d[ci, si + 1:si + 1 + w].mean() - d[ci, si + 1 - w:si + 1].mean()
print(f"largest sample-to-sample jump: {jump[ci, si]:.0f} uV on {raw.ch_names[ci]} at {t_pop:.2f} s "
      f"(largest jump on any other channel at that sample: {others:.0f} uV); step of the 0.1-s means: {step:+.0f} uV")
pp_pop = np.ptp(d[:, int((t_pop - 0.5) * SF):int((t_pop + 1.5) * SF)], axis=1)
fig = atlas_panel(raw, ["O1", "O2", "Oz", "PO4", "P8", "Cz"], (t_pop - 2, t_pop + 4),
                  focus="O2", reference=("channel", "O1"), values=pp_pop, unit="uV", map_title="peak-to-peak, 2 s around the pop",
                  title="Electrode pop (candidate) -- S011 R01", highlight=(t_pop - 0.05, t_pop + 0.9),
                  spacing_uV=150, psd_fmax=80)

raw = raws[("S005", "R02")]
spec = raw.compute_psd(method="welch", picks="eeg", fmin=1, fmax=80, n_fft=int(4 * SF), verbose=False)
psd, freqs = spec.get_data(return_freqs=True)
i60 = int(np.argmin(np.abs(freqs - 60)))
neigh = (np.abs(freqs - 60) > 2) & (np.abs(freqs - 60) < 8)
ratio60 = psd[:, i60] / np.median(psd[:, neigh], axis=1)
print(f"PSD at 60 Hz relative to the median of 52-58 and 62-68 Hz: median across channels x{np.median(ratio60):.0f}, "
      f"max x{ratio60.max():.0f} on {spec.ch_names[int(ratio60.argmax())]}")
fig = atlas_panel(raw, ["Fp1", "Fz", "Cz", "CPz", "T10", "O1"], (20, 30),
                  focus="Fz", reference=("channel", "O1"), band=(59, 61), title="Line noise, 60 Hz -- S005 R02",
                  spacing_uV=150, psd_fmax=80)
ax = helpers.plot_traces(raw, ["Fz", "Cz"], t0=20, duration=0.5, spacing_uV=100,
                         title="Line noise -- S005 R02: a 0.5-s window resolves the individual 60 Hz cycles")

raw = raws[("S004", "R01")]
names = raw.ch_names
bp = bandpass(raw.get_data() * 1e6, 0.5, 40)
pos = np.array([ch["loc"][:3] for ch in raw.info["chs"]])
pairs = [(i, j) for i in range(len(names)) for j in range(i + 1, len(names)) if np.linalg.norm(pos[i] - pos[j]) < 0.04]
ed = np.array([np.var(bp[i] - bp[j]) for i, j in pairs])          # electrical distance, uV^2
order = np.argsort(ed)
print(f"{len(pairs)} neighbouring pairs (< 4 cm apart); median electrical distance {np.median(ed):.0f} uV^2")
for k in order[:5]:
    i, j = pairs[k]
    print(f"  {names[i]}-{names[j]}: {ed[k]:6.1f} uV^2 ({ed[k] / np.median(ed):.3f} x median), correlation {np.corrcoef(bp[i], bp[j])[0, 1]:.4f}")
i, j = pairs[order[0]]
typical = next(pairs[k] for k in order[::-1] if i in pairs[k] and abs(ed[k] - np.median(ed)) < np.median(ed))
m = typical[1] if typical[0] == i else typical[0]

fig, axes = plt.subplots(1, 3, figsize=(15, 4.2), gridspec_kw=dict(width_ratios=[2.4, 1.4, 1.0]))
t0, t1 = 10, 15
sl = slice(int(t0 * SF), int(t1 * SF))
t = np.arange(sl.start, sl.stop) / SF
axes[0].plot(t, bp[i, sl], lw=0.8, label=names[i])
axes[0].plot(t, bp[j, sl], lw=0.8, alpha=0.7, label=names[j])
axes[0].plot(t, bp[i, sl] - bp[j, sl] - 100, lw=0.8, label=f"{names[i]} - {names[j]} (offset -100 uV)")
axes[0].plot(t, bp[i, sl] - bp[m, sl] - 200, lw=0.8, label=f"{names[i]} - {names[m]} (offset -200 uV), a typical neighbour")
axes[0].set(title=f"Bridging candidate -- S004 R01: {names[i]} and {names[j]} overlap (uV, 0.5-40 Hz)",
            xlabel="Time (s)", ylabel="Amplitude (uV)")
axes[0].legend(fontsize=8)
axes[0].grid(alpha=0.3)
top = order[:15]
axes[1].bar(range(len(top)), ed[top], color="tab:blue")
axes[1].axhline(np.median(ed), color="gray", ls="--", label="median of all pairs")
axes[1].set_xticks(range(len(top)))
axes[1].set_xticklabels([f"{names[a]}-{names[b]}" for a, b in (pairs[k] for k in top)], rotation=90, fontsize=7)
axes[1].set(title="15 smallest electrical distances (uV^2)", ylabel="var(x_i - x_j) (uV^2)", yscale="log")
axes[1].legend(fontsize=8)
min_ed = np.full(len(names), np.nan)
for k, (a, b) in enumerate(pairs):
    min_ed[a] = np.nanmin([min_ed[a], ed[k]])
    min_ed[b] = np.nanmin([min_ed[b], ed[k]])
min_ed = np.where(np.isnan(min_ed), np.nanmedian(min_ed), min_ed)
helpers.plot_topomap_values(raw, np.log10(min_ed), unit="log10 uV^2", title="smallest electrical distance to a neighbour\n", ax=axes[2])
for ax in axes:
    ax.title.set_fontsize(9)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

raw = raws[("S001", "R02")]
fig = atlas_panel(raw, ["Fp1", "Cz", "P3", "Pz", "P4", "O1", "Oz", "O2"], (30, 40),
                  focus="O1", reference=("channel", "Fp1"), band=(8, 12), title="Posterior alpha (not an artifact) -- S001 R02",
                  spacing_uV=150, psd_fmax=40)

demonstrated = {
    "blink": "S001 R01, Fp1/Fp2, same-sign 0.5-15 Hz deflection (section 1)",
    "horizontal saccade": "S002 R01, F7/F8 opposite polarity, HEOG-like step (section 2)",
    "EMG": "S002 R01, T7, sustained 25-72 Hz activity (section 3)",
    "drift / slow potential": "S003 R02, < 0.5 Hz excursion; source unattributed (section 4)",
    "electrode pop (candidate)": "S011 R01, O2, single-channel step with slow recovery (section 5)",
    "line noise 60 Hz": "S005 R02, all channels; only the fundamental is below Nyquist (section 6)",
    "electrode bridging (candidate)": "S004 R01, near-zero electrical distance between two neighbours (section 7)",
    "posterior alpha (contrast, not an artifact)": "S001 R02, occipital 8-12 Hz peak (section 8)",
}
missing = {
    "ECG": "no ECG channel; no QRS-like train identified -- TODO(confirm)",
    "pulse": "no single-channel pulse wave identified -- TODO(confirm)",
    "cable / motion": "not attributable without a movement log -- TODO(confirm)",
    "sweat drift (attributed)": "drift shown, source unattributed -- TODO(confirm)",
    "flat / dead channel": "none found in the runs used",
    "drowsiness-related slowing": "one-minute runs, no sleep scoring -- not labelled",
}
print("nb-0-5-artifact-atlas -- ds-eegbci (EEGMMIDB; PhysioNet DOI 10.13026/C28G6P; ODC-By 1.0)")
print(f"Demonstrated ({len(demonstrated)}; every label is label_source: algorithmic, TODO(confirm) at expert review):")
for k, v in demonstrated.items():
    print(f"  - {k}: {v}")
print(f"Not shown from these baselines ({len(missing)}):")
for k, v in missing.items():
    print(f"  - {k}: {v}")
print("Item bank: w-spot-the-artifact is algorithmically labelled pending expert review (label_source: algorithmic).")