# 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")

DATASET, SUBJECT = "ds-eegbci", "S001"
raw_eo = helpers.load_spine(DATASET, SUBJECT, "R01")   # R01: baseline, eyes open
raw_ec = helpers.load_spine(DATASET, SUBJECT, "R02")   # R02: baseline, eyes closed
print("R01:", helpers.EEGBCI_RUNS["R01"], "->", raw_eo)
print("R02:", helpers.EEGBCI_RUNS["R02"], "->", raw_ec)
print("Refused by the loader (catalog caveat):", ", ".join(helpers.EEGBCI_EXCLUDE_DEFAULT),
      "and, by default,", ", ".join(helpers.EEGBCI_EXCLUDE_OPTIONAL))

from IPython.display import HTML, display

KEYS = ["description", "sfreq_hz", "nyquist_hz", "n_channels", "channel_types", "duration_s",
        "highpass_hz_in_header", "lowpass_hz_in_header", "montage_attached", "annotation_counts", "amplitude_uV"]
for raw, label in ((raw_eo, "S001 R01, eyes open"), (raw_ec, "S001 R02, eyes closed")):
    rep = helpers.first_look_report(raw, DATASET)
    display(HTML(helpers.report_html({k: rep[k] for k in KEYS}, f"First look: {label}")))

from mne.datasets import eegbci

edf_path = eegbci.load_data(1, [1], update_path=False, verbose=False)[0]   # cached by load_spine
raw_edf = mne.io.read_raw_edf(edf_path, preload=False, verbose=False)
print("file:", Path(edf_path).name)
print("names in the EDF header :", raw_edf.ch_names[:6], "...")
eegbci.standardize(raw_edf)
print("standardised 10-10 names:", raw_edf.ch_names[:6], "...")
montage = mne.channels.make_standard_montage("standard_1005")
raw_edf.set_montage(montage, on_missing="raise", verbose=False)   # would raise if a name were unknown; none is
n_pos = sum(bool(np.any(ch["loc"][:3] != 0)) for ch in raw_eo.info["chs"])
print(f"channels with a position in raw_eo: {n_pos} of {len(raw_eo.ch_names)}")

fig = raw_eo.plot_sensors(show_names=True, show=False)
fig.set_size_inches(6, 6)
fig.suptitle("ds-eegbci S001: 64 electrodes at template 10-10 positions (standard_1005; a layout, no data units)", fontsize=9)
plt.show()   # render the static figure(s) of this cell inline

ANATOMICAL = ["Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "T7", "C3", "Cz", "C4", "T8",
              "P7", "P3", "Pz", "P4", "P8", "O1", "Oz", "O2"]


def browse(raw, start, duration, sensitivity_uV, title, channels=ANATOMICAL):
    """One page of MNE's raw browser as a static figure (sensitivity in uV per channel slot)."""
    order = [raw.ch_names.index(ch) for ch in channels]
    fig = raw.plot(start=start, duration=duration, n_channels=len(order), order=order,
                   scalings=dict(eeg=sensitivity_uV * 1e-6), show=False, block=False,
                   show_scrollbars=False, show_scalebars=True, verbose=False)
    fig.set_size_inches(11, 0.28 * len(order) + 1.4)
    fig.set_dpi(64)   # dense pages: keep the notebook file small
    pos = fig.mne.ax_main.get_position()
    fig.mne.ax_main.set_position([pos.x0, pos.y0, pos.width, pos.height - 0.06])
    fig.suptitle(f"{title} -- {duration:g}-s window, {sensitivity_uV:g} uV per channel slot", fontsize=10, y=0.99)
    return fig


fig = browse(raw_ec, 20, 10, 100, "S001 R02 (eyes closed), anatomical order")
plt.show()   # render the static figure(s) of this cell inline

fig = browse(raw_ec, 20, 2, 100, "S001 R02 (eyes closed): a 2-s window shows the alpha waveform")
plt.show()   # render the static figure(s) of this cell inline

fig = browse(raw_eo, 20, 10, 100, "S001 R01 (eyes open), the same settings")
plt.show()   # render the static figure(s) of this cell inline

fig = browse(raw_ec, 20, 10, 500, "S001 R02 (eyes closed): sensitivity too coarse -- the alpha trains flatten out")
plt.show()   # render the static figure(s) of this cell inline

t0, dur = 5, 10   # a window of the eyes-open run that contains a blink (section 5 locates it)
picks = ["Fp1", "F7", "T7", "P7", "O1"]
ref_avg = raw_eo.copy().set_eeg_reference("average", projection=False, verbose=False)
bipolar = mne.set_bipolar_reference(raw_eo, anode=["Fp1", "F7", "T7", "P7"], cathode=["F7", "T7", "P7", "O1"],
                                    drop_refs=True, verbose=False)
fig, axes = plt.subplots(3, 1, figsize=(11, 8.5))
helpers.plot_traces(raw_eo, picks, t0=t0, duration=dur, spacing_uV=150, ax=axes[0],
                    title="S001 R01, recording reference (as stored; reference electrode TODO(confirm))")
helpers.plot_traces(ref_avg, picks, t0=t0, duration=dur, spacing_uV=150, ax=axes[1],
                    title="S001 R01, average reference (each channel minus the mean of all 64)")
helpers.plot_traces(bipolar, ["Fp1-F7", "F7-T7", "T7-P7", "P7-O1"], t0=t0, duration=dur, spacing_uV=150, ax=axes[2],
                    title="S001 R01, bipolar left temporal chain (neighbour minus neighbour)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

from scipy import signal
import tempfile

sf = raw_eo.info["sfreq"]
fp = raw_eo.get_data(picks=["Fp1", "Fp2"]) * 1e6
b, a = signal.butter(2, [0.5, 15], btype="band", fs=sf)
same_sign = np.minimum(*signal.filtfilt(b, a, fp))          # positive only where both sides deflect together
t_blink = float(same_sign.argmax() / sf)
print(f"largest frontal deflection present on both Fp1 and Fp2: +{same_sign.max():.0f} uV at {t_blink:.2f} s "
      "(algorithmic, not reviewed)")

raw_ann = raw_eo.copy()
new = mne.Annotations(onset=[0.0, t_blink - 0.25], duration=[10.0, 0.5],
                      description=["inspect", "BAD_blink"],
                      orig_time=raw_ann.annotations.orig_time)   # same time origin as the file's own 'T0'
raw_ann.set_annotations(raw_ann.annotations + new)           # keeps 'T0'
for onset, duration, desc in zip(raw_ann.annotations.onset, raw_ann.annotations.duration, raw_ann.annotations.description):
    print(f"  {desc:10s} onset {onset:6.2f} s   duration {duration:5.2f} s")
fig = browse(raw_ann, 0, 15, 200, "S001 R01 with 'inspect' (0-10 s, grey) and 'BAD_blink' (red) annotations")
plt.show()   # render the static figure(s) of this cell inline

# Save the annotations (to a temporary folder here), read them back, and show
# that fixed-length epoching drops the span marked BAD_.
with tempfile.TemporaryDirectory() as tmp:
    f = Path(tmp) / "S001R01-annotations.csv"
    raw_ann.annotations.save(f, overwrite=True, verbose=False)
    print("read back:", mne.read_annotations(f))
epochs = mne.make_fixed_length_epochs(raw_ann, duration=2.0, reject_by_annotation=True, preload=True, verbose=False)
n_dropped = sum(1 for log in epochs.drop_log if log)
print(f"2-s fixed-length epochs: {len(epochs)} kept, {n_dropped} dropped because they overlap 'BAD_blink'")

OCC = ["O1", "Oz", "O2"]
fig, ax = plt.subplots(figsize=(10, 4.5))
helpers.plot_psd(raw_eo, OCC, fmin=1, fmax=40, ax=ax, label_prefix="eyes open ", title="")
for line in ax.get_lines():
    line.set_linestyle("--")
helpers.plot_psd(raw_ec, OCC, fmin=1, fmax=40, ax=ax, label_prefix="eyes closed ",
                 title="S001: occipital PSD, eyes open (dashed) vs eyes closed (solid); Welch, 4-s segments (dB re 1 uV^2/Hz)")
ax.axvspan(8, 12, color="tab:orange", alpha=0.15, label="8-12 Hz")
ax.legend(fontsize=8, ncol=2)
plt.show()   # render the static figure(s) of this cell inline

def band_power_db(raw, band):
    """Per-channel mean Welch PSD over `band` (2-s segments) in 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],
                           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)


v_eo, v_ec = band_power_db(raw_eo, (8, 12)), band_power_db(raw_ec, (8, 12))
vlim = (float(min(v_eo.min(), v_ec.min())), float(max(v_eo.max(), v_ec.max())))   # one colour scale for both
fig, axes = plt.subplots(1, 2, figsize=(9, 4))
helpers.plot_topomap_values(raw_eo, v_eo, unit="dB re 1 uV^2/Hz", title="Eyes open: 8-12 Hz power", ax=axes[0], vlim=vlim)
helpers.plot_topomap_values(raw_ec, v_ec, unit="dB re 1 uV^2/Hz", title="Eyes closed: 8-12 Hz power", ax=axes[1], vlim=vlim)
fig.tight_layout()

fig, axes = plt.subplots(1, 2, figsize=(12, 3), sharey=True)
for ax, raw, label in zip(axes, (raw_eo, raw_ec), ("eyes open (R01)", "eyes closed (R02)")):
    seg = raw.copy().crop(30, 33).pick("O1")
    ax.plot(seg.times + 30, seg.get_data()[0] * 1e6, "k", lw=0.8)
    ax.set(title=f"O1, {label}: 3 s (uV)", xlabel="Time (s)", ylabel="Amplitude (uV)")
    ax.grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

def band_power(raw, ch, band):
    """Mean Welch PSD (4-s segments, 50 % overlap) of channel `ch` over `band`, in uV^2/Hz."""
    n_fft = int(4 * raw.info["sfreq"])
    spec = raw.compute_psd(method="welch", picks=[ch], fmin=0.5, fmax=40.0,
                           n_fft=n_fft, n_overlap=n_fft // 2, verbose=False)
    psd, freqs = spec.get_data(return_freqs=True)
    m = (freqs >= band[0]) & (freqs <= band[1])
    return float(psd[0, m].mean() * 1e12)


ALPHA = (8.0, 12.0)
print(f"nb-0-4-browse -- {DATASET} (EEGMMIDB; PhysioNet DOI 10.13026/C28G6P; ODC-By 1.0), subject {SUBJECT}")
for raw, label in ((raw_eo, "R01 eyes open  "), (raw_ec, "R02 eyes closed")):
    n_eeg = len(mne.pick_types(raw.info, eeg=True))
    print(f"{label}: sampling rate {raw.info['sfreq']:.1f} Hz | {n_eeg} EEG channels | "
          f"duration {raw.times[-1] + 1 / raw.info['sfreq']:.2f} s")
for ch in ("O1", "O2", "Oz"):
    p_eo, p_ec = band_power(raw_eo, ch, ALPHA), band_power(raw_ec, ch, ALPHA)
    flag = "   <-- the number this notebook reports" if ch == "O1" else ""
    print(f"alpha {ALPHA[0]:g}-{ALPHA[1]:g} Hz at {ch}: eyes open {p_eo:7.2f} uV^2/Hz | eyes closed {p_ec:7.2f} uV^2/Hz | "
          f"ratio EC/EO = {p_ec / p_eo:.2f} ({10 * np.log10(p_ec / p_eo):.1f} dB){flag}")
print()
print("Exercise note: the L0.4 exercise (eyes-closed onset in a ds-lemon segment) is not computable from these two")
print("separate runs; the ratio above is the checkable output of this notebook. TODO(confirm): draft values until")
print("the author reviews them.")