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

_needed = ("mne", "scipy", "matplotlib", "pandas", "pooch", "specparam")
_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", "pooch>=1.8", "specparam==2.0.0rc4"]
    subprocess.check_call(_cmd)

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

HAVE_L7 = importlib.util.find_spec("helpers_l7") is not None
if HAVE_L7:
    import helpers_l7 as L7

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import mne
from scipy import signal as sps, stats

mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72

# Quiet the downloader: pooch logs the ABSOLUTE cache path at INFO and a stored notebook may not
# contain one (scripts/scrub-notebooks.py is a CI gate).  Every cell below prints the file names it
# fetched and the free disk before and after, so nothing is hidden by this.
try:
    import pooch

    pooch.get_logger().setLevel("WARNING")
except Exception:
    pass

SEED = L6.SEED
ALPHA = L6.ALPHA
rng = np.random.default_rng(SEED)

FULL_RUN = False
N_LEMON = 5 if not FULL_RUN else 12        # ds-lemon subjects (split-half only)
N_SRM = 10 if not FULL_RUN else 42         # ds-srm subjects with BOTH sessions (42 is the cohort cap)
N_DORT = 8 if not FULL_RUN else 40         # ds-dortmund subjects with BOTH sessions (208 is the cap)
N_PD = 8 if not FULL_RUN else 49           # ds-iowapd participants PER GROUP (49 controls is the cap)

print(f"MNE {mne.__version__}; helpers imported from notebooks/_shared (helpers_l7 present: {HAVE_L7})")
print(f"FULL_RUN = {FULL_RUN}: ds-lemon {N_LEMON}, ds-srm {N_SRM}, ds-dortmund {N_DORT}, "
      f"ds-iowapd {N_PD} per group")
print(f"Download budget: roughly {N_LEMON * 25 + N_SRM * 63 + N_DORT * 72 + N_PD * 2 * 36} MB in total, "
      f"fetched one recording at a time and deleted before the next, so at most one is ever on disk.")

MEASURE = {
    "harmonisation": "0.5-45 Hz zero-phase FIR band-pass on every dataset, so that recordings with "
                     "50 Hz and 60 Hz mains and with different online filters are compared over the "
                     "same band; the 40 Hz top of the exponent fit sits well inside it",
    "segment": "up to 120 s from t = 10 s (or the whole recording when it is shorter)",
    "bad channels": "an electrode whose robust amplitude is more than 6 MAD above the montage median, "
                    "or below a twentieth of it, is dropped BEFORE epoching (never more than a third "
                    "of the montage)",
    "epochs": "4-s non-overlapping epochs; an epoch is dropped if any retained channel's demeaned "
              "amplitude exceeds 150 uV",
    "reference": "average over the retained scalp channels, applied offline",
    "spectrum": "Welch, one 4-s Hann window per epoch (0.25 Hz resolution), averaged over epochs",
    "iaf": "largest residual in 7-13 Hz above a log-log line fitted over 2-40 Hz excluding 5-15 Hz, "
           "minimum 3 dB, over the posterior electrodes present",
    "theta_beta": "sum of PSD over 4-8 Hz divided by sum over 13-30 Hz, averaged over all retained channels",
    "exponent": "specparam fixed-mode aperiodic exponent over 1-40 Hz on the posterior spectrum",
}
EPOCH_S = 4.0
EPOCH_REJECT_UV = 150.0
SEGMENT_T0, SEGMENT_DUR = 10.0, 120.0
ALPHA_BAND, FIT_RANGE, MIN_DB = (7.0, 13.0), (2.0, 40.0), 3.0
THETA, BETA = (4.0, 8.0), (13.0, 30.0)
EXPONENT_RANGE = (1.0, 40.0)
POSTERIOR = ("O1", "Oz", "O2", "PO3", "POz", "PO4", "PO7", "PO8", "PO9", "PO10",
             "P1", "P2", "P3", "Pz", "P4", "P5", "P6", "P7", "P8")

for k, v in MEASURE.items():
    print(f"   {k:11s} : {v}")

NON_SCALP = L6.NON_SCALP_NAMES


HP_HZ, LP_HZ = 0.5, 45.0


def band_sum(freqs, psd, band):
    m = (freqs >= band[0]) & (freqs <= band[1])
    return float(np.trapezoid(np.asarray(psd)[..., m], np.asarray(freqs)[m], axis=-1).mean())


#: Largest plausible robust AC amplitude for scalp EEG, in the units MNE believes it is returning.
#: Ten millivolts of *typical* deviation is not a scalp signal in any montage, so a recording above
#: this is being read at the wrong scale.
MAX_PLAUSIBLE_UV = 1e4


def ensure_microvolts(raw):
    """Detect and correct a recording MNE has read at the wrong scale.

    An EDF states its physical dimension per signal.  When that field is blank, MNE has no unit to
    apply and returns the file's own numbers as volts.  ds-srm's EDFs are one such file: the BIDS
    channels.tsv says uV, the EDF header says nothing, and the data come back 10^6 too large.

    The test is on the data, not on the header, so it is the same test for every dataset: remove each
    channel's median (a DC-coupled amplifier's electrode offsets are real and are not the signal) and
    look at the typical remaining deviation.  Returns the factor applied and both amplitudes.
    """
    x = raw.get_data(picks="eeg")
    ac = float(np.median(np.abs(x - np.median(x, axis=1, keepdims=True)))) * 1e6
    factor = 1.0
    if ac > MAX_PLAUSIBLE_UV:
        factor = 1e-6
        raw.apply_function(lambda d: d * factor, picks="all")
    return {"scale_factor": factor, "ac_uv_before": ac, "ac_uv_after": ac * factor}


def clean_scalp(raw):
    """Keep scalp electrodes with a montage position, drop flat ones, band-pass.  (raw, notes)."""
    raw = raw.copy()
    scale = ensure_microvolts(raw)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        raw.pick("eeg", exclude=[])
        raw.set_montage("standard_1005", on_missing="ignore", verbose=False)
        mont = raw.get_montage()
        placed = set(mont.ch_names) if mont is not None else set(raw.ch_names)
        drop = [c for c in raw.ch_names if c not in placed or c.upper() in NON_SCALP]
        if drop:
            raw.drop_channels(drop)
        d = raw.get_data()
        flat = [raw.ch_names[i] for i in range(len(raw.ch_names)) if float(np.std(d[i])) < 1e-9]
        if flat:
            raw.drop_channels(flat)
            d = raw.get_data()
        # Drop grossly deviant electrodes BEFORE epoching.  An epoch criterion evaluated over the whole
        # montage rejects an epoch because of its worst electrode, so one bad contact can reject a whole
        # recording -- which is exactly the silent failure nb-7-5 section 8 measures.  The criterion is
        # robust (median absolute deviation of the per-channel robust amplitude), not a fixed threshold,
        # because the cohorts here differ in amplifier and in units.
        amp = np.median(np.abs(d - np.median(d, axis=1, keepdims=True)), axis=1)
        med = float(np.median(amp))
        mad = float(np.median(np.abs(amp - med))) or float(np.std(amp)) or 1e-30
        z = (amp - med) / (1.4826 * mad)
        bad = [raw.ch_names[i] for i in range(len(raw.ch_names))
               if z[i] > 6.0 or amp[i] < med / 20.0]
        if bad and len(bad) < len(raw.ch_names) // 3:
            raw.drop_channels(bad)
        # One band for every dataset.  Without it a 50 Hz cohort and a 60 Hz cohort are being compared
        # partly on their mains, and LEMON's raw release (0.015-1000 Hz, no notch) carries drift that
        # would dominate a peak-to-peak criterion.
        raw.filter(HP_HZ, LP_HZ, picks="eeg", method="fir", fir_design="firwin", phase="zero",
                   verbose=False)
    return raw, {"dropped": drop, "flat": flat, "bad": bad, "n_kept": len(raw.ch_names), **scale}


def epoch_spectra(raw, *, t0=SEGMENT_T0, duration=SEGMENT_DUR):
    """Per-epoch Welch spectra (uV^2/Hz) of one recording, average-referenced over the kept channels.

    Returns (freqs, psd of shape (n_epochs, n_channels, n_freqs), channel names, notes).
    """
    r, notes = clean_scalp(raw)
    sf = float(r.info["sfreq"])
    t1 = min(r.times[-1], t0 + duration)
    if t1 - t0 < 3 * EPOCH_S:
        t0, t1 = 0.0, min(r.times[-1], duration)
    x = r.get_data(tmin=t0, tmax=t1) * 1e6
    x = x - x.mean(axis=0, keepdims=True)            # average reference over the retained channels
    n = int(round(EPOCH_S * sf))
    n_ep = x.shape[1] // n
    if n_ep < 4:
        raise RuntimeError(f"only {n_ep} epochs of {EPOCH_S:g} s available")
    X = x[:, :n_ep * n].reshape(x.shape[0], n_ep, n).transpose(1, 0, 2)
    keep = np.abs(X - X.mean(axis=-1, keepdims=True)).max(axis=(1, 2)) <= EPOCH_REJECT_UV
    if keep.sum() < 4:
        raise RuntimeError(f"only {int(keep.sum())} of {n_ep} epochs survive the {EPOCH_REJECT_UV:g} uV "
                           f"criterion")
    freqs, psd = sps.welch(X[keep], fs=sf, window="hann", nperseg=n, noverlap=0,
                           detrend="constant", scaling="density", axis=-1)
    notes |= {"sfreq": sf, "t0_s": float(t0), "t1_s": float(t1), "n_epochs": int(n_ep),
              "n_epochs_kept": int(keep.sum()), "resolution_hz": float(sf / n)}
    return freqs, psd, list(r.ch_names), notes


def relative_alpha(raw, *, band=(8.0, 13.0), total=(1.0, 45.0)):
    """Relative posterior alpha power of a short probe -- no epoching, no rejection.

    Used only to decide which of LEMON's two marker codes is the eyes-closed one, which the
    directory does not record (label_source: algorithmic, TODO(confirm), as nb-4-6 does it).
    """
    r, _ = clean_scalp(raw)
    post = [c for c in r.ch_names if c in POSTERIOR]
    x = r.get_data(picks=post or None) * 1e6
    sf = float(r.info["sfreq"])
    n = int(round(2.0 * sf))
    f, P = sps.welch(x, fs=sf, window="hann", nperseg=n, noverlap=n // 2, detrend="constant",
                     scaling="density")
    return band_sum(f, P.mean(0), band) / band_sum(f, P.mean(0), total)


def peak_above_trend(freqs, psd, *, band=ALPHA_BAND, fit_range=FIT_RANGE, min_db=MIN_DB):
    """helpers.alpha_peak's algorithm applied to a spectrum rather than a Raw."""
    freqs = np.asarray(freqs, float)
    psd = np.asarray(psd, float)
    if psd.ndim > 1:
        psd = psd.mean(axis=tuple(range(psd.ndim - 1)))
    m = (freqs >= fit_range[0]) & (freqs <= fit_range[1]) & (psd > 0)
    f, p = freqs[m], 10 * np.log10(psd[m])
    keep = (f < 5) | (f > 15)
    slope, icpt = np.polyfit(np.log10(f[keep]), p[keep], 1)
    resid = p - (slope * np.log10(f) + icpt)
    bm = (f >= band[0]) & (f <= band[1])
    i = int(np.argmax(np.where(bm, resid, -np.inf)))
    # A maximum that sits ON the search limit is the limit's choice, not the data's: the residual is
    # still rising (or falling) where the band stops.  Eyes-open recordings without a posterior rhythm
    # do exactly this, and calling the edge an "alpha peak" would manufacture a value for every one of
    # them.  It is reported as absent, and the count of absences is itself a result.
    idx_band = np.where(bm)[0]
    at_edge = bool(i == idx_band[0] or i == idx_band[-1])
    ok = (resid[i] >= min_db) and not at_edge
    return {"freq_hz": float(f[i]) if ok else np.nan,
            "db_above_trend": float(resid[i]), "at_band_edge": at_edge}


def features(freqs, psd, ch_names, *, epochs=None):
    """The three pre-specified features from one set of epoch spectra."""
    sel = slice(None) if epochs is None else epochs
    P = np.asarray(psd)[sel]                                  # (n_epochs, n_channels, n_freqs)
    post = [i for i, c in enumerate(ch_names) if c in POSTERIOR]
    mean_all = P.mean(axis=0)                                  # (n_channels, n_freqs)
    mean_post = mean_all[post].mean(axis=0) if post else mean_all.mean(axis=0)
    out = {"n_epochs": int(P.shape[0]), "n_posterior": len(post)}
    out |= {"iaf_hz": peak_above_trend(freqs, mean_post)["freq_hz"],
            "iaf_db": peak_above_trend(freqs, mean_post)["db_above_trend"],
            "theta_beta": band_sum(freqs, mean_all, THETA) / band_sum(freqs, mean_all, BETA)}
    try:
        fit = L1.fit_specparam(freqs, mean_post, freq_range=EXPONENT_RANGE)
        out["exponent"] = float(fit["exponent"])
    except Exception as exc:                                   # noqa: BLE001
        out["exponent"] = np.nan
        out["exponent_error"] = f"{type(exc).__name__}: {exc}"
    return out


print("feature extractor defined; the same one is used for every dataset and for both halves of a "
      "split-half.")
print(f"unit check: a recording whose robust AC amplitude exceeds {MAX_PLAUSIBLE_UV:g} uV in the units "
      f"MNE returns is rescaled by 1e-6 and the correction is reported per recording.")

import yaml

_DIR_YAML = next((d / "data" / "directory.yaml" for d in (Path.cwd(), *Path.cwd().parents)
                  if (d / "data" / "directory.yaml").exists()), None)
DIRECTORY = {e["id"]: e for e in yaml.safe_load(_DIR_YAML.read_text())["datasets"]}

USED = ["ds-lemon", "ds-dortmund", "ds-srm", "ds-iowapd"]
DISCUSSED = ["ds-tdbrain", "ds-brainlat", "ds-chbmp", "ds-hbn", "ds-aszed"]
for group, ids in (("COMPUTED FROM", USED), ("NAMED, NOT LOADED", DISCUSSED)):
    print(f"========== {group} ==========")
    for ds in ids:
        e = DIRECTORY[ds]
        lic = e.get("license", {})
        print(f"\n{ds} -- {e['name']}")
        print(f"   licence {lic.get('name', 'TODO(confirm)')}, access {e.get('access')}, "
              f"sessions {e.get('sessions')}, population {e.get('population')}")
        print(f"   {e.get('channels', e.get('channels_note', '?'))} ch, "
              f"{e.get('sfreq_hz', e.get('sfreq_note', '?'))} Hz, mains {e.get('mains_hz')}, "
              f"reference {e.get('reference')}, online filters {e.get('online_filters')}")
        for c in e.get("caveats", []):
            print(f"      - {c}")
print()
L5.print_licences(*USED, *DISCUSSED)

print("Why the four in the second group are not loaded, and what each would have cost the number:\n")
BLOCKERS = {
    "ds-tdbrain": ("licence/access -- Data Use Agreement forbids redistribution (spec 10.7 class D)",
                   "47 healthy controls among 1,274 patients, comorbidity is common, and TWO CAP SYSTEMS "
                   "OVER 20 YEARS.  A test-retest across that span measures the caps as much as the "
                   "people (pf-site-device-confound)."),
    "ds-brainlat": ("licence/access -- Synapse registration and data-use terms the site cannot pass on "
                    "(class D)",
                    "MIXED 50 AND 60 Hz MAINS WITHIN ONE DATASET, by country.  Any measure whose "
                    "normalising band reaches the line frequency is partly a measure of geography."),
    "ds-chbmp": ("licence -- CC BY-NC-SA plus registration (class C and D)",
                 "channel counts vary per subject (58-121) and event markers are in Spanish; a "
                 "montage-dependent feature is not the same feature across participants."),
    "ds-hbn": ("licence decision -- CC BY-SA 4.0 with no license_decision recorded, so derive_snippets "
               "returns 'no' (class B).  A notebook MAY download it; no asset may derive from it, and "
               "this notebook derives nothing, so it downloads nothing either.",
               "the sample is TRANSDIAGNOSTIC despite the name 'Healthy Brain Network', channels are "
               "E1-E128 + Cz rather than 10-20 names, and each release is a separate accession.  A "
               "'healthy control' group drawn from it is not one."),
    "ds-aszed": ("nothing -- CC BY 4.0, open, and Level 4 computes from it",
                 "TWO AMPLIFIERS AT 200 AND 256 Hz AT TWO SITES.  It is in this list as the case where "
                 "the confound is in an OPEN dataset: the licence is not what makes an archive hard."),
}
for ds, (why, what) in BLOCKERS.items():
    print(f"{ds}")
    print(f"   why not loaded : {why}")
    print(f"   what it would  : {what}")
    print()
print("The pattern worth naming: THREE OF THE FIVE are hard for a reason that has nothing to do with the")
print("licence -- a cohort's composition, its hardware history or its montage.  A biomarker result that")
print("survives peer review can still fail to replicate because the second cohort was recorded on the")
print("other amplifier.  That is L7.6's replication problem, and it is a study-design problem.")

dl = L4.Downloads("nb-7-6").start()
_T_START = time.time()


def register_new(paths, downloads):
    """Register only files this run created -- a cache another notebook filled is not ours to delete."""
    new = [Path(p) for p in paths if p is not None and Path(p).is_file()
           and Path(p).stat().st_mtime >= _T_START]
    downloads.add(*new)
    return new


def split_half(freqs, psd, ch_names):
    """Features of the odd and the even epochs of one recording."""
    n = psd.shape[0]
    a = features(freqs, psd, ch_names, epochs=np.arange(0, n, 2))
    b = features(freqs, psd, ch_names, epochs=np.arange(1, n, 2))
    return a, b


LEMON_SUBJECTS = [f"sub-{10002 + i:06d}" for i in range(3 * N_LEMON + 6)]
PROBE_S = 12.0
lemon_rows = []
for subj in LEMON_SUBJECTS:
    if len(lemon_rows) >= N_LEMON:
        break
    tmp = L1.download_dir() / "l7b-lemon"
    try:
        hdr = L4.lemon_headers(subj)
        blocks = L4.lemon_blocks(hdr)
        # Which marker code is eyes-closed is NOT recorded in data/directory.yaml, so it is decided
        # from the signal, exactly as nb-4-6 does: probe one block of each code and take the one with
        # the higher relative posterior alpha.  label_source: algorithmic, TODO(confirm).
        probe = {}
        for code in ("S200", "S210"):
            b = next((b for b in blocks if b["code"] == code), None)
            if b is None:
                continue
            p = L4.fetch_lemon_window(subj, b["t0_s"], PROBE_S, tmp, dl, headers=hdr, verbose=False)
            if p is None:
                raise RuntimeError("too little free disk for the probe")
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                raw = mne.io.read_raw_brainvision(p, preload=True, verbose=False)
            probe[code] = {"rel_alpha": relative_alpha(raw), "block": b}
        ec_code = max(probe, key=lambda c: probe[c]["rel_alpha"])
        b = probe[ec_code]["block"]
        block_s = float(min(60.0, b["duration_s"]))
        vhdr = L4.fetch_lemon_window(subj, b["t0_s"], block_s, tmp, dl, headers=hdr, verbose=False)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            raw = mne.io.read_raw_brainvision(vhdr, preload=True, verbose=False)
        f, P, chs, notes = epoch_spectra(raw, t0=0.0, duration=block_s)
        whole = features(f, P, chs)
        a, bfeat = split_half(f, P, chs)
        lemon_rows.append({"subject": subj, "ec_code": ec_code,
                           "rel_alpha_ec": round(probe[ec_code]["rel_alpha"], 4),
                           "rel_alpha_other": round(min(v["rel_alpha"] for v in probe.values()), 4),
                           "n_ch": notes["n_kept"], "n_epochs": notes["n_epochs_kept"],
                           **{f"{k}_whole": whole[k] for k in ("iaf_hz", "theta_beta", "exponent")},
                           **{f"{k}_odd": a[k] for k in ("iaf_hz", "theta_beta", "exponent")},
                           **{f"{k}_even": bfeat[k] for k in ("iaf_hz", "theta_beta", "exponent")}})
        print(f"  {subj}: EC = {ec_code} (rel alpha {probe[ec_code]['rel_alpha']:.3f} vs "
              f"{min(v['rel_alpha'] for v in probe.values()):.3f}), "
              f"IAF {whole['iaf_hz']:.2f} Hz, theta/beta {whole['theta_beta']:.3f}", flush=True)
    except Exception as exc:                                   # noqa: BLE001
        print(f"  {subj}: SKIPPED -- {type(exc).__name__}: {exc}", flush=True)
    finally:
        dl.cleanup()
lemon = pd.DataFrame(lemon_rows)
print(f"\nds-lemon: {len(lemon)} subjects measured in {time.time() - _T_START:.0f} s")
L6.disk_report("after ds-lemon", folders={"course downloads": L1.download_dir()})

def spearman_brown(r):
    return 2 * r / (1 + r) if np.isfinite(r) else np.nan


def icc_two_way(x, y):
    """ICC(2,1) absolute agreement and ICC(3,1) consistency for two measurements of n subjects.

    Two-way random/mixed ANOVA with k = 2 raters (here: sessions).  Written out rather than taken from
    a package so that the notebook has no unpinned dependency; cross-checked against pingouin below
    when it is installed.
    """
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    m = np.isfinite(x) & np.isfinite(y)
    x, y = x[m], y[m]
    n, k = len(x), 2
    if n < 3:
        return {"n": n, "icc21": np.nan, "icc31": np.nan, "pearson_r": np.nan, "mean_shift": np.nan}
    M = np.column_stack([x, y])
    grand = M.mean()
    row = M.mean(axis=1)
    col = M.mean(axis=0)
    ms_r = k * ((row - grand) ** 2).sum() / (n - 1)                       # between subjects
    ms_c = n * ((col - grand) ** 2).sum() / (k - 1)                       # between sessions
    ss_e = ((M - row[:, None] - col[None, :] + grand) ** 2).sum()
    ms_e = ss_e / ((n - 1) * (k - 1))                                     # residual
    den21 = ms_r + (k - 1) * ms_e + k * (ms_c - ms_e) / n
    den31 = ms_r + (k - 1) * ms_e
    # A cohort with no between-subject variance and no measurement error makes both denominators
    # zero; that is a degenerate sample, not a reliability of 1, and it is reported as undefined.
    icc21 = (ms_r - ms_e) / den21 if abs(den21) > 1e-300 else np.nan
    icc31 = (ms_r - ms_e) / den31 if abs(den31) > 1e-300 else np.nan
    const = float(np.std(x)) < 1e-12 or float(np.std(y)) < 1e-12
    r = np.nan if const else float(stats.pearsonr(x, y)[0])
    return {"n": n, "icc21": float(icc21), "icc31": float(icc31), "pearson_r": r,
            "mean_shift": float(col[1] - col[0]),
            "sd_within": float(np.sqrt(ms_e)), "sd_between": float(np.sqrt(max(ms_r - ms_e, 0) / k))}


def icc_label(v):
    if not np.isfinite(v):
        return "n/a"
    return ("poor" if v < 0.5 else "moderate" if v < 0.75 else "good" if v < 0.9 else "excellent")


# Cross-check against pingouin when it happens to be installed.  It is NOT in
# notebooks/requirements.txt, so nothing here depends on it and nothing is installed for it.
HAVE_PINGOUIN = importlib.util.find_spec("pingouin") is not None
print(f"pingouin available for the cross-check: {HAVE_PINGOUIN}")

rows = []
for feat in ("iaf_hz", "theta_beta", "exponent"):
    r = icc_two_way(lemon[f"{feat}_odd"], lemon[f"{feat}_even"])
    rows.append({"feature": feat, "n": r["n"], "split-half r": round(r["pearson_r"], 4),
                 "Spearman-Brown": round(spearman_brown(r["pearson_r"]), 4),
                 "ICC(2,1) halves": round(r["icc21"], 4),
                 "verdict": icc_label(spearman_brown(r["pearson_r"]))})
lemon_rel = pd.DataFrame(rows)
print(f"\nSplit-half reliability, ds-lemon, one 60-s eyes-closed block per subject "
      f"({int(lemon.n_epochs.median())} epochs median, odd vs even):")
print(lemon_rel.to_string(index=False))
print()
print("THE SENTENCE THAT HAS TO BE SAID: ds-lemon has ONE EEG session per participant "
      f"(data/directory.yaml: sessions = {DIRECTORY['ds-lemon'].get('sessions')}).  Everything above is")
print("agreement between two halves of the SAME MINUTE of recording.  It is an upper bound on what a")
print("second session could give and it is not test-retest reliability.  Reporting it as though it were")
print("is the single most common way a resting-state feature is made to look more stable than it is.")
if HAVE_PINGOUIN:
    try:
        import pingouin as pg

        long = pd.DataFrame({"subject": np.r_[lemon.subject, lemon.subject],
                             "half": ["odd"] * len(lemon) + ["even"] * len(lemon),
                             "iaf": np.r_[lemon.iaf_hz_odd, lemon.iaf_hz_even]}).dropna()
        got = pg.intraclass_corr(data=long, targets="subject", raters="half", ratings="iaf")
        # pingouin 0.6 labels the two-way absolute-agreement single-measure ICC "ICC(A,1)"; older
        # versions call it "ICC2".  Accept either rather than pinning a label.
        types = got.Type.astype(str)
        row = got[types.isin(["ICC(A,1)", "ICC2"])]
        mine = icc_two_way(lemon.iaf_hz_odd, lemon.iaf_hz_even)
        ref = float(row["ICC"].iloc[0])
        print(f"\ncross-check vs pingouin {row['Type'].iloc[0]} on IAF: {mine['icc21']:.6f} "
              f"(this notebook) vs {ref:.6f} (pingouin), difference {mine['icc21'] - ref:+.2e}")
    except Exception as exc:                                    # noqa: BLE001
        # pingouin is NOT in notebooks/requirements.txt; nothing depends on it and an API change
        # in an unpinned package must not fail a notebook.  The reason is printed, not swallowed.
        print(f"\ncross-check with pingouin skipped: {type(exc).__name__}: {exc}")

def resting_features(dataset, subject, *, session="1", task="EyesClosed", acq="pre"):
    """Fetch one resting recording, measure the three features, delete the download in a finally."""
    files = {}
    try:
        if dataset == "ds-dortmund":
            raw, files = L1.load_dortmund(subject, session=session, task=task, acq=acq,
                                          return_paths=True, verbose=False)
        elif dataset == "ds-iowapd":
            raw, files = L1.load_iowapd(subject, return_paths=True, verbose=False)
        elif dataset == "ds-srm":
            base = f"ds003775/{subject}/ses-{session}/eeg/{subject}_ses-{session}_task-resteyesc_"
            for suffix in ("eeg.edf", "eeg.json", "channels.tsv"):
                files[suffix] = L1.fetch(f"{L1.OPENNEURO_S3}/{base}{suffix}", base + suffix,
                                         verbose=False)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                raw = mne.io.read_raw_edf(files["eeg.edf"], preload=True, verbose=False)
        else:
            raise ValueError(dataset)
        f, P, chs, notes = epoch_spectra(raw)
        out = features(f, P, chs)
        out |= {"subject": subject, "dataset": dataset, "session": session, "acq": acq,
                "n_ch": notes["n_kept"], "n_epochs_kept": notes["n_epochs_kept"],
                "duration_s": round(float(raw.times[-1]), 1), "sfreq": notes["sfreq"],
                "scale_factor": notes["scale_factor"], "ac_uv": round(notes["ac_uv_after"], 2)}
        return out
    finally:
        heavy = [p for p in files.values() if Path(p).suffix in (".edf", ".set", ".fdt", ".eeg")]
        if heavy:
            L1.cleanup(heavy, verbose=False)


def cohort(dataset, subjects, conditions, *, label=""):
    """resting_features over subjects x conditions; one recording is on disk at a time."""
    rows = []
    for s in subjects:
        got = {}
        for name, kw in conditions.items():
            try:
                got[name] = resting_features(dataset, s, **kw)
            except Exception as exc:                            # noqa: BLE001
                got[name] = {"error": f"{type(exc).__name__}: {exc}"}
        row = {"subject": s}
        ok = True
        for name, r in got.items():
            if "error" in r:
                ok = False
                row[f"{name}_error"] = r["error"]
            else:
                for k in ("iaf_hz", "theta_beta", "exponent", "n_ch", "n_epochs_kept",
                          "duration_s", "scale_factor", "ac_uv"):
                    row[f"{k}_{name}"] = r[k]
        row["ok"] = ok
        rows.append(row)
        msg = " | ".join(f"{n}: IAF {row.get(f'iaf_hz_{n}', float('nan')):.2f} Hz"
                         for n in conditions) if ok else row.get(f"{list(conditions)[0]}_error", "failed")
        print(f"  {label or dataset} {s}: {msg}  (free {helpers.free_disk_mb('.') / 1000:.2f} GB)",
              flush=True)
    return pd.DataFrame(rows)


dparts = L6.dortmund_participants()
print(f"ds-dortmund participants.tsv: {len(dparts)} rows, columns {list(dparts.columns)}")
have2 = dparts[(dparts.session1 == "yes") & (dparts.session2 == "yes")] \
    if "session2" in dparts.columns else dparts[dparts.session1 == "yes"]
dort_ids = have2.participant_id.tolist()[:N_DORT]
print(f"selection rule, stated before any recording is read: the first {N_DORT} participants with BOTH "
      f"sessions recorded, in participants.tsv order ({len(have2)} qualify).")
print(f"   {', '.join(dort_ids)}")
t0 = time.time()
dort_within = cohort("ds-dortmund", dort_ids,
                     {"pre": dict(session="1", task="EyesClosed", acq="pre"),
                      "post": dict(session="1", task="EyesClosed", acq="post")},
                     label="ds-dortmund ses-1")
print(f"\n{int(dort_within.ok.sum())} of {len(dort_within)} subjects in {time.time() - t0:.0f} s")
L6.disk_report("after ds-dortmund ses-1", folders={"course downloads": L1.download_dir()})

srm_parts = pd.read_csv(L1.fetch(f"{L1.OPENNEURO_S3}/ds003775/participants.tsv",
                                 "ds003775/participants.tsv", verbose=False), sep="\t")
print(f"ds-srm participants.tsv: {len(srm_parts)} rows; the release documents a second session for ~42.")
_ch = pd.read_csv(L1.fetch(f"{L1.OPENNEURO_S3}/ds003775/sub-001/ses-t1/eeg/"
                           f"sub-001_ses-t1_task-resteyesc_channels.tsv",
                           "ds003775/sub-001/ses-t1/eeg/sub-001_ses-t1_task-resteyesc_channels.tsv",
                           verbose=False), sep="\t")
print(f"   its channels.tsv declares units {sorted(set(_ch['units'].dropna()))} for "
      f"{len(_ch)} channels and types {sorted(set(_ch['type'].dropna()))}; the EDF signal headers "
      f"declare no unit at all, which is why ensure_microvolts has to test the data.")
srm_ids = []
for sid in srm_parts.participant_id:
    if len(srm_ids) >= N_SRM:
        break
    base = f"ds003775/{sid}/ses-t2/eeg/{sid}_ses-t2_task-resteyesc_eeg.edf"
    if helpers.http_head(f"{L1.OPENNEURO_S3}/{base}")["ok"]:
        srm_ids.append(sid)
print(f"selection rule, stated before any recording is read: the first {N_SRM} participants in "
      f"participants.tsv order whose ses-t2 file exists (checked with a HEAD request, no download).")
print(f"   {', '.join(srm_ids)}")
t0 = time.time()
srm = cohort("ds-srm", srm_ids, {"t1": dict(session="t1"), "t2": dict(session="t2")}, label="ds-srm")
print(f"\n{int(srm.ok.sum())} of {len(srm)} subjects in {time.time() - t0:.0f} s")
L6.disk_report("after ds-srm", folders={"course downloads": L1.download_dir()})

t0 = time.time()
dort_5y = cohort("ds-dortmund", dort_ids,
                 {"ses1": dict(session="1", task="EyesClosed", acq="pre"),
                  "ses2": dict(session="2", task="EyesClosed", acq="pre")},
                 label="ds-dortmund 5 y")
print(f"\n{int(dort_5y.ok.sum())} of {len(dort_5y)} subjects in {time.time() - t0:.0f} s")
L6.disk_report("after ds-dortmund ses-2", folders={"course downloads": L1.download_dir()})

RETESTS = [
    ("seconds  (split-half, same minute)", "ds-lemon", lemon, "odd", "even"),
    ("minutes  (pre vs post, one session)", "ds-dortmund", dort_within[dort_within.ok], "pre", "post"),
    ("2-3 months", "ds-srm", srm[srm.ok], "t1", "t2"),
    ("5 years", "ds-dortmund", dort_5y[dort_5y.ok], "ses1", "ses2"),
]
def icc_ci(x, y, *, n_boot=2000, seed=SEED):
    """Percentile bootstrap interval for ICC(2,1) over subjects.  Small n is the point, not a defect."""
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    m = np.isfinite(x) & np.isfinite(y)
    x, y = x[m], y[m]
    if len(x) < 3:
        return (np.nan, np.nan)
    g = np.random.default_rng(seed)
    vals = []
    warnings.simplefilter("ignore")   # a degenerate resample is handled above, not by a warning whose
                                      # traceback line would put a kernel temp path in the output
    for _ in range(n_boot):
        i = g.integers(0, len(x), len(x))
        if len(np.unique(i)) < 2:
            continue
        v = icc_two_way(x[i], y[i])["icc21"]
        if np.isfinite(v):
            vals.append(v)
    return (float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))) if vals else (np.nan, np.nan)


UNITS = {"iaf_hz": "Hz", "theta_beta": "ratio", "exponent": "exponent"}
rows = []
for label, ds, tab, a, b in RETESTS:
    for feat in ("iaf_hz", "theta_beta", "exponent"):
        r = icc_two_way(tab[f"{feat}_{a}"], tab[f"{feat}_{b}"])
        lo, hi = icc_ci(tab[f"{feat}_{a}"], tab[f"{feat}_{b}"])
        rows.append({"interval": label, "dataset": ds, "feature": feat, "n": r["n"],
                     "ICC(2,1)": round(r["icc21"], 3), "95% CI": f"[{lo:+.2f}, {hi:+.2f}]",
                     "ICC(3,1)": round(r["icc31"], 3),
                     "r": round(r["pearson_r"], 3), "mean shift": round(r["mean_shift"], 3),
                     "retest SD": round(r["sd_within"], 3), "units": UNITS[feat],
                     "between-subject SD": round(r["sd_between"], 3),
                     "verdict": icc_label(r["icc21"])})
rel = pd.DataFrame(rows)
print("Reliability as a function of the interval between the two measurements.")
print("ICC(2,1) is absolute agreement (a systematic shift counts against it); ICC(3,1) is consistency")
print("(a systematic shift does not).  Where they differ, the difference IS the shift.")
print("The RETEST SD column is the same information without the ratio: the typical difference between a")
print("person's two measurements, in the feature's own units.  It is there because an ICC is a RATIO of")
print("between-subject to total variance, so a cohort that happens to be homogeneous produces a low ICC")
print("for a measurement that is in fact precise.  Read the two together or neither.\n")
for feat in ("iaf_hz", "theta_beta", "exponent"):
    print(f"--- {feat} ---")
    print(rel[rel.feature == feat].drop(columns="feature").to_string(index=False))
    print()

fig, axes = plt.subplots(1, 3, figsize=(14.5, 4.3))
order = [r[0] for r in RETESTS]
for ax, feat, unit in zip(axes, ("iaf_hz", "theta_beta", "exponent"),
                          ("Hz", "dimensionless ratio", "dimensionless exponent")):
    sub = rel[rel.feature == feat].set_index("interval").loc[order]
    x = np.arange(len(sub))
    ax.bar(x - 0.2, sub["ICC(2,1)"], 0.4, label="ICC(2,1) absolute agreement")
    ax.bar(x + 0.2, sub["ICC(3,1)"], 0.4, label="ICC(3,1) consistency")
    for thr, name in ((0.5, "poor / moderate"), (0.75, "moderate / good"), (0.9, "good / excellent")):
        ax.axhline(thr, color="0.7", lw=0.7, ls=":")
        ax.text(len(sub) - 0.45, thr + 0.012, name, fontsize=6.4, color="0.45", ha="right")
    ax.set_xticks(x, [f"{s}\n(n = {int(n)})" for s, n in zip(sub.index, sub["n"])],
                  rotation=18, ha="right", fontsize=7)
    ax.set(ylim=(-0.2, 1.05), ylabel="Intraclass correlation (dimensionless)",
           title=f"{feat}  ({unit})")
    ax.axhline(0, color="k", lw=0.8)
    ax.grid(alpha=0.3, axis="y")
axes[0].legend(fontsize=7.4, loc="lower left")
fig.suptitle("Test–retest reliability against the interval between measurements (ICC, dimensionless)",
             y=1.02, fontsize=11)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

parts = L6.iowapd_participants()
print(f"ds-iowapd participants.tsv: {len(parts)} rows, columns {list(parts.columns)}")
print(parts.GROUP.value_counts().to_string())
for col in ("AGE", "MOCA", "UPDRS"):
    if col in parts.columns:
        print(f"\n{col} by group:")
        print(parts.groupby("GROUP")[col].agg(["count", "mean", "std"]).round(2).to_string())
pd_ids = parts.loc[parts.GROUP == "PD", "participant_id"].tolist()[:N_PD]
hc_ids = parts.loc[parts.GROUP == "Control", "participant_id"].tolist()[:N_PD]
print(f"\nselection rule, stated before any recording is read: the first {N_PD} PD and the first {N_PD} "
      f"Control participants in participants.tsv order.  No recording is inspected first.")
t0 = time.time()
iowa = cohort("ds-iowapd", pd_ids + hc_ids, {"rest": {}}, label="ds-iowapd")
iowa["group"] = ["PD"] * len(pd_ids) + ["Control"] * len(hc_ids)
iowa = iowa.merge(parts[["participant_id", "AGE", "GENDER", "MOCA"]],
                  left_on="subject", right_on="participant_id", how="left")
print(f"\n{int(iowa.ok.sum())} of {len(iowa)} recordings in {time.time() - t0:.0f} s")
L6.disk_report("after ds-iowapd", folders={"course downloads": L1.download_dir()})

ok = iowa[iowa.ok]
rows = []
for feat in ("iaf_hz", "theta_beta", "exponent"):
    a = ok.loc[ok.group == "PD", f"{feat}_rest"].to_numpy(float)
    b = ok.loc[ok.group == "Control", f"{feat}_rest"].to_numpy(float)
    a, b = a[np.isfinite(a)], b[np.isfinite(b)]
    if min(len(a), len(b)) < 3:
        rows.append({"feature": feat, "n PD": len(a), "n HC": len(b), "PD mean": np.nan,
                     "HC mean": np.nan, "Hedges g": np.nan, "95% CI": "n/a", "Welch p": np.nan})
        continue
    g = L6.hedges_g(a, b)
    t, p = stats.ttest_ind(a, b, equal_var=False)
    rows.append({"feature": feat, "n PD": len(a), "n HC": len(b),
                 "PD mean": round(a.mean(), 3), "HC mean": round(b.mean(), 3),
                 "Hedges g": round(g["g"], 3),
                 "95% CI": f"[{g['ci'][0]:+.2f}, {g['ci'][1]:+.2f}]",
                 "Welch p": round(float(p), 4)})
groupcon = pd.DataFrame(rows)
print(f"ds-iowapd group contrast, eyes-open rest, {int((ok.group == 'PD').sum())} PD (on medication) vs "
      f"{int((ok.group == 'Control').sum())} Control:")
print(groupcon.to_string(index=False))
print()
n_nan = int((~np.isfinite(ok[f'iaf_hz_rest'].to_numpy(float))).sum())
print(f"IAF is undefined for {n_nan} of {len(ok)} eyes-open recordings -- either no residual clears "
      f"{MIN_DB:g} dB,")
print(f"or the largest one sits ON the 7 or 13 Hz search limit, which means the band chose it and not")
print("the data.  That is not a failure of the code: eyes-open rest was chosen by these authors")
print("BECAUSE it suppresses the posterior dominant rhythm.  A feature that cannot be measured in the")
print("condition the cohort was recorded in has an effect size of nothing, whatever its reliability.")
print()
print("A confound check, because the catalog names it: is the analysed sample age-imbalanced?")
age_pd = ok.loc[ok.group == "PD", "AGE"].to_numpy(float)
age_hc = ok.loc[ok.group == "Control", "AGE"].to_numpy(float)
t_age, p_age = stats.ttest_ind(age_pd, age_hc, equal_var=False, nan_policy="omit")
print(f"   PD {np.nanmean(age_pd):.1f} +- {np.nanstd(age_pd, ddof=1):.1f} y, "
      f"Control {np.nanmean(age_hc):.1f} +- {np.nanstd(age_hc, ddof=1):.1f} y; "
      f"Welch t = {t_age:.2f}, p = {p_age:.3f}")
print(f"   sex: {ok.groupby(['group', 'GENDER']).size().to_dict()}")

best = {}
for feat in ("iaf_hz", "theta_beta", "exponent"):
    r5 = rel[(rel.feature == feat) & (rel.interval == "5 years")]["ICC(2,1)"].iloc[0]
    rm = rel[(rel.feature == feat) & (rel.interval == "2-3 months")]["ICC(2,1)"].iloc[0]
    rs = rel[(rel.feature == feat) & (rel.interval.str.startswith("seconds"))]["ICC(2,1)"].iloc[0]
    g = groupcon[groupcon.feature == feat]["Hedges g"]
    g = float(g.iloc[0]) if len(g) and np.isfinite(g.iloc[0]) else np.nan
    best[feat] = {"ICC seconds": rs, "ICC 2-3 months": rm, "ICC 5 years": r5,
                  "observed g (ds-iowapd)": g,
                  "ceiling on |r| with any outcome (sqrt ICC_3mo)":
                      round(np.sqrt(max(rm, 0)), 3) if np.isfinite(rm) else np.nan,
                  "true g implied if attenuation is the only cause":
                      round(g / np.sqrt(max(rm, 1e-9)), 3) if np.isfinite(g) and rm > 0 else np.nan}
fit = pd.DataFrame(best).T
print("Reliability and effect, side by side:\n")
print(fit.to_string())
print()
for feat in fit.index:
    rm = fit.loc[feat, "ICC 2-3 months"]
    r5 = fit.loc[feat, "ICC 5 years"]
    print(f"{feat}:")
    print(f"   at 2-3 months  ICC = {rm:+.3f} ({icc_label(rm)});  at 5 years ICC = {r5:+.3f} "
          f"({icc_label(r5)})")
    print(f"   a correlation with any external outcome cannot exceed "
          f"{np.sqrt(max(rm, 0)):.2f} in expectation, however large the true association")
    print()