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

# 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

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

# 4. Quiet the downloader.  pooch, which MNE uses to fetch datasets, logs
#    "Downloading file '...' from '...' to '<cache directory>'" at INFO, and that last field is an
#    ABSOLUTE PATH from whichever machine executed the notebook.  Absolute paths are not allowed in a
#    stored notebook (scripts/scrub-notebooks.py is a CI gate) and re-executing would put them straight
#    back, so the message is suppressed at the source rather than cleaned up afterwards.  Nothing is
#    hidden by this: every cell below prints the file NAMES it fetched and helpers_l4.Downloads prints
#    free disk before and after.  Please do not delete this as noise.
try:
    import pooch

    pooch.get_logger().setLevel("WARNING")
except Exception:  # pooch absent or its API moved: the scrub script is the backstop
    pass
print(f"MNE {mne.__version__}; helpers_l4 imported from notebooks/_shared")
print(f"downloads go to {helpers_l1.download_dir().name}/ (resolved relative to the working directory, "
      "or $EEG_COURSE_DOWNLOADS) and are deleted at the end of this notebook")

for key, value in L4.MI_PIPELINE.items():
    print(f"{key:16s} : {value}")
print()
print(f"default baseline  : {L4.TF_BASELINE[0]:g} .. {L4.TF_BASELINE[1]:g} s      "
      f"default active window: {L4.TF_ACTIVE[0]:g} .. {L4.TF_ACTIVE[1]:g} s")
print(f"frequency axis    : {L4.TF_FREQS[0]:g} .. {L4.TF_FREQS[-1]:g} Hz in "
      f"{L4.TF_FREQS[1] - L4.TF_FREQS[0]:g} Hz steps ({len(L4.TF_FREQS)} rows), "
      f"n_cycles = max(3, f/2), decim = {L4.TF_DECIM}")

SUBJECT = "S001"
CH = "C3"
COND = "T2"          # right-fist imagery; C3 is contralateral to the right hand

# Every download this notebook makes is registered with `dl` and deleted in the `finally`
# of the last cell (helpers_l4.Downloads.finish), which also prints free disk afterwards.
dl = L4.Downloads("nb-4-1").start()
epochs, info = L4.load_imagery_epochs(SUBJECT, L4.MI_RUNS, dl)
for k, v in info.items():
    print(f"  {k:16s}: {v}")
print()
print("downloaded for this notebook: " + ", ".join(p.name for p in dl.paths))
print("conditions in these runs: " + "; ".join(f"{k} = {v}" for k, v in L4.MI_CONDITIONS.items()))

x = epochs[COND].get_data(picks=[CH])[:, 0, :] * 1e6      # (n_trials, n_times) in uV
times_raw = epochs.times
sf = float(epochs.info["sfreq"])
print(f"{SUBJECT} {COND} ({L4.MI_CONDITIONS[COND]}): {x.shape[0]} trials x {x.shape[1]} samples at {sf:g} Hz, "
      f"channel {CH}, epoch {times_raw[0]:g} .. {times_raw[-1]:g} s")

parts = L4.evoked_and_induced(x, sf, times=times_raw)
freqs, times = parts["freqs"], parts["times"]
edge = L4.edge_seconds(freqs, L4.TF_CYCLES)
print(f"time-frequency call: {parts['call']}")
print(f"units: {parts['units']}")
print(f"wavelet half-width at the lowest frequency ({freqs[0]:g} Hz): {edge[0]:.4f} s -- "
      f"estimates inside {times[0]:g}..{times[0] + edge[0]:.3f} s and "
      f"{times[-1] - edge[0]:.3f}..{times[-1]:g} s are partly made of zero padding (pf-tf-edge-effects)")

fig, axes = plt.subplots(1, 3, figsize=(16.5, 4.2))
ax = axes[0]
ax.plot(times_raw, x.mean(0), lw=1.6, color="k", label=f"trial average, n = {x.shape[0]}")
ax.plot(times_raw, x[0], lw=0.6, color="0.6", alpha=0.8, label="one trial")
ax.axvline(0, color="k", lw=1.0, ls="--")
ax.axhline(0, color="gray", lw=0.6)
ax.set(xlabel="Time from cue (s)", ylabel="Amplitude (uV)",
       title=f"{SUBJECT} {CH}: the ERP of an imagery cue (uV, positive up)")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)

for ax, key, label in ((axes[1], "evoked", "Evoked power (power of the trial average)"),
                       (axes[2], "total", "Total power (average of single-trial power)")):
    P = L4.baseline_normalise(parts[key][0], times, L4.TF_BASELINE, mode="db")
    L4.plot_tfr(P, freqs, times, ax=ax, edge_s=edge, baseline=L4.TF_BASELINE, vlim=(-6, 6),
                cbar_label="Power change from baseline (dB)",
                title=f"{label}\n{SUBJECT} {CH}, {COND}, dB re {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

mu = (freqs >= L4.MU_BAND[0]) & (freqs <= L4.MU_BAND[1])
act = (times >= L4.TF_ACTIVE[0]) & (times <= L4.TF_ACTIVE[1])
base = (times >= L4.TF_BASELINE[0]) & (times <= L4.TF_BASELINE[1])
rows = []
for name in ("total", "evoked", "induced"):
    P = parts[name][0]
    pb = P[mu].mean(0)
    rows.append({"component": name,
                 f"baseline power {L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s (library units)": pb[base].mean(),
                 f"active power {L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s (library units)": pb[act].mean(),
                 "change (%)": 100 * (pb[act].mean() - pb[base].mean()) / pb[base].mean()})
print(f"Mu band ({L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz) at {CH}, {SUBJECT}, condition {COND}:")
print(L4.fmt_table(rows, floatfmt="{:.4g}"))
ev_frac = parts["evoked_fraction"][0][mu][:, act].mean()
print()
print(f"Evoked power is {100 * ev_frac:.2f} % of total power in the mu band over the active window. "
      f"The ERP is not where this effect lives.")
print("NOTE: 'library units' means MNE's Morlet power for an input in uV -- proportional to uV^2 with a "
      "frequency-dependent constant, not band power in uV^2. Ratios and percentages are unaffected by it; "
      "absolute values would be, so none is quoted as uV^2 here.")

SYN = dict(sfreq=250.0, tmin=-1.0, tmax=1.5, f_hz=10.0, n_cycles=7.0, onset_s=0.3,
           amplitude_uv=20.0, noise_exponent=1.343, noise_amplitude_uv=19.8, seed=20260918)
N_TRIALS = 100
JITTERS = (0.0, 90.0, 217.2, 360.0)

fig, axes = plt.subplots(2, len(JITTERS), figsize=(17, 7.2))
syn_freqs = np.arange(4.0, 20.5, 0.5)
syn_cycles = np.maximum(3.0, syn_freqs / 2.0)
peak_amps = {}
for j, W in enumerate(JITTERS):
    tr, tt = L4.synthetic_trials(N_TRIALS, phase_jitter_deg=W, **SYN)
    avg = tr.mean(0)
    near = np.abs(tt - SYN["onset_s"]) <= 2 * SYN["n_cycles"] / (2 * np.pi * SYN["f_hz"])
    peak_amps[W] = float(np.abs(avg[near]).max())
    ax = axes[0, j]
    ax.plot(tt, tr[0], lw=0.5, color="0.7", label="one trial")
    ax.plot(tt, avg, lw=1.6, color="k", label=f"average of {N_TRIALS}")
    ax.set_ylim(-45, 45)
    ax.set(xlabel="Time (s)", ylabel="Amplitude (uV)" if j == 0 else "",
           title=f"SYNTHETIC: phase jitter {W:g} deg (full width)\ntrial average peaks at "
                 f"{peak_amps[W]:.1f} uV")
    ax.axhline(0, color="gray", lw=0.6)
    ax.grid(alpha=0.3)
    if j == 0:
        ax.legend(fontsize=8)
    r = L4.morlet_power(tr[:, None, :], sfreq=SYN["sfreq"], freqs=syn_freqs, n_cycles=syn_cycles,
                        decim=2, times=tt)
    tot = r["power"].mean(0)[0]
    tt2 = r["times"]
    P = L4.baseline_normalise(tot, tt2, (-0.9, -0.4), mode="db")
    L4.plot_tfr(P, syn_freqs, tt2, ax=axes[1, j], vlim=(-8, 8), event_s=SYN["onset_s"],
                edge_s=L4.edge_seconds(syn_freqs, syn_cycles),
                cbar_label="Power change from baseline (dB)",
                title=f"SYNTHETIC total power, jitter {W:g} deg\n(dB re -0.9..-0.4 s)")
fig.suptitle("The burst leaves the average and stays in the power (synthetic; 10 Hz, 7 cycles, 20 uV on "
             "1/f noise of 19.8 uV SD)", y=1.01)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print("Peak absolute amplitude of the trial average near the burst (uV), synthetic:")
for W, a in peak_amps.items():
    print(f"  jitter {W:6.1f} deg (full width): {a:6.2f} uV   "
          f"({100 * a / peak_amps[0.0]:5.1f} % of the zero-jitter value)")

an = L4.half_amplitude_jitter()
print("Analytic:")
for k, v in an.items():
    print(f"  {k:26s}: {v if isinstance(v, str) else f'{v:.4f}'}")

GRID = np.arange(0.0, 360.0 + 1e-9, 5.0)
N_SEEDS = 20
tmpl, _ = L4.synthetic_trials(1, phase_jitter_deg=0.0, **{**SYN, "noise_amplitude_uv": 0.0})
template = tmpl[0]
tt = np.linspace(SYN["tmin"], SYN["tmax"], template.size)
near = np.abs(tt - SYN["onset_s"]) <= 2 * SYN["n_cycles"] / (2 * np.pi * SYN["f_hz"])
tnorm = float(template @ template)

curves = {"peak-abs": np.zeros((N_SEEDS, len(GRID))), "template": np.zeros((N_SEEDS, len(GRID)))}
for s in range(N_SEEDS):
    for g, W in enumerate(GRID):
        tr, _ = L4.synthetic_trials(N_TRIALS, phase_jitter_deg=W,
                                    **{**SYN, "seed": 20260918 + 1000 * s + g})
        avg = tr.mean(0)
        curves["peak-abs"][s, g] = np.abs(avg[near]).max()
        curves["template"][s, g] = (avg @ template) / tnorm * SYN["amplitude_uv"]

def first_crossing(grid, y, level):
    """First linear crossing of `level` -- what reading a slider gives."""
    for i in range(len(y) - 1):
        if (y[i] - level) * (y[i + 1] - level) <= 0 and y[i] != y[i + 1]:
            f = (level - y[i]) / (y[i + 1] - y[i])
            return grid[i] + f * (grid[i + 1] - grid[i])
    return float("nan")

half = {}
for name, C_ in curves.items():
    vals = np.array([first_crossing(GRID, C_[s], C_[s, 0] / 2) for s in range(N_SEEDS)])
    half[name] = (float(np.nanmean(vals)), float(np.nanstd(vals, ddof=1)), vals)
    print(f"  measured half-amplitude jitter, {name:9s}: {half[name][0]:6.1f} +/- {half[name][1]:4.1f} deg "
          f"over {N_SEEDS} seeds (grid step {GRID[1] - GRID[0]:g} deg, {N_TRIALS} trials per point)")

fig, ax = plt.subplots(figsize=(9.2, 4.4))
for name, colour in (("peak-abs", "tab:orange"), ("template", "tab:blue")):
    m = curves[name].mean(0)
    sd = curves[name].std(0, ddof=1)
    ax.plot(GRID, m, color=colour, lw=1.6, label=f"{name} reading")
    ax.fill_between(GRID, m - sd, m + sd, color=colour, alpha=0.2, lw=0)
    ax.axhline(m[0] / 2, color=colour, lw=0.8, ls=":")
ax.plot(GRID, SYN["amplitude_uv"] * np.abs(np.sinc(np.deg2rad(GRID) / (2 * np.pi))), color="k", lw=1.2,
        ls="--", label="analytic  A |sin(W/2)/(W/2)|")
ax.axvline(an["uniform_full_width_deg"], color="k", lw=1.0)
ax.annotate(f"{an['uniform_full_width_deg']:.1f} deg", (an["uniform_full_width_deg"], 18),
            xytext=(6, 0), textcoords="offset points", fontsize=9)
ax.axvline(an["gaussian_sigma_deg"], color="tab:red", lw=1.0, ls="-.")
ax.annotate(f"{an['gaussian_sigma_deg']:.1f} deg\n(the Gaussian-sigma reading)", (an["gaussian_sigma_deg"], 6),
            xytext=(6, 0), textcoords="offset points", fontsize=8, color="tab:red")
ax.set(xlabel="Phase jitter (degrees, full width of the uniform draw)",
       ylabel="Amplitude of the trial average (uV)",
       title=f"SYNTHETIC: how the trial average shrinks with phase jitter (uV; {N_TRIALS} trials, "
             f"{N_SEEDS} seeds, mean +/- SD)")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

WIDGET_KEY_DEG = 217.2       # site/notes/integration-phase3.md, data-p3a and widgets-I
WIDGET_TOLERANCE_DEG = 30.0

print(f"widget key                     : {WIDGET_KEY_DEG:.1f} deg (tolerance +/- {WIDGET_TOLERANCE_DEG:g})")
print(f"this notebook, analytic        : {an['uniform_full_width_deg']:.4f} deg")
print(f"this notebook, peak-abs sweep  : {half['peak-abs'][0]:.1f} +/- {half['peak-abs'][1]:.1f} deg")
print(f"this notebook, template sweep  : {half['template'][0]:.1f} +/- {half['template'][1]:.1f} deg")
agree = abs(an["uniform_full_width_deg"] - WIDGET_KEY_DEG) < 0.05
print()
print("AGREES with the widget key to within 0.05 deg." if agree else
      "DISAGREES with the widget key -- report both figures and find the difference before tuning either.")
print(f"Both sweeps lie inside the +/- {WIDGET_TOLERANCE_DEG:g} deg tolerance: "
      f"{all(abs(half[k][0] - WIDGET_KEY_DEG) < WIDGET_TOLERANCE_DEG for k in half)}")
print(f"The Gaussian-sigma reading, {an['gaussian_sigma_deg']:.1f} deg, is {abs(an['gaussian_sigma_deg'] - WIDGET_KEY_DEG):.0f} deg away "
      f"and is excluded by that tolerance -- which is what the tolerance is for.")

try:
    print("nb-4-1-evoked-vs-induced -- L4.1 exercise numbers (draft; TODO(confirm) at author review)")
    print(f"Real data: ds-eegbci {SUBJECT}, runs {'+'.join(info['runs'])} "
          f"(PhysioNet DOI {L4.DATASETS_L4['ds-eegbci']['dataset_doi']}; "
          f"{L4.DATASETS_L4['ds-eegbci']['license']}); {info['n_eeg']} EEG channels at {info['sfreq']:g} Hz, "
          f"average reference, no filtering.")
    print(f"Condition {COND} = {L4.MI_CONDITIONS[COND]}, {x.shape[0]} trials; channel {CH}.")
    print(f"Synthetic half: {N_TRIALS} trials of a {SYN['f_hz']:g} Hz {SYN['n_cycles']:g}-cycle burst, "
          f"{SYN['amplitude_uv']:g} uV, on 1/f^{SYN['noise_exponent']:g} noise of "
          f"{SYN['noise_amplitude_uv']:g} uV SD, seeded; helpers_l4.synthetic_trials.")
    print()
    print("ANSWER KEY -- ex-4-1 (numeric): at what phase jitter does the ERP amplitude halve?")
    print(f"    {an['uniform_full_width_deg']:.1f} degrees, where the parameter is the FULL WIDTH of a "
          f"uniform phase draw.")
    print(f"    Exact in expectation: the average is the zero-jitter burst times sin(W/2)/(W/2), and "
          f"sin(x)/x = 1/2 at x = {an['uniform_full_width_rad'] / 2:.6f} rad.")
    print(f"    Measured on this notebook's own simulation: "
          f"{half['peak-abs'][0]:.1f} +/- {half['peak-abs'][1]:.1f} deg read as a peak, "
          f"{half['template'][0]:.1f} +/- {half['template'][1]:.1f} deg read as a template projection "
          f"({N_SEEDS} seeds, {N_TRIALS} trials, {GRID[1] - GRID[0]:g} deg grid).")
    print(f"    THE PROMPT MUST NAME THE CONVENTION. Read as a Gaussian standard deviation the same "
          f"question answers {an['gaussian_sigma_deg']:.1f} degrees; read as half the slider it answers 180; "
          f"the jitter at which the ERP reaches ZERO (not half) is 360.")
    print(f"    Cross-check: site/notes/integration-phase3.md records {WIDGET_KEY_DEG:g} deg from data-p3a "
          f"and widgets-I. This notebook {'AGREES' if agree else 'DISAGREES'}.")
    print()
    print("Supporting -- the dissociation on real trials (the point of the lesson, not a keyed answer):")
    for r in rows:
        print(f"    {r['component']:8s} power in {L4.MU_BAND[0]:g}-{L4.MU_BAND[1]:g} Hz over "
              f"{L4.TF_ACTIVE[0]:g}..{L4.TF_ACTIVE[1]:g} s changes by {r['change (%)']:+7.1f} % from the "
              f"{L4.TF_BASELINE[0]:g}..{L4.TF_BASELINE[1]:g} s baseline")
    print(f"    evoked power is {100 * ev_frac:.2f} % of total power there: the ERP carries almost none "
          f"of it.")
    print(f"    Peak absolute ERP amplitude at {CH} over the active window: "
          f"{np.abs(x.mean(0)[(times_raw >= 0) & (times_raw <= 4)]).max():.2f} uV.")
    print()
    print("Pitfalls: none named for L4.1 in the spec; pf-tf-edge-effects applies to every map drawn here "
          "and the edge region is shaded in each. Widget: w-evoked-vs-induced.")
    print("TODO(confirm): the catalog carries no published ERD percentage for this dataset, so nothing "
          "above is compared with a literature value.")
finally:
    dl.finish()