# 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 = "ds-eegbci", "S001"
root = helpers.data_dir()                                              # None = MNE's default data directory
paths = eegbci.load_data(1, [1, 3], update_path=False, verbose=False, **({"path": str(root)} if root else {}))
raw_r01 = mne.io.read_raw_edf(paths[0], preload=True, verbose=False)
raw_r03 = mne.io.read_raw_edf(paths[1], preload=True, verbose=False)
for r in (raw_r01, raw_r03):
    eegbci.standardize(r)                                              # Sharbrough labels -> 10-10 names (nb-0-2)
    r.set_montage("standard_1005", on_missing="ignore", verbose=False)
print("R01:", helpers.EEGBCI_RUNS["R01"], "|", Path(paths[0]).name)
print(raw_r01.info)

try:
    import edfio
except ImportError:
    edfio = None
    print("edfio is not installed (pip install edfio): skipping the EDF header fields")
if edfio is not None:
    edf = edfio.read_edf(paths[0])
    print(f"EDF header of {Path(paths[0]).name}: version {getattr(edf, 'version', '?')}, start {edf.startdate} {edf.starttime}, "
          f"{edf.num_data_records} data records of {edf.data_record_duration} s, {edf.num_signals} signals")
    print(f"{'label':8s} {'unit':5s} {'phys min':>9s} {'phys max':>9s} {'dig min':>8s} {'dig max':>7s} {'uV/step':>8s} {'Hz':>5s}  prefiltering")
    for s in list(edf.signals)[:4]:
        step = (s.physical_max - s.physical_min) / (s.digital_max - s.digital_min)
        print(f"{s.label:8s} {s.physical_dimension:5s} {s.physical_min:9.1f} {s.physical_max:9.1f} {s.digital_min:8d} {s.digital_max:7d} "
              f"{step:8.4f} {s.sampling_frequency:5.0f}  {s.prefiltering!r}")
    print("... (one row per signal); an empty prefiltering field is why MNE reports high-pass 0 Hz and low-pass = Nyquist,")
    print("    consistent with the catalog's 'no hardware filters' for this dataset")

events_r03, event_id_r03 = mne.events_from_annotations(raw_r03, verbose=False)
print("R03:", helpers.EEGBCI_RUNS["R03"])
print("event_id:", {str(k): v for k, v in event_id_r03.items()}, "| events array shape:", events_r03.shape, "(sample, previous value, code)")
for label, r in (("R01", raw_r01), ("R03", raw_r03)):
    tc = helpers.event_timing_check(r)
    line = (f"{label}: {tc['n_samples']} samples / {tc['sfreq_hz']:g} Hz = {tc['duration_s']:.2f} s; "
            f"{tc['n_annotations']} annotations {tc['counts']}")
    if "interval_s" in tc:
        line += (f"; onset intervals median {tc['interval_s']['median']:.2f} s (min {tc['interval_s']['min']:.2f}, "
                 f"max {tc['interval_s']['max']:.2f}); overlapping spans {tc.get('n_overlapping_spans', 0)}, "
                 f"gaps between spans {tc.get('n_gaps_between_spans', 0)}, long gaps {tc['n_long_gaps']}")
    print(line)

fig, ax = plt.subplots(figsize=(12, 2.2))
colours = {"T0": "0.75", "T1": "tab:blue", "T2": "tab:orange"}
for onset, dur, desc in zip(raw_r03.annotations.onset, raw_r03.annotations.duration, raw_r03.annotations.description):
    ax.barh(0, dur, left=onset, height=0.6, color=colours.get(desc, "k"), edgecolor="w", lw=0.5)
ax.set(yticks=[], xlabel="Time (s)", xlim=(0, raw_r03.times[-1]),
       title="S001 R03: T0 (rest, grey), T1 (blue) and T2 (orange) spans tile the run without gaps or overlaps")
plt.show()   # render the static figure(s) of this cell inline

LEMON_SUBJECT, LEMON_MAX_MINUTES = "sub-010002", 5
print(f"free disk on the download volume: {helpers.free_disk_mb(helpers.download_root()):.0f} MB")
vhdr = helpers.fetch_lemon_raw(LEMON_SUBJECT, max_minutes=LEMON_MAX_MINUTES)
LEMON_OK = vhdr is not None
if not LEMON_OK:
    print("ds-lemon: skipped in this run (see the reason above); sections 2 and the LEMON lines of section 5 are omitted")
else:
    text = vhdr.read_text(encoding="utf-8", errors="replace").splitlines()
    print(f"--- {vhdr.name}: [Common Infos], [Binary Infos] and the first channel entries")
    for line in text[:16] + [l for l in text if l.startswith(("Ch1=", "Ch2=", "Ch17=", "Ch25="))]:
        print("   ", line)
    vmrk = vhdr.with_suffix(".vmrk")
    marks = [l for l in vmrk.read_text(encoding="utf-8", errors="replace").splitlines() if l.startswith("Mk")]
    print(f"--- {vmrk.name}: {len(marks)} markers in the file (the whole 17-min recording); the first eight:")
    for line in marks[:8]:
        print("   ", line)
    info_path = helpers.download_root() / "MNE-lemon-data" / "EEG_Info"
    if info_path.exists() or helpers.download_url(helpers.LEMON_EEG_INFO, info_path, quiet=True):
        coding = [l.strip() for l in info_path.read_text(errors="replace").splitlines() if l.strip().startswith(("S200", "S210"))]
        print("--- marker coding from the dataset's EEG_Info file:", "; ".join(coding))

if LEMON_OK:
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        raw_lemon = mne.io.read_raw_brainvision(vhdr, preload=False, verbose=False)
    for c in caught:
        if "annotation" in str(c.message).lower():
            print("MNE says:", str(c.message)[:120])
    raw_lemon.set_channel_types({"VEOG": "eog"})                      # the header types every line as EEG; VEOG is not
    raw_lemon.set_montage("standard_1005", on_missing="ignore", verbose=False)
    print(raw_lemon.info)
    types = raw_lemon.get_channel_types()
    print(f"channel types: {{t: types.count(t) for t in sorted(set(types))}}; FCz (the documented online reference) stored as a channel? "
          f"{'FCz' in raw_lemon.ch_names}")
    events_lemon, event_id_lemon = mne.events_from_annotations(raw_lemon, verbose=False)
    tc = helpers.event_timing_check(raw_lemon)
    print(f"{tc['n_samples']} samples / {tc['sfreq_hz']:g} Hz = {tc['duration_s']:.1f} s in the fetched prefix; "
          f"{tc['n_annotations']} markers: {tc['counts']}")
    print(f"onset intervals: median {tc['interval_s']['median']:.3f} s, min {tc['interval_s']['min']:.4f} s, "
          f"max {tc['interval_s']['max']:.3f} s; intervals > 2 x median: {tc['n_long_gaps']}")
    print("measurement start (from the 'New Segment' marker):", raw_lemon.info["meas_date"])

if LEMON_OK:
    onsets, descs = raw_lemon.annotations.onset, raw_lemon.annotations.description
    fig, ax = plt.subplots(figsize=(12, 2.4))
    for code, y, colour, label in (("Stimulus/S210", 0, "tab:blue", "S210 eyes closed"), ("Stimulus/S200", 1, "tab:orange", "S200 eyes open")):
        t = onsets[descs == code]
        ax.plot(t, np.full_like(t, y), "|", ms=14, color=colour, label=f"{label} ({len(t)} markers)")
    other = onsets[~np.isin(descs, ["Stimulus/S210", "Stimulus/S200"])]
    ax.plot(other, np.full_like(other, 0.5), "x", color="k", label=f"other markers ({len(other)})")
    ax.set(yticks=[0, 0.5, 1], yticklabels=["EC", "other", "EO"], xlabel="Time (s)", xlim=(0, raw_lemon.times[-1]),
           title=f"ds-lemon {LEMON_SUBJECT}, first {LEMON_MAX_MINUTES} min: a marker every 2 s names the condition of each one-minute block")
    ax.legend(fontsize=8, loc="center right")
    plt.show()   # render the static figure(s) of this cell inline

try:
    helpers.load_spine("ds-erpcore", 1)
except NotImplementedError as e:
    print("NotImplementedError:", e)

SEED = 20260917          # fixed: the recording, and therefore the answer, is reproducible
OFFSET_MS = 33.0         # planted: each stimulus appears this long after its trigger (plus the jitter below)
JITTER_MS = 1.0          # per-trial uniform jitter in [-1, +1] ms (a refresh-timing wobble)
N_STIM = 36
SFREQ_SYN = 1000.0


def make_synthetic_recording(raw_eeg, *, seed, offset_ms, jitter_ms, n_stim, sfreq, stim_s=0.2, first_s=2.0, isi_s=(1.0, 1.6)):
    """Real EEG (resampled to `sfreq`) + a simulated photodiode (misc) and trigger (stim) channel with a planted offset."""
    rng = np.random.default_rng(seed)
    rec = raw_eeg.copy().resample(sfreq, verbose=False)
    sf, n = rec.info["sfreq"], rec.n_times
    t_trig = first_s + np.cumsum(np.r_[0.0, rng.uniform(*isi_s, size=n_stim - 1)])
    t_trig = np.round(t_trig * sf) / sf                                   # the recorder samples the TTL line on its grid
    offsets = offset_ms + rng.uniform(-jitter_ms, jitter_ms, size=n_stim)
    t_stim = np.round((t_trig + offsets / 1000.0) * sf) / sf
    assert t_stim[-1] + stim_s + 0.1 < rec.times[-1], "the stimulus schedule must fit inside the recording"
    trig = np.zeros(n)
    diode = rng.normal(0.0, 0.02, n)                                       # dark level with sensor noise (arbitrary units)
    k, tail = int(stim_s * sf), int(0.02 * sf)
    tau = 0.001 * sf                                                       # 1-ms rise time constant
    for a, b in zip(t_trig, t_stim):
        i, j = int(round(a * sf)), int(round(b * sf))
        trig[i:i + int(0.01 * sf)] = 1.0                                  # 10-ms TTL pulse
        diode[j:j + k] += 1.0 - np.exp(-(np.arange(k) + 1) / tau)         # light on: exponential rise to 1
        diode[j + k:j + k + tail] += np.exp(-(np.arange(tail) + 1) / tau) # light off
    aux = mne.io.RawArray(np.vstack([diode, trig]), mne.create_info(["PHOTODIODE", "TRIG"], sf, ["misc", "stim"]), verbose=False)
    rec.add_channels([aux], force_update_info=True)
    rec.info["description"] = (f"SYNTHETIC recording for the L0.3 exercise: ds-eegbci S001 R01 EEG resampled to {sf:g} Hz + simulated "
                               f"PHOTODIODE and TRIG channels; seed {seed}; planted stimulus-trigger offset {offset_ms:g} ms +/- {jitter_ms:g} ms jitter")
    return rec, t_trig, t_stim, offsets


syn, t_trig, t_stim, planted = make_synthetic_recording(raw_r01, seed=SEED, offset_ms=OFFSET_MS, jitter_ms=JITTER_MS,
                                                        n_stim=N_STIM, sfreq=SFREQ_SYN)
print(syn)
print("channel types:", {t: syn.get_channel_types().count(t) for t in sorted(set(syn.get_channel_types()))})
print(f"{N_STIM} stimuli between {t_trig[0]:.2f} s and {t_trig[-1]:.2f} s; planted offsets: mean {planted.mean():.2f} ms, "
      f"range {planted.min():.2f}-{planted.max():.2f} ms (the exact values are hidden from the measurement below)")

import tempfile

with tempfile.TemporaryDirectory() as tmp:
    fif = Path(tmp) / "synthetic-eegbci-S001R01-photodiode_raw.fif"
    syn.save(fif, overwrite=True, verbose=False)
    back = mne.io.read_raw_fif(fif, preload=False, verbose=False)
    print(f"{fif.name}: {fif.stat().st_size / 1e6:.1f} MB on disk; read back -> {back}")
    print("types preserved:", {t: back.get_channel_types().count(t) for t in sorted(set(back.get_channel_types()))},
          "| montage:", back.get_montage() is not None, "| annotations:", back.annotations,
          "| description starts:", back.info["description"][:40], "...")

sf = syn.info["sfreq"]
events_syn = mne.find_events(syn, stim_channel="TRIG", verbose=False)         # rising edges of the TTL line
trig_samples = events_syn[:, 0]
diode = syn.get_data(picks="PHOTODIODE")[0]
baseline_sd = diode[: int(1.5 * sf)].std()                                    # dark-level noise before the first stimulus
threshold = 5 * baseline_sd
diode_samples = np.array([i + int(np.argmax(diode[i:i + int(0.5 * sf)] > threshold)) for i in trig_samples])
measured_ms = (diode_samples - trig_samples) / sf * 1000.0
print(f"{len(events_syn)} triggers found; photodiode onset = first sample above 5 x the dark-level noise ({threshold:.3f} a.u.)")
print(f"stimulus - trigger offset: mean {measured_ms.mean():.2f} ms, median {np.median(measured_ms):.1f} ms, "
      f"SD {measured_ms.std(ddof=1):.2f} ms, range {measured_ms.min():.0f}-{measured_ms.max():.0f} ms (n = {len(measured_ms)})")

k = 4                                                                          # one stimulus, close up
t0 = trig_samples[k] / sf
seg = syn.copy().crop(t0 - 0.05, t0 + 0.35)
fig, axes = plt.subplots(3, 1, figsize=(11, 6.5), sharex=True)
t = seg.times + (t0 - 0.05)
axes[0].plot(t, seg.get_data(picks="PHOTODIODE")[0], "k", lw=0.9)
axes[0].axhline(threshold, color="tab:red", lw=0.8, ls="--", label="detection threshold")
axes[0].set(ylabel="photodiode (a.u.)", title=f"Stimulus {k + 1}: trigger pulse, light onset and the EEG at Oz (synthetic recording)")
axes[1].plot(t, seg.get_data(picks="TRIG")[0], "k", lw=0.9)
axes[1].set(ylabel="TRIG (TTL)")
axes[2].plot(t, seg.get_data(picks="Oz")[0] * 1e6, "k", lw=0.9)
axes[2].set(ylabel="Oz (uV)", xlabel="Time (s)")
for ax in axes:
    ax.axvline(trig_samples[k] / sf, color="tab:blue", lw=1, label="trigger onset")
    ax.axvline(diode_samples[k] / sf, color="tab:orange", lw=1, label="light onset")
    ax.grid(alpha=0.3)
axes[0].legend(fontsize=8, loc="upper right")
fig.tight_layout()

fig, axes = plt.subplots(1, 2, figsize=(11, 3.2), gridspec_kw=dict(width_ratios=[2, 1]))
axes[0].plot(np.arange(1, len(measured_ms) + 1), measured_ms, "o-", ms=4, lw=0.8)
axes[0].axhline(measured_ms.mean(), color="tab:red", lw=0.8, label=f"mean {measured_ms.mean():.2f} ms")
axes[0].set(xlabel="stimulus number", ylabel="offset (ms)", title="Stimulus - trigger offset per trial (ms)")
axes[0].legend(fontsize=8)
axes[0].grid(alpha=0.3)
axes[1].hist(measured_ms, bins=np.arange(measured_ms.min() - 0.5, measured_ms.max() + 1.5, 1.0), color="tab:blue")
axes[1].set(xlabel="offset (ms)", ylabel="trials", title="Distribution across trials (ms)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

tc1, tc3 = helpers.event_timing_check(raw_r01), helpers.event_timing_check(raw_r03)
print(f"nb-0-3-first-load -- {DATASET} (EEGMMIDB; PhysioNet DOI 10.13026/C28G6P; ODC-By 1.0), subject {SUBJECT}")
print(f"EDF+ S001 R01: {tc1['sfreq_hz']:g} Hz | {len(mne.pick_types(raw_r01.info, eeg=True))} EEG channels | {tc1['duration_s']:.2f} s | "
      f"annotations {tc1['counts']}")
print(f"EDF+ S001 R03: {tc3['sfreq_hz']:g} Hz | {tc3['duration_s']:.2f} s | annotations {tc3['counts']} | onset intervals median "
      f"{tc3['interval_s']['median']:.2f} s (min {tc3['interval_s']['min']:.2f}, max {tc3['interval_s']['max']:.2f}) | "
      f"overlaps {tc3.get('n_overlapping_spans', 0)}, gaps {tc3.get('n_gaps_between_spans', 0)}")
if LEMON_OK:
    tcl = helpers.event_timing_check(raw_lemon)
    print(f"BrainVision ds-lemon {LEMON_SUBJECT} (first {LEMON_MAX_MINUTES} min): {tcl['sfreq_hz']:g} Hz | "
          f"{len(mne.pick_types(raw_lemon.info, eeg=True))} EEG + {len(mne.pick_types(raw_lemon.info, eog=True))} EOG channels | "
          f"{tcl['duration_s']:.1f} s | markers {tcl['counts']} | onset intervals median {tcl['interval_s']['median']:.3f} s | "
          f"FCz stored: {'FCz' in raw_lemon.ch_names}")
else:
    print("BrainVision ds-lemon: skipped in this run (not fetched)")
print("ds-erpcore: TODO(confirm) -- not in the catalog (spec 13.22); loader raises NotImplementedError")
print()
print("L0.3 exercise (SYNTHETIC recording: ds-eegbci S001 R01 EEG resampled to 1000 Hz + simulated PHOTODIODE and TRIG; "
      f"seed {SEED}; n = {N_STIM} stimuli)")
print(f"  planted stimulus-trigger offset : {OFFSET_MS:.1f} ms constant, +/- {JITTER_MS:g} ms uniform jitter per trial "
      f"(planted mean {planted.mean():.2f} ms)")
print(f"  measured from the photodiode    : mean {measured_ms.mean():.2f} ms | median {np.median(measured_ms):.1f} ms | "
      f"SD {measured_ms.std(ddof=1):.2f} ms | range {measured_ms.min():.0f}-{measured_ms.max():.0f} ms"
      "   <-- the number this notebook reports (answer key, tolerance +/- 2 ms)")
print("  the stimulus lags the trigger: events time-locked to TRIG must be shifted later by this amount")
print("TODO(confirm): draft values until the author reviews them; the recording is synthetic because no catalog dataset")
print("documents a photodiode channel.")