# 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", "edfio")
_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")

DATASET, SUBJECT, RUN = "ds-eegbci", "S001", "R01"
raw = helpers.load_spine(DATASET, SUBJECT, RUN)


def print_checks(rows):
    for r in rows:
        print(f"[{r['status'].upper():4s}] {r['item']:30s} {r['value']}" + (f"\n       {r['note']}" if r["note"] else ""))


checks = helpers.first_look_checks(raw, DATASET)
print_checks(checks)

print("flat_channels (std < 0.5 uV):", helpers.flat_channels(raw) or "none")
dc = helpers.dc_offsets(raw)
print("dc_offsets, three largest |mean| (uV):", sorted(((round(v, 1), k) for k, v in dc.items()), key=lambda t: -abs(t[0]))[:3])
print("line_noise_peak:", helpers.line_noise_peak(raw, mains_hz=60))
print("alpha_peak:", helpers.alpha_peak(raw))
print("event_timing_check:", helpers.event_timing_check(raw))
cons = helpers.channel_position_consistency(raw)
print("channel_position_consistency, the three lowest rho:", [(c["name"], round(c["rho"], 2), c["best_match"]) for c in cons[:3]])

import hashlib
import json
import tempfile
from datetime import date
from IPython.display import HTML, display

from mne.datasets import eegbci

root = helpers.data_dir()
edf_path = Path(eegbci.load_data(1, [1], update_path=False, verbose=False, **({"path": str(root)} if root else {}))[0])
bids_notes = [
    "EEG-BIDS names a recording sub-<label>[_ses-<label>]_task-<label>[_run-<index>]_eeg.<ext> and describes it in sidecars: "
    "_eeg.json (acquisition metadata), _channels.tsv (names, types, units), _electrodes.tsv (positions), _events.tsv (events).",
    f"This file is not BIDS-named ({edf_path.name}); the catalog lists a BIDS mirror of the dataset (OpenNeuro ds004362) -- see section 3.",
]
html = helpers.first_look_html(raw, DATASET, title=f"First look: {DATASET} {SUBJECT} {RUN}", bids_notes=bids_notes)
display(HTML(html))

qc_log = {
    "date": date.today().isoformat(),
    "dataset": DATASET,
    "citation": helpers.DATASETS[DATASET]["citation"],
    "doi": helpers.DATASETS[DATASET]["doi"],
    "license": helpers.DATASETS[DATASET]["license"],
    "file": edf_path.name,
    "sha256_16": hashlib.sha256(edf_path.read_bytes()).hexdigest()[:16],
    "software": f"MNE {mne.__version__}",
    "checks": checks,
    "decision": "TODO: your judgment after looking at the traces (nb-0-4) and the artifact atlas (nb-0-5)",
}
with tempfile.TemporaryDirectory() as tmp:
    (Path(tmp) / "first-look-S001R01.html").write_text(html, encoding="utf-8")
    (Path(tmp) / "qc-log-S001R01.json").write_text(json.dumps(qc_log, indent=1), encoding="utf-8")
    print("written to a temporary folder:", ", ".join(f"{p.name} ({p.stat().st_size / 1e3:.0f} kB)" for p in sorted(Path(tmp).iterdir())))

for ds_id, prefix, label in (("ds004362", "sub-001/eeg/", "ds-eegbci BIDS mirror (OpenNeuro ds004362, catalog)"),
                             ("ds004584", "sub-001/eeg/", "ds-iowapd (OpenNeuro ds004584, CC0)")):
    files = helpers.list_openneuro_files(ds_id, prefix)
    print(f"{label}: {prefix}")
    if not files:
        print("   (listing not reachable in this run)")
        continue
    for f in files[:7]:
        print(f"   {f['name']:52s} {f['size'] / 1e3:9.1f} kB")
    if len(files) > 7:
        print(f"   ... {len(files)} files in total")

def make_corrupted_copy(raw_in, folder):
    """A labelled synthetic derivative of raw_in with two planted problems; returns the path of the written file."""
    data = raw_in.get_data().copy()
    names = list(raw_in.ch_names)
    # planted problem 1: the data of two channels are swapped under their labels (a mislabelled channel pair)
    i, j = names.index("Fp1"), names.index("O1")
    data[[i, j]] = data[[j, i]]
    # planted problem 2: the header claims 200 Hz for data sampled at 160 Hz
    info = mne.create_info(names, 200.0, "eeg")
    bad = mne.io.RawArray(data, info, verbose=False)
    bad.set_annotations(mne.Annotations([0.0], [bad.times[-1]], ["T0"]))
    bad.info["description"] = "SYNTHETIC derivative of ds-eegbci S001 R01 with two planted problems (L0.6 exercise)"
    try:
        import edfio  # noqa: F401  (MNE's EDF export needs it)
        out = Path(folder) / "S001R01-corrupted.edf"
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")       # EDF export pads the last data record; not our concern here
            bad.export(out, fmt="edf", overwrite=True, verbose=False)
    except ImportError:
        out = Path(folder) / "S001R01-corrupted_raw.fif"
        bad.save(out, overwrite=True, verbose=False)
        print("edfio is not installed: the copy is written as FIF instead of EDF")
    return out


tmpdir = tempfile.TemporaryDirectory()             # deleted in the last cell of the notebook
bad_path = make_corrupted_copy(raw, tmpdir.name)
if bad_path.suffix == ".edf":
    handed = mne.io.read_raw_edf(bad_path, preload=True, verbose=False)
else:
    handed = mne.io.read_raw_fif(bad_path, preload=True, verbose=False)
handed.set_montage("standard_1005", on_missing="raise", verbose=False)
print(f"{bad_path.name}: {bad_path.stat().st_size / 1e6:.1f} MB")
print(handed)
print(f"sfreq {handed.info['sfreq']:g} Hz | {handed.info['nchan']} channels | {handed.times[-1] + 1 / handed.info['sfreq']:.1f} s | "
      f"annotations {helpers.event_timing_check(handed)['counts']}")

html_bad = helpers.first_look_html(handed, DATASET, title=f"First look: {bad_path.name} (a file you were handed)")
display(HTML(html_bad))
checks_bad = helpers.first_look_checks(handed, DATASET)
print("rows that are not OK:")
print_checks([r for r in checks_bad if r["status"] in ("warn", "fail")])

fig, axes = plt.subplots(1, 2, figsize=(14, 4), gridspec_kw=dict(width_ratios=[1.6, 1]))
helpers.plot_traces(handed, ["Fp1", "Fp2", "O1", "O2"], t0=5, duration=10, spacing_uV=250, ax=axes[0],
                    title=f"{bad_path.name}: which channel carries the blinks?")
axes[0].set_xlabel("Time (s, at the header's 200 Hz)")
helpers.plot_psd(raw, ["Cz"], fmin=1, fmax=raw.info["sfreq"] / 2 - 1, ax=axes[1], label_prefix="original 160 Hz: ", title="")
helpers.plot_psd(handed, ["Cz"], fmin=1, fmax=handed.info["sfreq"] / 2 - 1, ax=axes[1], label_prefix="header 200 Hz: ",
                 title="Cz PSD, original vs corrupted header, each up to its own Nyquist (dB re 1 uV^2/Hz)")
for f, colour in ((60, "tab:blue"), (75, "tab:orange")):
    axes[1].axvline(f, color=colour, lw=0.8, ls="--")
axes[1].legend(fontsize=8)
for ax in axes:
    ax.title.set_fontsize(9)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

p88 = Path(eegbci.load_data(88, [3], update_path=False, verbose=False, **({"path": str(root)} if root else {}))[0])
raw88 = mne.io.read_raw_edf(p88, preload=True, verbose=False)
eegbci.standardize(raw88)
raw88.set_montage("standard_1005", on_missing="ignore", verbose=False)
print(f"{p88.name}: {raw88}")
checks88 = helpers.first_look_checks(raw88, DATASET)
print_checks([r for r in checks88 if r["status"] in ("warn", "fail")])
print("\nthe first eight annotations (onset, duration, end, in s):")
ann = raw88.annotations
for k in range(8):
    print(f"   {ann.description[k]:3s} {ann.onset[k]:8.3f} {ann.duration[k]:7.3f} {ann.onset[k] + ann.duration[k]:8.3f}")
t88 = helpers.event_timing_check(raw88)
raw_r03 = helpers.load_spine(DATASET, SUBJECT, "R03")
t01 = helpers.event_timing_check(raw_r03)
print(f"\nS088 R03: {t88['sfreq_hz']:g} Hz, {t88['n_annotations']} annotations, overlapping spans {t88['n_overlapping_spans']}, "
      f"gaps between spans {t88['n_gaps_between_spans']}")
print(f"S001 R03: {t01['sfreq_hz']:g} Hz, {t01['n_annotations']} annotations, overlapping spans {t01['n_overlapping_spans']}, "
      f"gaps between spans {t01['n_gaps_between_spans']}  (the same run of an unaffected subject)")

fails = {r["item"]: r["value"] for r in checks_bad if r["status"] == "fail"}
ln_bad = helpers.line_noise_peak(handed, mains_hz=60)
print(f"nb-0-6-first-look-report -- {DATASET} (EEGMMIDB; PhysioNet DOI 10.13026/C28G6P; ODC-By 1.0)")
print(f"baseline S001 R01: {raw.info['sfreq']:g} Hz | {len(mne.pick_types(raw.info, eeg=True))} EEG channels | "
      f"{raw.times[-1] + 1 / raw.info['sfreq']:.2f} s | flat channels: {len(helpers.flat_channels(raw))} | "
      f"largest |DC| {max(abs(v) for v in helpers.dc_offsets(raw).values()):.1f} uV | "
      f"line {helpers.line_noise_peak(raw, mains_hz=60)['freq_hz']:.2f} Hz | "
      f"checklist: {sum(r['status'] == 'ok' for r in checks)} ok, {sum(r['status'] == 'warn' for r in checks)} warn, "
      f"{sum(r['status'] == 'fail' for r in checks)} fail")
print()
print(f"L0.6 exercise -- the two planted problems in {bad_path.name} (SYNTHETIC derivative of S001 R01; answer key):")
print(f"  1. mislabelled channel: the labels Fp1 and O1 are swapped -- the channel labelled O1 carries the frontal-polar "
      f"(blink) data and the channel labelled Fp1 the occipital data   <-- answer 1")
print(f"     tell-tale: channel-position consistency flags {', '.join(n for n in ('O1', 'Fp1') if n in fails.get('channel-position consistency', ''))}; "
      f"the blinks appear on 'O1' in the traces")
print(f"  2. wrong sampling rate in the header: the header says {handed.info['sfreq']:g} Hz; the data were sampled at "
      f"{raw.info['sfreq']:g} Hz   <-- answer 2")
print(f"     tell-tale: the sharpest spectral line is at {ln_bad['freq_hz']:.2f} Hz instead of 60 Hz "
      f"({handed.info['sfreq']:g} x 60 / {ln_bad['freq_hz']:.0f} = {handed.info['sfreq'] * 60 / ln_bad['freq_hz']:.0f} Hz), "
      f"the file lasts {handed.times[-1] + 1 / handed.info['sfreq']:.1f} s instead of about a minute, and the catalog documents 160 Hz")
print(f"  checklist on the corrupted copy: {sum(r['status'] == 'fail' for r in checks_bad)} FAIL rows "
      f"({', '.join(r['item'] for r in checks_bad if r['status'] == 'fail')})")
print()
print(f"second file, real -- S088 R03 (catalog caveat): sampling rate {t88['sfreq_hz']:g} Hz (documented 160 Hz); "
      f"{t88['n_annotations']} annotations with {t88['n_overlapping_spans']} overlapping spans and {t88['n_gaps_between_spans']} gaps; "
      f"S001 R03 for comparison: {t01['sfreq_hz']:g} Hz, 0 overlaps, 0 gaps")
print("TODO(confirm): draft values until the author reviews them.")
tmpdir.cleanup()                                   # remove the temporary folder with the corrupted copy