# 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", "edfio>=0.4"]
    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 mne.datasets import eegbci

DATASET, SUBJECT, RUN = "ds-eegbci", "S001", "R01"
root = helpers.data_dir()                                              # None = MNE's default data directory
edf_path = eegbci.load_data(1, [1], update_path=False, verbose=False, **({"path": str(root)} if root else {}))[0]
raw = mne.io.read_raw_edf(edf_path, preload=True, verbose=False)      # the file as stored: names untouched
print("file:", Path(edf_path).name, "|", raw)
names_in_file = list(raw.ch_names)

montage = mne.channels.make_standard_montage("standard_1005")           # a superset of the 10-10 positions
try:
    raw.copy().set_montage(montage, on_missing="raise", verbose=False)
except ValueError as e:
    print("set_montage with the raw labels ->", str(e).split("\n")[0][:110], "...")

eegbci.standardize(raw)                                                # dataset-specific: strip the dots, fix the case
mapping = dict(zip(names_in_file, raw.ch_names))
changed = {k: v for k, v in mapping.items() if k != v}
print(f"{len(changed)} of {len(mapping)} labels changed; examples:")
for k in ["Fc5.", "Fcz.", "Fp1.", "Afz.", "T9..", "Iz..", "Po7.", "Cp3."]:
    print(f"   {k!r:8s} -> {mapping[k]!r}")
raw.set_montage(montage, on_missing="raise", verbose=False)           # no ValueError now: every name is known
n_pos = sum(bool(np.any(ch["loc"][:3] != 0)) for ch in raw.info["chs"])
print(f"montage attached: {n_pos} of {len(raw.ch_names)} channels have a position")

fig = plt.figure(figsize=(12, 5.5))
ax2d = fig.add_subplot(1, 2, 1)
ax3d = fig.add_subplot(1, 2, 2, projection="3d")
raw.plot_sensors(show_names=True, axes=ax2d, show=False)
raw.plot_sensors(kind="3d", axes=ax3d, show=False)
ax2d.set_title("ds-eegbci S001: 64 electrodes, template 10-10 positions (standard_1005), top view")
ax3d.set_title("the same positions in 3-D (head coordinates, m)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

# The odd/even/z convention, checked against the x coordinate (+x = right ear).
agree, disagree = 0, []
for ch in raw.info["chs"]:
    name, x_mm = ch["ch_name"], ch["loc"][0] * 1000
    digit = next((c for c in reversed(name) if c.isdigit()), None)
    if name.endswith("z"):
        ok = abs(x_mm) < 5
    elif digit is not None:
        ok = (x_mm < 0) if int(digit) % 2 else (x_mm > 0)
    else:
        ok = True
    agree += ok
    if not ok:
        disagree.append((name, round(x_mm)))
print(f"odd = left, even = right, z = midline: {agree} of {len(raw.ch_names)} labels agree with their template position"
      + (f"; exceptions: {disagree}" if disagree else ""))

LEGACY_1020 = {"T3": "T7", "T4": "T8", "T5": "P7", "T6": "P8"}      # 10-20 legacy names -> 10-10 names

# --- a vendor-style copy of the recording (synthetic names; the data are S001 R01 unchanged)
messy = raw.copy()
vendor_names = {}
for name in messy.ch_names:
    legacy = {v: k for k, v in LEGACY_1020.items()}.get(name, name)  # T7 -> T3 etc.
    vendor_names[name] = f"EEG {legacy.upper()}-REF"
messy.rename_channels(vendor_names)
status = mne.io.RawArray(np.zeros((1, messy.n_times)), mne.create_info(["Status"], messy.info["sfreq"], "eeg"), verbose=False)
messy.add_channels([status], force_update_info=True)
print("vendor-style labels:", messy.ch_names[:4], "...", messy.ch_names[-3:])
try:
    messy.copy().set_montage(montage, on_missing="raise", verbose=False)
except ValueError as e:
    print("set_montage ->", str(e).split("\n")[0][:100], "...")


def normalise_name(label, known):
    """Map an exported label onto a montage name: strip prefix/suffix, ignore case, translate legacy 10-20 names."""
    core = label.strip()
    for prefix in ("EEG ", "EEG:", "EEG_"):
        if core.upper().startswith(prefix):
            core = core[len(prefix):]
    core = core.split("-")[0].strip()                    # 'FP1-REF' -> 'FP1'
    core = LEGACY_1020.get(core.upper(), core)
    by_upper = {k.upper(): k for k in known}             # case-insensitive lookup into the montage names
    return by_upper.get(core.upper(), label)             # unknown labels are left alone (and reported below)


known = set(montage.ch_names)
fixed = {ch: normalise_name(ch, known) for ch in messy.ch_names}
unresolved = [ch for ch, new in fixed.items() if new not in known]
print(f"renamed {sum(k != v for k, v in fixed.items())} labels; unresolved: {unresolved}")
messy.rename_channels(fixed)
messy.set_channel_types({"Status": "stim"})               # not EEG: excluded from EEG picks, references, topographies
messy.set_montage(montage, on_missing="raise", verbose=False)   # the stim channel needs no position
n_eeg, n_stim = len(mne.pick_types(messy.info, eeg=True)), len(mne.pick_types(messy.info, stim=True))
print(f"after the repair: {n_eeg} EEG channels with positions + {n_stim} stim channel; "
      f"identical data to the original: {np.array_equal(messy.get_data(picks='eeg'), raw.get_data())}")

import json

setf = helpers.fetch_iowapd("sub-001")                     # prints what it fetches or reuses; None if unreachable
IOWA_OK = setf is not None
if not IOWA_OK:
    print("ds-iowapd could not be fetched; the rest of this section is skipped")
else:
    sidecar = json.loads((setf.parent / setf.name.replace("_eeg.set", "_eeg.json")).read_text())
    channels_tsv = (setf.parent / setf.name.replace("_eeg.set", "_channels.tsv")).read_text().splitlines()
    print("BIDS sidecar (selected keys):")
    for k in ("EEGReference", "EEGGround", "EEGChannelCount", "SamplingFrequency", "PowerLineFrequency",
              "RecordingDuration", "SoftwareFilters"):
        print(f"   {k:20s} {sidecar.get(k)}")
    print(f"channels.tsv: {len(channels_tsv) - 1} rows (header: {channels_tsv[0]!r})")

    iowa = helpers.load_iowapd("sub-001", quiet=True)
    rep = helpers.first_look_report(iowa, "ds-iowapd")
    ann_counts = {str(k): v for k, v in rep["annotation_counts"].items()}
    print(f"\nfile: {setf.name} | {rep['sfreq_hz']:g} Hz | {rep['n_channels']} channels ({rep['channel_types']}) | "
          f"{rep['duration_s']:.1f} s | annotations {ann_counts}")
    print("header filters:", rep["highpass_hz_in_header"], "Hz high-pass,", rep["lowpass_hz_in_header"], "Hz low-pass "
          "(the header records no filter; the catalog documents a 0.1 Hz online high-pass -- the file does not carry it)")
    print("is the documented reference 'Pz' among the stored channels?", "Pz" in iowa.ch_names)

if IOWA_OK:
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")                                 # MNE reminds us to set the montage again; we do
        with_ref = mne.add_reference_channels(iowa, "Pz", copy=True)   # the reference as an explicit, all-zero channel
    with_ref.set_montage("standard_1005", on_missing="ignore", verbose=False)
    with_ref.info["bads"] = ["Pz"]                                  # for the plot: bads are drawn in red
    fig = plt.figure(figsize=(12, 5.5))
    ax_a = fig.add_subplot(1, 2, 1)
    ax_b = fig.add_subplot(1, 2, 2)
    iowa.plot_sensors(show_names=True, axes=ax_a, show=False)
    with_ref.plot_sensors(show_names=True, axes=ax_b, show=False)
    ax_a.set_title("ds-iowapd sub-001: the 63 stored channels (template 10-10 positions)")
    ax_b.set_title("the same, with the documented reference Pz added back (red)")
    fig.tight_layout()
    plt.show()   # render the static figure(s) of this cell inline
    with_ref.info["bads"] = []

if IOWA_OK:
    eeg_names = [iowa.ch_names[i] for i in mne.pick_types(iowa.info, eeg=True)]
    sd = iowa.get_data(picks=eeg_names).std(axis=1) * 1e6
    order = np.argsort(sd)
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.6), gridspec_kw=dict(width_ratios=[1, 1.8]))
    helpers.plot_topomap_values(iowa, sd, unit="uV", title="Standard deviation per channel", ax=axes[0], vlim=(0, float(np.percentile(sd, 95))))
    axes[1].bar(range(len(sd)), sd[order], color=["tab:red" if i < 2 else "tab:blue" for i in range(len(sd))])
    axes[1].set_xticks(range(len(sd)))
    axes[1].set_xticklabels([eeg_names[i] for i in order], rotation=90, fontsize=7)
    axes[1].set(ylabel="std (uV)", title="Channel standard deviation, sorted (uV); the two smallest in red")
    axes[1].grid(alpha=0.3, axis="y")
    fig.tight_layout()
    plt.show()   # render the static figure(s) of this cell inline

    lowest = [(eeg_names[i], float(sd[i])) for i in order[:2]]
    pos = {ch["ch_name"]: ch["loc"][:3] for ch in with_ref.info["chs"]}
    for name, s in lowest:
        d_cm = np.linalg.norm(pos[name] - pos["Pz"]) * 100
        print(f"{name}: std {s:.1f} uV (median over channels {np.median(sd):.1f} uV), {d_cm:.1f} cm from the Pz position")

if IOWA_OK:
    print("flat channels in the file as stored     :", helpers.flat_channels(iowa) or "none")
    print("flat channels with Pz added back        :", helpers.flat_channels(with_ref))
    avg = with_ref.copy().set_eeg_reference("average", projection=False, verbose=False)
    print(f"Pz after an average reference           : std {float(avg.get_data(picks='Pz').std() * 1e6):.1f} uV (no longer flat)")
    eeg_before = len(mne.pick_types(iowa.info, eeg=True))
    print(f"EEG channels: {eeg_before} stored -> {len(mne.pick_types(with_ref.info, eeg=True))} with the reference added back")

print(f"nb-0-2-montages -- {DATASET} (EEGMMIDB; PhysioNet DOI 10.13026/C28G6P; ODC-By 1.0), subject {SUBJECT} run {RUN}")
print(f"labels mapped from the EDF header to 10-10 names: {len(changed)} of {len(mapping)} changed; "
      f"{n_pos} of {len(raw.ch_names)} channels positioned by standard_1005; "
      f"odd/even/z convention agrees with the template positions for {agree} of {len(raw.ch_names)} labels")
print(f"vendor-style repair: {sum(k != v for k, v in fixed.items())} labels renamed, 1 channel retyped (Status: eeg -> stim), "
      f"{n_eeg} EEG channels positioned")
print()
if IOWA_OK:
    ref_name = sidecar.get("EEGReference")
    print("ds-iowapd (OpenNeuro ds004584; CC0; DOI 10.18112/openneuro.ds004584.v1.0.0), sub-001 task-Rest:")
    print(f"  {rep['sfreq_hz']:g} Hz | {rep['n_channels']} stored channels (catalog: 64-channel cap) | {rep['duration_s']:.1f} s")
    flat_after = helpers.flat_channels(with_ref)
    print(f"  recording reference: {ref_name}  <-- the answer this notebook reports (L0.2 exercise, 'identify the recording reference')")
    print(f"     named by the documentation (sidecar EEGReference = {ref_name!r}; catalog: Pz online reference) and confirmed by the data:")
    print(f"     (1) {rep['n_channels']} stored channels for a 64-electrode cap, and the empty position in the montage is the parietal midline;")
    print(f"     (2) no stored channel is flat (smallest std {sd[order[0]]:.1f} uV on {eeg_names[order[0]]}), so the reference electrode was not "
          f"stored -- added back it is the flat channel ({flat_after[0][0]}, {flat_after[0][1]:.2f} uV);")
    print(f"     (3) the two lowest-amplitude channels, {lowest[0][0]} ({lowest[0][1]:.1f} uV) and {lowest[1][0]} ({lowest[1][1]:.1f} uV), "
          f"lie beside that position (median channel std {np.median(sd):.1f} uV)")
else:
    print("ds-iowapd: not fetched in this run (see the note above); the reference answer is TODO(confirm) until the file is checked")
print()
print("Note: in this file the reference electrode is absent rather than stored flat; the flat channel appears only once the")
print("reference is added back. TODO(confirm): draft answer until the author reviews it.")