# 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", "pooch>=1.8"]
    subprocess.check_call(_cmd)

# 2. Shared helpers, 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_l5.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L5/ (or notebooks/) so that _shared/helpers_l5.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l5 as L5

# 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 each figure cell ends with plt.show().
import matplotlib.pyplot as plt
import numpy as np
import mne
import pooch

mne.set_log_level("WARNING")
pooch.get_logger().setLevel("WARNING")   # no download chatter: it would print local paths
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers_l5 imported from notebooks/_shared")
print("mne-connectivity available:", L5.have_mne_connectivity())
print(L5.disk_line("disk at the start"))

print("Licences, from data/directory.yaml (never from memory):")
L5.print_licences("ds-lemon", notes=True)
print()
L5.print_source_model_block()

import warnings
from mne.minimum_norm import make_inverse_operator
from mne.minimum_norm.inverse import _assemble_kernel, prepare_inverse_operator

FULL_COHORT = False
SUBSET = ("sub-010002", "sub-010003", "sub-010004", "sub-010005")
FULL_LIST = ("sub-010002", "sub-010003", "sub-010004", "sub-010005", "sub-010006", "sub-010007",
             "sub-010008", "sub-010009", "sub-010010", "sub-010011", "sub-010012", "sub-010013",
             "sub-010014", "sub-010015", "sub-010016", "sub-010017", "sub-010018", "sub-010019",
             "sub-010020", "sub-010021")
SUBJECTS = FULL_LIST if FULL_COHORT else SUBSET

MAX_MINUTES = 3.2
BAND = L5.ALPHA_BAND
CSD_PARAMS = dict(lambda2=1e-5, stiffness=4, n_legendre_terms=50)
N_REGIONS, MAX_DEPTH_MM = 12, 35.0
NOISE_UV, SNR = 1.0, 3.0
LAMBDA2 = 1.0 / SNR ** 2
N_SURROGATE = 100

ANALYSES = ("A sensor coherence", "B sensor wPLI", "C CSD then wPLI", "D source regions, wPLI + orth.")

print(f"FULL_COHORT = {FULL_COHORT}; {len(SUBJECTS)} subjects: {list(SUBJECTS)}")
print(f"  the full list is {len(FULL_LIST)} subjects and has NOT been executed for these outputs")
print(f"download per subject: the first {MAX_MINUTES:g} min of the raw BrainVision file, about 60 MB, "
      f"deleted before the next subject")
print()
for k, v in L5.LEMON_PIPELINE.items():
    print(f"  {k:16s}: {v}")
print(f"  {'band_hz':16s}: {BAND}")
print(f"  {'csd':16s}: mne.preprocessing.compute_current_source_density(sphere='auto', "
      + ", ".join(f"{a}={b}" for a, b in CSD_PARAMS.items()) + ")")
print(f"  {'source model':16s}: four-shell concentric sphere, minimum-norm inverse, lambda^2 = {LAMBDA2:.4f}, "
      f"{N_REGIONS} GEOMETRIC regions")
print(f"  {'surrogates':16s}: {N_SURROGATE} epoch shuffles per analysis per subject, max statistic over "
      f"pairs, 95th percentile")
print()
print(L5.disk_line("disk before any download"))

bem = L5.sphere_models()["four-shell"]
rr_all = L5.source_grid(spacing_m=0.012, max_fraction=0.85, min_radius_m=0.012)
depth_all = (L5.HEAD_RADIUS_M - np.linalg.norm(rr_all, axis=1)) * 1000
rr = rr_all[depth_all <= MAX_DEPTH_MM]
nn = rr / np.linalg.norm(rr, axis=1, keepdims=True)

rng = np.random.default_rng(L5.SEED)
centres = rr[rng.choice(len(rr), N_REGIONS, replace=False)]
for _ in range(200):
    lab = np.argmin(((rr[:, None, :] - centres[None]) ** 2).sum(-1), axis=1)
    new = np.stack([rr[lab == k].mean(axis=0) if (lab == k).any() else centres[k] for k in range(N_REGIONS)])
    if np.allclose(new, centres):
        break
    centres = new
labels = np.argmin(((rr[:, None, :] - centres[None]) ** 2).sum(-1), axis=1)
REGIONS = [f"R{k + 1:02d}" for k in range(N_REGIONS)]
print(f"{len(rr)} sources within {MAX_DEPTH_MM:.0f} mm of the scalp, {N_REGIONS} geometric regions "
      f"(k-means on position, seed {L5.SEED})")
print("  region labels are coordinates, not anatomy; no structure is named anywhere in this capstone")


def source_region_kernel(info_l):
    """Sensor-to-region weights: one row per geometric region, from a minimum-norm inverse."""
    fwd_l, G_l = L5.leadfield(info_l, bem, rr, fixed_normals=nn)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        cov = mne.make_ad_hoc_cov(info_l, std=NOISE_UV * 1e-6, verbose=False)
        inv = make_inverse_operator(info_l, fwd_l, cov, loose=0.0, depth=None, fixed=True, verbose=False)
        prep = prepare_inverse_operator(inv, nave=1, lambda2=LAMBDA2, method="MNE", verbose=False)
        K, _, _, _ = _assemble_kernel(prep, None, "MNE", pick_ori=None, verbose=False)
    return K


def region_rows(K, data_cov):
    W = np.zeros((N_REGIONS, K.shape[1]))
    for k in range(N_REGIONS):
        Kk = K[labels == k]
        w, V = np.linalg.eigh(Kk @ data_cov @ Kk.T)
        row = V[:, -1] @ Kk
        W[k] = row * np.sign(row[np.argmax(np.abs(row))])
    return W


def surrogate_threshold(tc, method, n=N_SURROGATE, rng=None):
    rng = rng or np.random.default_rng(L5.SEED)
    out = np.empty(n)
    for k in range(n):
        sh = np.stack([tc[rng.permutation(tc.shape[0]), c, :] for c in range(tc.shape[1])], axis=1)
        c_k, _ = L5.connectivity(sh, sfreq=SFREQ_ANALYSIS, methods=(method,), fmin=BAND[0], fmax=BAND[1],
                                 backend="numpy")
        out[k] = np.abs(L5.upper_pairs(np.abs(c_k[method]))).max()
    return float(np.percentile(out, 95))


SFREQ_ANALYSIS = L5.LEMON_PIPELINE["resample_hz"]
print(f"analyses: {ANALYSES}")

results, per_subject = {}, []
K_cache = {}
for subject in SUBJECTS:
    paths = []
    try:
        epochs, meta = L5.lemon_condition_epochs(subject, "eyes-closed", max_minutes=MAX_MINUTES,
                                                 verbose=False)
        paths = L5.lemon_files(subject)
        if epochs is None:
            print(f"  {subject}: SKIPPED -- {meta.get('error')}")
            continue
        sensor = epochs.get_data(copy=True)
        csd = mne.preprocessing.compute_current_source_density(epochs.copy(), **CSD_PARAMS,
                                                              verbose=False).get_data(copy=True)
        ch = tuple(epochs.ch_names)
        info_l = epochs.info.copy()
    finally:
        L5.delete_files(paths, verbose=False)

    if ch not in K_cache:
        K_cache[ch] = source_region_kernel(info_l)
    W = region_rows(K_cache[ch], np.mean([np.cov(x) for x in sensor], axis=0))
    tc = np.einsum("rc,ect->ert", W, sensor)
    flat = tc.transpose(1, 0, 2).reshape(N_REGIONS, -1)
    tc_orth = L5.symmetric_orthogonalise(flat).reshape(N_REGIONS, tc.shape[0],
                                                       tc.shape[2]).transpose(1, 0, 2)

    con_sensor, backend = L5.connectivity(sensor, sfreq=SFREQ_ANALYSIS, methods=("coh", "wpli"),
                                          fmin=BAND[0], fmax=BAND[1])
    con_csd, _ = L5.connectivity(csd, sfreq=SFREQ_ANALYSIS, methods=("wpli",), fmin=BAND[0], fmax=BAND[1])
    con_src, _ = L5.connectivity(tc_orth, sfreq=SFREQ_ANALYSIS, methods=("wpli",), fmin=BAND[0], fmax=BAND[1])
    M = {ANALYSES[0]: np.abs(con_sensor["coh"]), ANALYSES[1]: np.abs(con_sensor["wpli"]),
         ANALYSES[2]: np.abs(con_csd["wpli"]), ANALYSES[3]: np.abs(con_src["wpli"])}
    thr = {ANALYSES[0]: surrogate_threshold(sensor, "coh"), ANALYSES[1]: surrogate_threshold(sensor, "wpli"),
           ANALYSES[2]: surrogate_threshold(csd, "wpli"), ANALYSES[3]: surrogate_threshold(tc_orth, "wpli")}
    results[subject] = {"matrices": M, "thresholds": thr, "n_epochs": len(sensor), "channels": list(ch)}
    row = {"subject": subject, "n_epochs": len(sensor), "alpha_ratio": meta["alpha_check"]["ratio"]}
    for a in ANALYSES:
        v = L5.upper_pairs(M[a])
        row[a] = (float(np.median(v)), int((v > thr[a]).sum()), int(v.size), float(thr[a]))
    per_subject.append(row)
    print(f"  {subject}: {len(sensor)} epochs, {len(ch)} channels, eyes-closed/open alpha ratio "
          f"{meta['alpha_check']['ratio']:.2f}")
    for a in ANALYSES:
        med, k, n, t = row[a]
        print(f"      {a:32s} median {med:.4f}, {k:4d} of {n:4d} pairs above the null ({t:.4f})")
print()
print(f"backend for every connectivity estimate: {backend['backend']}")
print(L5.disk_line("disk after every download was deleted"))

print(f"{len(results)} subjects, alpha band {BAND[0]:g}-{BAND[1]:g} Hz")
print()
print(f"{'analysis':>34s} {'unit':>14s} {'pairs':>7s} {'median (mean over subjects)':>28s} "
      f"{'significant pairs, mean':>25s} {'%':>7s}")
UNITS = {ANALYSES[0]: "channel pair", ANALYSES[1]: "channel pair", ANALYSES[2]: "channel pair",
         ANALYSES[3]: "region pair"}
group = {}
for a in ANALYSES:
    meds = np.array([r[a][0] for r in per_subject])
    ks = np.array([r[a][1] for r in per_subject], float)
    ns = np.array([r[a][2] for r in per_subject], float)
    group[a] = (meds, ks, ns)
    print(f"{a:>34s} {UNITS[a]:>14s} {int(ns[0]):7d} {meds.mean():28.4f} {ks.mean():25.1f} "
          f"{100 * (ks / ns).mean():7.1f}")
print()
print("The three sensor analyses share a channel set, so their pair counts are identical and comparable.")
print(f"The source analysis has {N_REGIONS} regions and therefore {int(group[ANALYSES[3]][2][0])} pairs; its")
print("percentage is comparable, its count is not.")

# How much do the three sensor analyses agree, pair by pair, within a subject?
print("Agreement between the three sensor-level analyses, as a correlation across channel pairs "
      "(per subject):")
print(f"{'subject':>12s} " + " ".join(f"{a.split()[0] + '-' + b.split()[0]:>10s}"
                                      for a in ANALYSES[:3] for b in ANALYSES[:3] if a < b))
agree = {}
for subject, res in results.items():
    line, vals = f"{subject:>12s} ", []
    for a in ANALYSES[:3]:
        for b in ANALYSES[:3]:
            if a < b:
                r = float(np.corrcoef(L5.upper_pairs(res["matrices"][a]),
                                      L5.upper_pairs(res["matrices"][b]))[0, 1])
                vals.append(((a, b), r))
                line += f"{r:10.3f} "
    agree[subject] = vals
    print(line)
print()
pairs = [k for k, _ in agree[list(agree)[0]]]
for k in pairs:
    rs = [dict(v)[k] for v in agree.values()]
    print(f"  {k[0]} vs {k[1]}: mean r = {np.mean(rs):+.3f} over {len(rs)} subjects")
print()
print("Two analyses of the same epochs correlating at r below 0.5 across pairs are not two views of one")
print("result.  They are different results, and the capstone's job is to say why.")

subject0 = list(results)[0]
fig, axes = plt.subplots(1, 4, figsize=(19, 4.4))
for ax, a in zip(axes, ANALYSES):
    M = results[subject0]["matrices"][a]
    lab = REGIONS if a == ANALYSES[3] else None
    L5.plot_matrix(M, lab, title=f"{a}\n(median {np.median(L5.upper_pairs(M)):.3f})", ax=ax, vmin=0,
                   vmax=1.0 if "coherence" in a else 0.5, cbar_label="dimensionless (0 to 1)")
fig.suptitle(f"C5, ds-lemon {subject0}, eyes closed, {BAND[0]:g}-{BAND[1]:g} Hz: the same epochs, four "
             f"analyses", y=1.02, fontsize=11)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2))
x = np.arange(len(ANALYSES))
axes[0].bar(x, [group[a][0].mean() for a in ANALYSES],
            yerr=[group[a][0].std() for a in ANALYSES], capsize=4, color="C0")
axes[0].set_xticks(x); axes[0].set_xticklabels([a.split()[0] for a in ANALYSES])
axes[0].set_ylabel("median connectivity over pairs (dimensionless)")
axes[0].set_title(f"mean and SD over {len(results)} subjects", fontsize=9)
axes[1].bar(x, [100 * (group[a][1] / group[a][2]).mean() for a in ANALYSES],
            yerr=[100 * (group[a][1] / group[a][2]).std() for a in ANALYSES], capsize=4, color="C1")
axes[1].set_xticks(x); axes[1].set_xticklabels([a.split()[0] for a in ANALYSES])
axes[1].set_ylabel("pairs above the surrogate null (% of all pairs)")
axes[1].set_title("how many connections each analysis calls significant", fontsize=9)
for ax in axes:
    ax.grid(alpha=0.25, axis="y")
fig.suptitle("C5 group summary: four analyses of identical epochs", y=1.02, fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

A, B, Cc, D = ANALYSES
g = {a: (group[a][0].mean(), 100 * (group[a][1] / group[a][2]).mean()) for a in ANALYSES}
mean_r = {k: float(np.mean([dict(v)[k] for v in agree.values()])) for k in pairs}

print("DIVERGENCE PARAGRAPH (draft; TODO(confirm) at author review)")
print()
print(f"Four analyses of the same {int(np.mean([r['n_epochs'] for r in per_subject])):.0f} eyes-closed epochs "
      f"per subject, over {len(results)} subjects, in the {BAND[0]:g}-{BAND[1]:g} Hz band, give four different")
print("answers, and every difference has a cause that can be named.")
print()
print(f"(1) A against B -- coherence {g[A][0]:.3f} against wPLI {g[B][0]:.3f}, correlating at "
      f"r = {mean_r[(A, B)]:+.3f} across pairs,")
print(f"    and {g[A][1]:.1f} % of pairs significant against {g[B][1]:.1f} %.  Coherence keeps everything two "
      f"channels share,")
print("    lagged or not; wPLI keeps only the lagged part.  Volume conduction is instantaneous, so the gap")
print("    between them is the share of the shared variance that has no lag in it -- which nb-5-1 shows is")
print("    everything, when there is one generator.  The honest reading of A is 'how much do these two")
print("    electrodes have in common', not 'how connected are these two places'.")
print()
print(f"(2) B against C -- wPLI {g[B][0]:.3f} against wPLI-after-CSD {g[Cc][0]:.3f}, correlating at "
      f"r = {mean_r[(B, Cc)]:+.3f}.")
print("    The surface Laplacian makes each channel a contrast with its neighbourhood, so the part of a")
print("    channel that is shared with the whole head is removed BEFORE the connectivity estimate.  It also")
print("    removes any genuinely broad or deep generator (nb-5-3 measures both), and it is reference-free by")
print("    construction, so C is the only one of the four whose value does not depend on the reference at all.")
print()
print(f"(3) C against D -- sensor wPLI after CSD {g[Cc][0]:.3f} against leakage-corrected source wPLI "
      f"{g[D][0]:.3f}.")
print("    These are not the same units of analysis: C is a pair of electrodes and D is a pair of regions, so")
print("    the counts are not comparable and only the proportions are.  D is the only one of the four in")
print("    which the mixing has been modelled rather than avoided -- an inverse operator plus symmetric")
print("    orthogonalisation -- and it is also the one that depends on the most assumptions: a head model, a")
print("    regularization parameter, a region definition and an orthogonalisation convention.")
print()
print(f"(4) The reference.  A and B are computed on the average reference of "
      f"{len(results[subject0]['channels'])} channels; changing it")
print("    changes every value in both (nb-5-2 measures the size of that), C is immune to it, and D inherits")
print("    it through the inverse operator's leadfield.  Any of these numbers quoted without the reference is")
print("    incomplete.")
print()
print(f"(5) What NONE of the four can do.  The source analysis uses geometric regions on a concentric sphere,")
print("    because the template anatomy the brief names does not clear the licence gate.  No result here can")
print("    be attached to a brain structure, and a claim that named one would not be supported by anything")
print("    computed in this notebook.")

RUBRIC = [
    ("Reference stated", f"yes -- {L5.LEMON_PIPELINE['reference']}; {L5.LEMON_PIPELINE['reference_note']}. "
                         f"Analyses A, B and D depend on it; C does not."),
    ("CSD choices stated", "yes -- mne.preprocessing.compute_current_source_density(sphere='auto', "
                           + ", ".join(f"{k}={v}" for k, v in CSD_PARAMS.items())
                           + "); units uV/m^2 by the display convention of nb-5-3."),
    ("Leakage addressed", f"yes -- analysis D uses a minimum-norm inverse on a four-shell sphere and symmetric "
                          f"(Colclough-style) orthogonalisation of the {N_REGIONS} region time courses; "
                          f"nb-5-6 measures the leakage of the same operator directly from its resolution "
                          f"matrix."),
    ("A null that pays for multiple comparisons",
     f"yes -- {N_SURROGATE} epoch-shuffled surrogates per analysis per subject, max statistic over pairs, "
     f"95th percentile."),
    ("No claim exceeds what the method supports",
     "the regions are geometric and no anatomical name appears anywhere; coherence is described as shared "
     "variance rather than as connection; and every count is reported beside the null it was tested against."),
    ("Subset documented and a FULL_COHORT switch",
     f"yes -- SUBSET = {list(SUBSET)} ({len(SUBSET)} subjects); FULL_COHORT = True runs {len(FULL_LIST)} and "
     f"has not been executed for these outputs."),
    ("Downloads deleted", "yes -- every subject's BrainVision files are deleted in a finally before the next "
                          "subject is fetched; peak disk is one subject."),
]
def wrap(text, width=100):
    lines, cur = [], ""
    for w in text.split():
        if len(cur) + len(w) + 1 > width:
            lines.append(cur)
            cur = w
        else:
            cur = f"{cur} {w}".strip()
    if cur:
        lines.append(cur)
    return lines


for item, answer in RUBRIC:
    print(f"[x] {item}")
    for line in wrap(answer):
        print(f"      {line}")

print("nb-c5-connectivity -- C5 numbers (draft; TODO(confirm) at author review)")
print(f"Data: ds-lemon, eyes closed, first {MAX_MINUTES:g} min per subject; "
      f"{L5.licence_line('ds-lemon')}")
print(f"Subjects: {list(SUBJECTS)} (FULL_COHORT = {FULL_COHORT}; the full list has "
      f"{len(FULL_LIST)} subjects and was not executed)")
print(f"Band {BAND[0]:g}-{BAND[1]:g} Hz; {L5.LEMON_PIPELINE['filter_hz']} Hz filter; "
      f"{L5.LEMON_PIPELINE['reference']}; {int(np.mean([r['n_epochs'] for r in per_subject]))} epochs per "
      f"subject on average")
print()
print("PER SUBJECT")
print(f"{'subject':>12s} {'epochs':>7s} {'EC/EO alpha':>12s} " +
      " ".join(f"{a.split()[0]:>26s}" for a in ANALYSES))
for r in per_subject:
    print(f"{r['subject']:>12s} {r['n_epochs']:7d} {r['alpha_ratio']:12.2f} " +
          " ".join(f"{r[a][0]:.4f} ({r[a][1]:3d}/{r[a][2]:<4d}){'':>4s}" for a in ANALYSES))
print()
print("GROUP (mean over subjects)")
print(f"{'analysis':>34s} {'median':>9s} {'significant pairs':>19s} {'% of pairs':>12s}")
for a in ANALYSES:
    meds, ks, ns = group[a]
    print(f"{a:>34s} {meds.mean():9.4f} {ks.mean():19.1f} {100 * (ks / ns).mean():12.1f}")
print()
print("AGREEMENT between the three sensor-level analyses (mean correlation across channel pairs):")
for k in pairs:
    print(f"      {k[0]} vs {k[1]}: r = {mean_r[k]:+.3f}")
print()
print("WHAT THE LICENCE FINDING COST THIS CAPSTONE, AND WHAT IT DID NOT:")
print("  cost: the fourth analysis is on geometric regions of a sphere, not on an anatomical parcellation, so")
print("        no connection found here can be named.  The brief's 'parcel-level ... on fsaverage' is not")
print("        possible while ds-fsaverage and ds-mne-sample fail the spec 10.7 gate.")
print("  kept: the inverse operator, the leakage, the orthogonalisation, the surrogate null and every")
print("        comparison between the four analyses.  If the author settles the ds-mne-sample licence")
print("        (spec section 13), only the region definition changes.")
print()
print(L5.disk_line("disk at the end"))