# 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/capstones/ (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")

FULL_COHORT = False          # False: the documented subset (one file per spine dataset). True: the sets below, local runs only.
LEMON_MAX_MINUTES = 5        # the ds-lemon data file is fetched as a prefix of this many minutes (None = whole recording, 302 MB)
LEMON_DELETE_AFTER_USE = FULL_COHORT   # in a full-cohort run, delete each LEMON subject's files after its report (disk)

SUBSET = {
    "ds-eegbci": [("S001", "R01")],
    "ds-lemon": ["sub-010002"],
    "ds-erpcore": ["TODO(confirm)"],
}
# Documented full-cohort sets (spec 11: 10-20 subjects; none of the eegbci subjects below is in the catalog's exclusion
# list; the LEMON ids are the first ten initial ids of the mirror's name_match.csv, each fetched as a 5-min prefix).
FULL = {
    "ds-eegbci": [(f"S{i:03d}", "R01") for i in range(1, 13)],
    "ds-lemon": [f"sub-0100{i:02d}" for i in range(2, 12)],
    "ds-erpcore": ["TODO(confirm)"],
}
PLAN = FULL if FULL_COHORT else SUBSET

OUT_DIR = _shared.parent / "capstones"          # next to this notebook, found relative to the working directory
REPORT_NAME = "nb-c0-recording-report.html"
print("plan:", {k: (v if len(v) <= 3 else f"{len(v)} files") for k, v in PLAN.items()})
print(f"report file: {REPORT_NAME} (written into the capstones folder)")

from datetime import date
from html import escape
from IPython.display import HTML, display

TECHNICAL = ("emg", "pop", "flat")          # spans that count toward "artifact coverage" (noisy is a channel flag, not a span)
OCULAR = ("blink", "saccade")               # reported, expected in eyes-open data, handled later (L2.6), not a caution by themselves


def coverage_fraction(events, duration_s, labels=TECHNICAL):
    """Fraction of the recording covered by the union of the events with the given labels."""
    spans = sorted((e["onset_s"], e["onset_s"] + e["duration_s"]) for e in events if e["label"] in labels)
    covered, cur = 0.0, None
    for a, b in spans:
        if cur is None or a > cur[1]:
            if cur is not None:
                covered += cur[1] - cur[0]
            cur = [a, b]
        else:
            cur[1] = max(cur[1], b)
    if cur is not None:
        covered += cur[1] - cur[0]
    return covered / duration_s if duration_s else 0.0


def recommend(checks, events, duration_s):
    """The go/no-go template: a decision and generated reasons (to be rewritten after viewing the traces)."""
    fails = [c for c in checks if c["status"] == "fail"]
    warns = [c for c in checks if c["status"] == "warn"]
    counts = {lab: sum(e["label"] == lab for e in events) for lab in helpers.ARTIFACT_LABELS}
    cov = coverage_fraction(events, duration_s)
    ocular = coverage_fraction(events, duration_s, OCULAR)
    reasons = []
    if fails:
        decision = "NO-GO"
        reasons += [f"checklist FAIL -- {c['item']}: {c['value']}" for c in fails]
        reasons.append("resolve the metadata/integrity problems (or document them) before any analysis")
    else:
        flags = []
        if counts["noisy"]:
            flags.append(f"{counts['noisy']} noisy channel(s): {', '.join(e['channels'][0] for e in events if e['label'] == 'noisy')}"
                         " -- exclude or interpolate (L2.2) before anything else")
        if counts["flat"] >= 2:
            flags.append(f"{counts['flat']} flat-channel events")
        if cov > 0.20:
            flags.append(f"technical artifact coverage {100 * cov:.0f} % of the recording (EMG/pop/flat)")
        decision = "CAUTION" if flags else "GO"
        reasons += flags or [f"no failed checklist item; technical artifact coverage {100 * cov:.0f} % (EMG/pop/flat), no noisy channel"]
        reasons += [f"checklist WARN -- {c['item']}: {c['value']}" for c in warns]
        reasons.append(f"ocular events cover {100 * ocular:.0f} % ({counts['blink']} blinks, {counts['saccade']} saccades): expected with the eyes open, "
                       "to be handled by EOG regression or ICA (L2.6), not a reason to reject")
        reasons.append(f"{counts['emg']} EMG bursts, {counts['drift']} drift windows, {counts['pop']} pops, "
                       f"{counts['line-noise']} line-noise windows located (algorithmic); drift and line noise are filtering matters (L1.6, L2.4)")
    return {"decision": decision, "reasons": reasons, "coverage": cov, "ocular": ocular, "counts": counts}


def recording_report(raw, dataset, label):
    """One-page report (HTML section) + the numbers behind it."""
    cat = helpers.DATASETS.get(dataset, {})
    mains = cat.get("mains_hz")
    sf, dur = float(raw.info["sfreq"]), raw.n_times / raw.info["sfreq"]
    meta = helpers.first_look_report(raw, dataset)
    checks = helpers.first_look_checks(raw, dataset)
    timing = helpers.event_timing_check(raw)
    events = helpers.detect_artifacts(raw, mains_hz=mains)
    psd = helpers.psd_summary(raw, mains_hz=mains)
    rec = recommend(checks, events, dur)
    types = raw.get_channel_types()
    tcount = ", ".join(f"{types.count(t)} {t}" for t in sorted(set(types)))

    # figures
    fig_m, ax = plt.subplots(figsize=(4.6, 4.6))
    raw.plot_sensors(show_names=True, axes=ax, show=False)
    ax.set_title(f"{label}: {len(mne.pick_types(raw.info, eeg=True))} EEG channels, template 10-10 positions", fontsize=9)
    fig_t, ax = plt.subplots(figsize=(11, 3.2))
    helpers.artifact_timeline(events, dur, ax=ax, title=f"{label}: artifact timeline (algorithmic labels; s)")
    picks = [c for c in ("Fp1", "Fz", "Cz", "Pz", "O1", "O2") if c in raw.ch_names]
    fig_p, ax = plt.subplots(figsize=(8, 3.4))
    helpers.plot_psd(raw, picks, fmin=0.5, fmax=min(sf / 2, 100), ax=ax, title=f"{label}: Welch PSD, 4-s segments (dB re 1 uV^2/Hz)")
    imgs = {k: helpers.fig_to_base64(f) for k, f in (("montage", fig_m), ("timeline", fig_t), ("psd", fig_p))}
    for f in (fig_m, fig_t, fig_p):
        plt.close(f)

    # tables
    ann = raw.annotations
    ev_rows = [(f"{o:.3f}", f"{d:.3f}", str(desc)) for o, d, desc in list(zip(ann.onset, ann.duration, ann.description))[:8]]
    count_rows = [(k, v) for k, v in timing["counts"].items()]
    art_rows = []
    for lab in helpers.ARTIFACT_LABELS:
        evs = [e for e in events if e["label"] == lab]
        if not evs:
            continue
        chans = sorted({c for e in evs for c in e["channels"]})
        art_rows.append((lab, len(evs), f"{sum(e['duration_s'] for e in evs):.1f}", ", ".join(f"{e['onset_s']:.1f}" for e in evs[:6]) + (" ..." if len(evs) > 6 else ""),
                         ", ".join(chans[:8]) + (" ..." if len(chans) > 8 else ""), evs[0]["detector"]))
    tl = timing.get("interval_s")
    timing_text = (f"{timing['n_samples']} samples / {sf:g} Hz = {dur:.1f} s; {timing['n_annotations']} annotations"
                   + (f"; onset intervals median {tl['median']:.3f} s (min {tl['min']:.3f}, max {tl['max']:.3f}); "
                      f"overlaps {timing.get('n_overlapping_spans', 0)}, gaps {timing.get('n_gaps_between_spans', 0)}, "
                      f"long gaps {timing.get('n_long_gaps', 0)}" if tl else ""))
    status_colour = {"ok": "#2a7", "warn": "#d90", "fail": "#c33", "info": "#888"}
    check_rows = "".join(f"<tr><td style='padding:2px 8px'>{escape(c['item'])}</td><td style='padding:2px 8px;color:{status_colour[c['status']]};font-weight:bold'>"
                         f"{c['status'].upper()}</td><td style='padding:2px 8px'>{escape(str(c['value']))}</td><td style='padding:2px 8px;color:#555'>{escape(c['note'])}</td></tr>" for c in checks)
    decision_colour = {"GO": "#2a7", "CAUTION": "#d90", "NO-GO": "#c33"}[rec["decision"]]
    html = f"""
<section style="font-family:sans-serif;max-width:64em;page-break-after:always;border-top:2px solid #999;padding-top:1em">
<h2>Recording report -- {escape(label)}</h2>
<p><b>{escape(cat.get('name', dataset))}</b> -- {escape(cat.get('citation', ''))}; DOI {escape(str(cat.get('doi', 'TODO(confirm)')))}; license {escape(str(cat.get('license', 'TODO(confirm)')))}.<br>
File: {escape(str(raw.filenames[0].name if raw.filenames and raw.filenames[0] else label))}; report generated by nb-c0-recording-report on {date.today().isoformat()} with MNE {mne.__version__}.
Artifact labels are <b>algorithmic</b> (threshold rules, not reviewed); the recommendation is a <b>template</b>.</p>
<h3>1. Montage</h3>
<img alt="sensor layout" src="data:image/png;base64,{imgs['montage']}" style="max-width:340px">
<p>Documented reference: {escape(str(cat.get('reference', 'not documented -- TODO(confirm)')))}. Montage: {escape(str(cat.get('montage', 'TODO(confirm)')))}.</p>
<h3>2. Sampling and duration</h3>
{helpers.html_table([
    ("sampling rate (header)", f"{sf:g} Hz (Nyquist {sf / 2:g} Hz)", f"documented: {cat.get('sfreq', 'TODO(confirm)')} Hz"),
    ("channels", f"{raw.info['nchan']} ({tcount})", f"documented: {cat.get('n_channels', 'TODO(confirm)')}"),
    ("duration", f"{dur:.1f} s ({dur / 60:.2f} min)", "documented run length: see the dataset page"),
    ("header filters", f"high-pass {meta['highpass_hz_in_header']:g} Hz, low-pass {meta['lowpass_hz_in_header']:g} Hz", f"documented: {cat.get('hardware_filters', 'TODO(confirm)')}"),
    ("mains", f"{cat.get('mains_hz', 'TODO(confirm)')} Hz (documented)", ""),
], header=("item", "value", "documentation"))}
<h3>3. Events</h3>
{helpers.html_table(count_rows, header=("description", "count"))}
<p>{escape(timing_text)}</p>
{helpers.html_table(ev_rows, header=("onset (s)", "duration (s)", "description")) if ev_rows else "<p>no annotations</p>"}
<h3>4. Artifact timeline (algorithmic)</h3>
<img alt="artifact timeline" src="data:image/png;base64,{imgs['timeline']}" style="max-width:100%">
{helpers.html_table(art_rows, header=("type", "n", "total s", "first onsets (s)", "channels", "detector")) if art_rows else "<p>nothing flagged</p>"}
<p>Coverage of the recording by technical spans (EMG/pop/flat): {100 * rec['coverage']:.0f} %; by ocular spans (blink/saccade): {100 * rec['ocular']:.0f} %.</p>
<h3>5. First-look PSD</h3>
<img alt="PSD" src="data:image/png;base64,{imgs['psd']}" style="max-width:100%">
<p>{escape(psd['alpha_text'])} {escape(psd['line_text'])}</p>
<h3>6. Checklist</h3>
<table style="font-size:0.9em;border-collapse:collapse"><tr><th style='text-align:left;padding:2px 8px'>item</th><th style='text-align:left;padding:2px 8px'>status</th><th style='text-align:left;padding:2px 8px'>value</th><th style='text-align:left;padding:2px 8px'>note</th></tr>{check_rows}</table>
<h3>7. Go / no-go (template -- rewrite after viewing the traces)</h3>
<p style="font-size:1.2em"><b style="color:{decision_colour}">{rec['decision']}</b></p>
<ul>{''.join(f'<li>{escape(r)}</li>' for r in rec['reasons'])}</ul>
<p style="color:#555">Template rules: any failed checklist item -> NO-GO; a noisy channel, two or more flat-channel events, or more than 20 % technical coverage (EMG/pop/flat) -> CAUTION; otherwise GO. Ocular events, drift and line noise are reported but do not decide. Your reasons should name what you saw in the traces (nb-0-4) and which atlas entries (nb-0-5) the flagged spans match.</p>
</section>"""
    summary = {"label": label, "dataset": dataset, "sfreq": sf, "n_channels": raw.info["nchan"], "duration_s": dur,
               "event_counts": timing["counts"], "artifact_counts": rec["counts"], "coverage": rec["coverage"], "ocular": rec["ocular"],
               "alpha": psd["alpha"], "line_noise": psd["line_noise"], "alpha_text": psd["alpha_text"],
               "line_text": psd["line_text"], "decision": rec["decision"], "reasons": rec["reasons"],
               "n_fail": sum(c["status"] == "fail" for c in checks), "n_warn": sum(c["status"] == "warn" for c in checks)}
    return summary, html


def stub_report(dataset, label, message):
    """A report section for a dataset that cannot be loaded yet."""
    return ({"label": label, "dataset": dataset, "decision": "TODO(confirm)", "reasons": [message]},
            f"<section style='font-family:sans-serif;max-width:64em;border-top:2px solid #999;padding-top:1em'>"
            f"<h2>Recording report -- {escape(label)}</h2><p>{escape(message)}</p></section>")

import shutil

summaries, sections = [], []

for subject, run in PLAN["ds-eegbci"]:
    raw = helpers.load_spine("ds-eegbci", subject, run)
    s, h = recording_report(raw, "ds-eegbci", f"ds-eegbci {subject} {run} ({helpers.EEGBCI_RUNS[run]})")
    summaries.append(s)
    sections.append(h)
    display(HTML(h))

for subject in PLAN["ds-lemon"]:
    raw = helpers.load_lemon_raw(subject, max_minutes=LEMON_MAX_MINUTES, preload=True)
    if raw is None:
        s, h = stub_report("ds-lemon", f"ds-lemon {subject}", "Skipped in this run: the raw files could not be fetched from the public "
                           "mirror (see the message above); re-run with network access, or set LEMON_MAX_MINUTES lower.")
    else:
        s, h = recording_report(raw, "ds-lemon", f"ds-lemon {subject} (raw, first {LEMON_MAX_MINUTES} min)" if LEMON_MAX_MINUTES
                                else f"ds-lemon {subject} (raw)")
        if LEMON_DELETE_AFTER_USE:
            folder = helpers.download_root() / "MNE-lemon-data" / "EEG_Raw_BIDS_ID" / subject
            del raw
            shutil.rmtree(folder, ignore_errors=True)
            print(f"deleted the files of {subject} after use")
    summaries.append(s)
    sections.append(h)
    display(HTML(h))

s, h = stub_report("ds-erpcore", "ds-erpcore (TODO(confirm))",
                   "ERP CORE is not in the site's dataset catalog (spec 13.22): license and per-subject download are TODO(confirm); "
                   "no file is fetched and no report is produced until the author resolves it.")
summaries.append(s)
sections.append(h)
display(HTML(h))

OUT_DIR.mkdir(parents=True, exist_ok=True)
page = ("<!DOCTYPE html><html><head><meta charset='utf-8'><title>C0 recording reports</title></head><body>"
        f"<h1 style='font-family:sans-serif'>Capstone C0 -- recording reports ({date.today().isoformat()})</h1>"
        + "".join(sections) + "</body></html>")
(OUT_DIR / REPORT_NAME).write_text(page, encoding="utf-8")
print(f"\nwritten: {REPORT_NAME} ({(OUT_DIR / REPORT_NAME).stat().st_size / 1e3:.0f} kB, {len(sections)} sections) next to the notebook")

print("nb-c0-recording-report -- capstone C0 (subset run)" if not FULL_COHORT else "nb-c0-recording-report -- capstone C0 (FULL_COHORT run)")
for s in summaries:
    print(f"\n{s['label']}:")
    if "sfreq" not in s:
        print(f"  {s['decision']}: {s['reasons'][0]}")
        continue
    print(f"  {s['sfreq']:g} Hz | {s['n_channels']} channels | {s['duration_s']:.1f} s | events {s['event_counts']}")
    print(f"  artifacts (algorithmic): {s['artifact_counts']} | technical coverage {100 * s['coverage']:.0f} % | ocular coverage {100 * s['ocular']:.0f} %")
    print(f"  {s['alpha_text']}")
    print(f"  {s['line_text']}")
    print(f"  checklist: {s['n_fail']} fail, {s['n_warn']} warn | recommendation (template): {s['decision']} -- {s['reasons'][0]}")
print(f"\nrendered HTML: {REPORT_NAME} in the capstones folder; every label is label_source: algorithmic; TODO(confirm) at author review.")