# Setup: dependencies, the shared helpers, non-interactive plotting, a quiet downloader.
import importlib.util
import subprocess
import sys
import time
import warnings
from collections import Counter
from pathlib import Path

_needed = ("mne", "scipy", "matplotlib", "pandas", "pooch", "sklearn")
_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", "yasa", "scikit-learn"]
    subprocess.check_call(_cmd)

_shared = next((d / "_shared" for d in (Path.cwd(), *Path.cwd().parents)
                if (d / "_shared" / "helpers_l7.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L7/ (or notebooks/) so that _shared/helpers_l7.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l7 as L7

import matplotlib.pyplot as plt
import numpy as np
import mne
import pooch

mne.set_log_level("WARNING")
# Silence the downloader at the source rather than scrubbing cache paths out of stored outputs
# afterwards: a re-execution would put them straight back (Phase 3 note in notebooks/README.md).
pooch.get_logger().setLevel("WARNING")
plt.rcParams["figure.dpi"] = 72

HAS_YASA = importlib.util.find_spec("yasa") is not None
EPOCH_S = float(L7.SLEEP_SPEC["epoch_s"])
PAD_EPOCHS = 60                       # 30 minutes either side of the sleep period, for section 4.2

print(f"MNE {mne.__version__}; helpers_l7 imported from notebooks/_shared")
print(f"yasa installed: {HAS_YASA}"
      + ("" if HAS_YASA else "  -> the documented fallback stager will run instead"))
print(L7.dataset_line("ds-sleep-edf"))

for k, v in L7.SLEEP_SPEC.items():
    print(f"  {k:16s} {v}")

L7.disk_report("before any download", folders={"Sleep-EDF cache": L7.sleep_edf_root()})
already_on_disk = [p for p in L7.sleep_edf_root().rglob("*") if p.is_file()] \
    if L7.sleep_edf_root().exists() else []
print(f"{len(already_on_disk)} Sleep-EDF files were cached before this run; they are kept and "
      f"anything this run fetches is deleted")

t0 = time.time()
paths = L7.fetch_sleep_night()
try:
    signal_rates = L7.edf_signal_rates(paths[0])  # from the header, before the file is deleted
    raw, ann = L7.read_sleep_night(*paths)        # preload=True: the signal is now in memory
    print(f"read in {time.time() - t0:.0f} s")
finally:
    freed = L7.drop_sleep_edf(paths, keep=already_on_disk, verbose=True)
L7.disk_report("after the download was deleted", folders={"Sleep-EDF cache": L7.sleep_edf_root()})

print(f"\n{raw.times[-1] / 3600:.2f} hours, {raw.n_times:,} samples at "
      f"{raw.info['sfreq']:.0f} Hz nominal")
print(f"channels: {raw.ch_names}")
demographics = L7.subject_demographics(raw)
print(f"subject (from the EDF header): {demographics}")
n_epochs = int(raw.n_times // (EPOCH_S * raw.info["sfreq"]))
print(f"{n_epochs} scoring epochs of {EPOCH_S:.0f} s")

print(f"MNE reports one sampling rate for the whole file: {raw.info['sfreq']:.0f} Hz, and every "
      f"channel comes back with {raw.n_times:,} samples.  The EDF HEADER says otherwise:")
print(f"\n{'channel':20s} {'header rate':>12s}   what it really is")
for ch, r in signal_rates.items():
    print(f"{ch:20s} {r:9.2f} Hz   "
          + ("a signal" if r >= 10 else "a 1 Hz channel filled in to 100 Hz on read"))
print(f"\nThis is not visible after the read — MNE brings every signal up to the file's highest rate, "
      f"so counting distinct values per second finds 100 in every channel.  The header is the only "
      f"place the difference survives, which is why helpers_l7.edf_signal_rates reads it directly.")
print(f"\n{L7.SLEEP_SPEC['emg_note']}")

rk = L7.expert_hypnogram(raw, ann)
counts = Counter(rk.tolist())
print(f"R&K stage per 30-s epoch, {len(rk)} epochs:")
for s in L7.RK_STAGES:
    if counts.get(s):
        print(f"  {s:2s}  {counts[s]:5d} epochs  {counts[s] * EPOCH_S / 60:7.1f} min  "
              f"{100 * counts[s] / len(rk):5.1f} %")
print(f"\nwake is {100 * counts.get('W', 0) / len(rk):.1f} % of the scored epochs. "
      f"{L7.SLEEP_SPEC['recording_note']}")

print(f"\nThe hypnogram FILE covers {sum(ann.duration) / 3600:.1f} h against the PSG's "
      f"{raw.times[-1] / 3600:.2f} h.")
tail = [(o, d, x) for o, d, x in zip(ann.onset, ann.duration, ann.description)
        if o + d > raw.times[-1]]
for o, d, x in tail:
    print(f"  the annotation starting at {o / 3600:.2f} h for {d / 3600:.2f} h is '{x}' — "
          f"{int(d / EPOCH_S)} epochs of it, all of them after the last sample")

print(L7.RK_TO_AASM_NOTE)
print("\nThe mapping this notebook applies:")
for k in L7.RK_STAGES:
    v = L7.RK_TO_AASM[k]
    print(f"  R&K {k:2s} -> {v if v else 'excluded by policy (see below)'}")
print("\nThe policies available for M and ?, and what each one does:")
for k, v in L7.STAGE_POLICIES.items():
    print(f"  [{k}] {v}")

aasm, keep, policy = L7.map_rk_to_aasm(rk, movement="exclude", unscored="exclude")
print(f"policy: movement={policy['movement']}, unscored={policy['unscored']}")
print(f"  movement-time epochs in this recording: {policy['n_movement']}")
print(f"  unscored epochs in this recording:      {policy['n_unscored']}")
print(f"  epochs excluded:                        {policy['n_excluded']}")
print(f"  epochs kept:                            {policy['n_kept']}")
print(f"\n  R&K 3 ({policy['n_rk3']} epochs) + R&K 4 ({policy['n_rk4']} epochs) -> "
      f"{policy['merged_3_and_4']} N3 epochs.  IRREVERSIBLE.")

print("\nOn this night the policy for M and ? changes nothing, because the recording contains none of "
      "either — every epoch inside the PSG carries a sleep stage.  It is still stated, because the "
      "same code on a different night of this dataset would change the denominator and therefore the "
      "kappa, and because a reader cannot tell which case they are in without being told:")
for mov in L7.STAGE_POLICIES:
    _, k2, p2 = L7.map_rk_to_aasm(rk, movement=mov, unscored=mov)
    print(f"  movement={mov:8s} unscored={mov:8s} -> {p2['n_kept']} epochs compared, "
          f"{p2['n_excluded']} excluded")

t0 = time.time()
auto, stager = L7.stage_night(raw, metadata={"age": demographics["age"],
                                             "male": demographics["male"]})
print(f"  {time.time() - t0:.1f} s, path = {stager['path']}")
print(f"  channels: {stager['channels']}")
if stager["path"] == "yasa":
    print(f"  classes: {stager['classes']}")
    print(f"  mean posterior of the chosen class: {np.mean(stager['confidence']):.3f}")
print(f"\nThe fallback this notebook would have used instead:")
for k in ("what_it_is", "what_it_is_not", "why_it_exists", "why_not_a_threshold_rule"):
    print(f"  [{k}] {L7.FALLBACK_STAGER[k]}")

print(f"\nautomatic stage counts: "
      + ", ".join(f"{k} {v}" for k, v in sorted(Counter(np.asarray(auto).tolist()).items())))

fig, _ = L7.plot_hypnograms(aasm, auto, mask=keep,
                            title="One night scored twice: expert (R&K mapped to AASM) above, "
                                  "automatic below, disagreement below that "
                                  "(stage per 30-s epoch)")
plt.show()

fig, _ = L7.plot_night_spectrogram(raw, fmax=25.0)
plt.show()

whole = L7.staging_agreement(aasm, auto, mask=keep)
print(f"kappa      {whole['kappa']:.4f}")
print(f"accuracy   {whole['accuracy']:.4f}")
print(f"epochs     {whole['n_epochs']}")
print(f"\n{'stage':6s} {'expert':>8s} {'automatic':>10s} {'recall':>8s} {'precision':>10s}")
for s in whole["labels"]:
    print(f"{s:6s} {whole['expert_counts'][s]:8d} {whole['auto_counts'][s]:10d} "
          f"{whole['recall'][s]:8.3f} {whole['precision'][s]:10.3f}")

fig, _ = L7.plot_confusion(whole,
                           title=f"Whole recording, {whole['n_epochs']} epochs "
                                 f"(count per cell; kappa = {whole['kappa']:.3f})")
plt.show()

bed, bed_info = L7.in_bed_mask(aasm, pad_epochs=PAD_EPOCHS, keep=keep)
print(f"window: epochs {bed_info['lo']}-{bed_info['hi']} = {bed_info['hours']:.2f} h "
      f"({bed_info['n_epochs']} epochs); {bed_info['definition']}")
in_bed = L7.staging_agreement(aasm, auto, mask=bed)
print(f"\n{'':22s} {'kappa':>8s} {'accuracy':>9s} {'epochs':>8s}")
print(f"{'whole recording':22s} {whole['kappa']:8.4f} {whole['accuracy']:9.4f} {whole['n_epochs']:8d}")
print(f"{'sleep period +- 30 min':22s} {in_bed['kappa']:8.4f} {in_bed['accuracy']:9.4f} "
      f"{in_bed['n_epochs']:8d}")

coarse_whole = L7.staging_agreement(L7.coarse_stages(aasm), L7.coarse_stages(auto), mask=keep,
                                    labels=["WAKE", "NREM", "REM"])
coarse_bed = L7.staging_agreement(L7.coarse_stages(aasm), L7.coarse_stages(auto), mask=bed,
                                  labels=["WAKE", "NREM", "REM"])
print(f"\nthree classes instead of five (WAKE / NREM / REM):")
print(f"{'whole recording':22s} {coarse_whole['kappa']:8.4f} {coarse_whole['accuracy']:9.4f} "
      f"{coarse_whole['n_epochs']:8d}")
print(f"{'sleep period +- 30 min':22s} {coarse_bed['kappa']:8.4f} {coarse_bed['accuracy']:9.4f} "
      f"{coarse_bed['n_epochs']:8d}")
print("\nFour numbers, one night, one stager.  Each is correct for the question it answers, and none "
      "of them can be quoted without saying which question that was.")

variants = []
if stager["path"] == "yasa":
    meta = {"age": demographics["age"], "male": demographics["male"]}
    for label, kw in (
            (f"{L7.SLEEP_SPEC['eeg']} + EOG + demographics (pre-specified)", dict(metadata=meta)),
            (f"{L7.SLEEP_SPEC['eeg']} + EOG, no demographics", dict(metadata=None)),
            (f"{L7.SLEEP_SPEC['eeg']} only, no EOG", dict(eog=None, metadata=meta)),
            ("EEG Pz-Oz + EOG + demographics", dict(eeg="EEG Pz-Oz", metadata=meta))):
        lab, _ = L7.yasa_stage(raw, **kw)
        a = L7.staging_agreement(aasm, lab, mask=keep)
        b = L7.staging_agreement(aasm, lab, mask=bed)
        variants.append((label, a["kappa"], b["kappa"], a["accuracy"]))
    print(f"{'variant':46s} {'kappa (all)':>12s} {'kappa (sleep)':>14s} {'accuracy':>9s}")
    for label, ka, kb, acc in variants:
        print(f"{label:46s} {ka:12.4f} {kb:14.4f} {acc:9.4f}")
    print("\nThe EOG channel is worth more than the demographics here, and the demographics are worth "
          "a NEGATIVE amount on this subject — dropping them raises kappa.  One night is one draw, so "
          "that is a fact about this recording and not a recommendation; it is reported because it is "
          "what happened.")
else:
    print("sensitivity runs need yasa; on the fallback path only the headline number is computed")

grapho = L7.detect_graphoelements(raw, aasm)
if grapho:
    print(f"\n  detector           {grapho['detector']}")
    print(f"  derivation         {grapho['channel']}")
    print(f"  label_source       {grapho['label_source']}")
    print(f"  N2+N3 time         {grapho['n2n3_minutes']:.1f} min")
    print(f"  spindles           {grapho['n_spindles']} "
          f"({grapho['spindle_density_per_min']:.3f} per minute of N2+N3)")
    print(f"  slow waves         {grapho['n_slow_waves']} "
          f"({grapho['slow_wave_density_per_min']:.3f} per minute of N2+N3)")
    print(f"\n  {grapho['caveat']}")
    print("  TODO(confirm): published spindle densities are quoted per minute of N2 on a CENTRAL "
          "derivation.  No published value is quoted here from memory, and this number is not "
          "comparable with one until somebody checks what the comparison would have to hold fixed.")

if grapho and grapho.get("spindles") is not None and len(grapho["spindles"]):
    sp = grapho["spindles"]
    t_h = sp["Start"].to_numpy() / 3600.0
    fig, axes = plt.subplots(2, 1, figsize=(11, 4.4), sharex=True)
    axes[0].hist(t_h, bins=44, color="C0")
    axes[0].set_ylabel("Spindles\n(count per 0.5 h)")
    sw_t = grapho["slow_waves"]["Start"].to_numpy() / 3600.0
    axes[1].hist(sw_t, bins=44, color="C1")
    axes[1].set_ylabel("Slow waves\n(count per 0.5 h)")
    axes[1].set_xlabel("Time from the start of the recording (hours)")
    for a in axes:
        a.grid(alpha=0.3)
    fig.suptitle(f"Algorithmic detections across the night on {grapho['channel']} "
                 f"(counts per half hour; no expert marks exist to check them against)", fontsize=10)
    fig.tight_layout()
    plt.show()

print("=" * 96)
print("nb-7-2-sleep — the numbers the L7.2 exercise asks for")
print("=" * 96)
print(f"\n[1] COHEN'S KAPPA = {whole['kappa']:.4f}  (accuracy {whole['accuracy']:.4f}, "
      f"{whole['n_epochs']} epochs)")
print("    and it does not exist without all of this:")
print(f"      record            {L7.SLEEP_SPEC['record']} ({raw.times[-1] / 3600:.2f} h, "
      f"{len(rk)} epochs of {EPOCH_S:.0f} s)")
print(f"      expert scoring    {L7.SLEEP_SPEC['scoring']}")
print(f"      automatic stager  {stager['stager']} on {', '.join(stager['channels'])}"
      + (f", metadata age {demographics['age']}, male {demographics['male']}"
         if stager["path"] == "yasa" else ""))
print(f"      mapping           W->WAKE, 1->N1, 2->N2, 3->N3, 4->N3, R->REM "
      f"(3 and 4 merged: {policy['n_rk3']} + {policy['n_rk4']} = {policy['merged_3_and_4']} epochs, "
      f"IRREVERSIBLE)")
print(f"      M and ? policy    movement={policy['movement']}, unscored={policy['unscored']}; "
      f"this night has {policy['n_movement']} movement and {policy['n_unscored']} unscored epochs, "
      f"so {policy['n_excluded']} epochs were excluded")
print(f"      window            the whole recording, {whole['n_epochs']} epochs")

print(f"\n[2] The same night on three other windows or vocabularies:")
print(f"      sleep period +- 30 min, five classes   kappa {in_bed['kappa']:.4f} "
      f"({in_bed['n_epochs']} epochs, {bed_info['hours']:.2f} h)")
print(f"      whole recording, three classes         kappa {coarse_whole['kappa']:.4f}")
print(f"      sleep period, three classes            kappa {coarse_bed['kappa']:.4f}")

print(f"\n[3] Where the disagreement is (whole recording, recall per expert stage):")
for s in whole["labels"]:
    print(f"      {s:5s} {whole['recall'][s]:6.3f}   "
          f"expert {whole['expert_counts'][s]:5d} epochs, automatic {whole['auto_counts'][s]:5d}")
print(f"      N1 is where staging always disagrees and it is {100 * whole['expert_counts']['N1'] / whole['n_epochs']:.1f} % "
      f"of this night, so it costs kappa little and would cost a clinical read a great deal.")

if variants:
    print(f"\n[4] What each pre-specified choice was worth (kappa over the whole recording):")
    for label, ka, kb, acc in variants:
        print(f"      {label:46s} {ka:.4f}")

if grapho:
    print(f"\n[5] Graphoelements, label_source ALGORITHMIC, no expert marks exist: "
          f"{grapho['n_spindles']} spindles ({grapho['spindle_density_per_min']:.3f}/min) and "
          f"{grapho['n_slow_waves']} slow waves ({grapho['slow_wave_density_per_min']:.3f}/min) "
          f"over {grapho['n2n3_minutes']:.0f} min of N2+N3 on {grapho['channel']}")

print(f"\n[6] The night: {100 * counts.get('W', 0) / len(rk):.1f} % wake, "
      f"{len(rk) * EPOCH_S / 3600:.1f} h. " + str(L7.SLEEP_SPEC["recording_note"]))
print("=" * 96)

# Nothing this notebook downloaded is left on disk; the fetch cell deleted it in a finally.
L7.disk_report("at the end", folders={"Sleep-EDF cache": L7.sleep_edf_root()})
left = [p.name for p in L7.sleep_edf_root().rglob("*") if p.is_file()] \
    if L7.sleep_edf_root().exists() else []
print(f"Sleep-EDF files still on disk: {sorted(left) or 'none'} "
      f"(pre-existing before this run: {sorted(p.name for p in already_on_disk) or 'none'})")
assert set(left) <= {p.name for p in already_on_disk}, "this run left a download behind"
print("no download made by this run was left behind")