# 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 (notebooks/_shared/helpers.py and helpers_l1.py), located
#    relative to the working directory -- notebooks/<level>/ or notebooks/ --
#    never through an absolute path.
_shared = next((d / "_shared" for d in (Path.cwd(), *Path.cwd().parents)
                if (d / "_shared" / "helpers_l1.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L1/ (or notebooks/) so that _shared/helpers_l1.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l1

# 3. Plotting: Jupyter's default inline backend renders static PNGs through Agg
#    (no windows, nothing blocks); outside Jupyter the helpers select Agg.  Every
#    MNE figure is requested with show=False, and plt.show() renders each cell's
#    figures in place.
import matplotlib.pyplot as plt
import numpy as np
import mne
import pooch

mne.viz.set_browser_backend("matplotlib", verbose=False)
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 and helpers_l1 imported from notebooks/_shared; downloads go to the course "
      "data directory (EEG_COURSE_DOWNLOADS or EEG_COURSE_DATA if set, else MNE's data directory under eeg-course/)")

from IPython.display import HTML, display
from scipy import signal

DATASET, SUBJECT = "ds-iowapd", "sub-007"
raw_pd, pd_paths = helpers_l1.load_iowapd(SUBJECT, return_paths=True)
SF = raw_pd.info["sfreq"]
sidecar = helpers_l1.read_sidecar(pd_paths["eeg.json"])
rep = helpers.first_look_report(raw_pd, DATASET)
display(HTML(helpers.report_html({k: rep[k] for k in ("description", "sfreq_hz", "nyquist_hz", "n_channels", "duration_s",
                                                       "highpass_hz_in_header", "lowpass_hz_in_header", "bads", "amplitude_uV")},
                                 f"First look: {DATASET} {SUBJECT}")))
print("BIDS sidecar:", {k: sidecar.get(k) for k in ("PowerLineFrequency", "EEGReference", "SamplingFrequency", "SoftwareFilters")})

WELCH = dict(seg_s=4.0, overlap=0.5, window="hann")
f_pd, psd_pd, ch_pd = helpers_l1.welch_raw(raw_pd, **WELCH)
peaks = helpers_l1.line_peaks(f_pd, psd_pd, fmin=5, min_prominence_db=6)
print("sharp peaks in the channel-mean PSD: " + ", ".join(f"{p['freq_hz']:.2f} Hz (+{p['prominence_db']:.0f} dB, {p['width_hz']:.2f} Hz wide)" for p in peaks))
MAINS = int(sidecar.get("PowerLineFrequency", helpers_l1.DATASETS_L1[DATASET]["mains_hz"]))
candidates = [k * MAINS for k in range(1, int(SF / 2 // MAINS) + 1)]
LINES = [f for f in candidates if any(abs(p["freq_hz"] - f) < 1.0 for p in peaks)]
print(f"multiples of {MAINS} Hz below the {SF / 2:g} Hz Nyquist frequency: {candidates}; found as sharp peaks: {LINES}")
print(f"200 Hz (removed by the authors in their analysis): {helpers_l1.peak_prominence_db(f_pd, psd_pd, 200.0):+.1f} dB over its flanks here -- not a sharp peak in this subject")
prom60 = np.array([helpers_l1.peak_prominence_db(f_pd, p, 60.0) for p in psd_pd])
CH = ch_pd[int(np.argmax(prom60))]
print(f"60 Hz prominence per channel: median +{np.median(prom60):.0f} dB, strongest {CH} +{prom60.max():.0f} dB, weakest {ch_pd[int(np.argmin(prom60))]} +{prom60.min():.0f} dB")

fig = plt.figure(figsize=(13, 4))
ax1 = fig.add_axes([0.06, 0.15, 0.6, 0.75])
helpers_l1.plot_spectra({f"mean of {len(ch_pd)} channels": (f_pd, psd_pd), f"{CH}": (f_pd, psd_pd[ch_pd.index(CH)], dict(lw=0.7, alpha=0.8))},
                        ax=ax1, lines=LINES, title=f"{DATASET} {SUBJECT}: the line and its harmonics below Nyquist")
ax2 = fig.add_axes([0.7, 0.05, 0.28, 0.9])
helpers.plot_topomap_values(raw_pd, prom60, unit="dB over flanks", title="60 Hz prominence", ax=ax2)
plt.show()   # render the static figure(s) of this cell inline

def residual_table(rows, f, psd_before, lines, half_width=1.0):
    """Print residual-power and collateral metrics for {label: psd_after}."""
    print(f"{'method':46s} {'ratio':>7s} {'excess':>7s} {'flank':>6s} {'1-45 Hz':>8s}   per-line ratio")
    out = {}
    df = f[1] - f[0]
    bb = (f >= 1) & (f <= 45)
    p0 = psd_before.mean(axis=0)[bb].sum() * df
    for label, psd_after in rows.items():
        r = helpers_l1.residual_power_ratio(f, psd_before, psd_after, lines, half_width=half_width)
        r["broadband_ratio"] = psd_after.mean(axis=0)[bb].sum() * df / p0
        out[label] = r
        print(f"{label:46s} {r['ratio']:7.4f} {r['ratio_excess']:7.3f} {r['flank_ratio']:6.3f} {r['broadband_ratio']:8.4f}   "
              + " ".join(f"{v:.3f}" for v in r["ratio_per_line"]))
    return out


def trace_panel(ax, sig_uV, sf, t0, dur, title):
    n0, n1 = int(t0 * sf), int((t0 + dur) * sf)
    ax.plot(np.arange(n0, n1) / sf, sig_uV[n0:n1], "k", lw=0.8)
    ax.set(title=title, xlabel="Time (s)", ylabel="uV")
    ax.grid(alpha=0.3)


print("Metric definitions (used throughout):")
print("  ratio        = line-band power after / before, summed over +/-1 Hz around each line (1 = nothing removed, 0 = nothing left, EEG included)")
print("  excess       = the same for the power ABOVE the local background (median PSD 2-5 Hz on either side); 0 = exactly the background is left, < 0 = a hole")
print("  flank        = power 2-5 Hz on either side of each line, after / before (collateral damage next to the lines; 1 = none)")
print("  1-45 Hz      = broadband power after / before (collateral damage in the EEG band; 1 = none)")

raw_fir = raw_pd.copy().notch_filter(LINES, picks="eeg", verbose=False)                         # MNE defaults
raw_wide = raw_pd.copy().notch_filter(LINES, picks="eeg", notch_widths=4.0, verbose=False)      # too wide, for contrast

imp = np.zeros(int(30 * SF))
imp[len(imp) // 2] = 1.0
h_notch = mne.filter.notch_filter(imp[None], SF, LINES, verbose=False)[0]          # effective (zero-phase) impulse response
h_wide = mne.filter.notch_filter(imp[None], SF, LINES, notch_widths=4.0, verbose=False)[0]
w, H = signal.freqz(h_notch, worN=1 << 17, fs=SF)
_, Hw = signal.freqz(h_wide, worN=1 << 17, fs=SF)
gain = 20 * np.log10(np.maximum(np.abs(H), 1e-12))
gain_w = 20 * np.log10(np.maximum(np.abs(Hw), 1e-12))
for lab, g in (("default", gain), ("4 Hz wide", gain_w)):
    m = (w > 55) & (w < 65)
    below3 = w[m][g[m] <= -3]
    print(f"notch {lab:9s}: gain at 60 Hz {g[np.argmin(np.abs(w - 60))]:.0f} dB; -3 dB width around 60 Hz {below3.max() - below3.min():.2f} Hz; "
          f"impulse response above 1 % of its peak spans {np.ptp(np.where(np.abs(h_notch if lab == 'default' else h_wide) > 0.01 * np.abs(h_notch if lab == 'default' else h_wide).max())[0]) / SF:.2f} s")

fig, axes = plt.subplots(1, 2, figsize=(12, 3.6))
axes[0].plot(w, gain, label="MNE default (notch_widths = f/200)")
axes[0].plot(w, gain_w, label="notch_widths = 4 Hz")
axes[0].set(xlim=(50, 70), ylim=(-70, 5), xlabel="Frequency (Hz)", ylabel="Gain (dB)", title="Notch at 60 Hz: magnitude response, zoom (dB)")
axes[0].grid(alpha=0.3); axes[0].legend(fontsize=8)
tt = (np.arange(len(imp)) - len(imp) // 2) / SF
axes[1].plot(tt, h_notch, lw=0.6, label="MNE default")
axes[1].plot(tt, h_wide, lw=0.6, label="4 Hz wide")
axes[1].set(xlim=(-4, 4), xlabel="Time (s)", ylabel="Impulse response (a.u.)", title="Notch impulse responses (zero-phase, centred on 0)")
axes[1].grid(alpha=0.3); axes[1].legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

import time

t0 = time.time()
raw_sf = raw_pd.copy().notch_filter(LINES, method="spectrum_fit", picks="eeg", verbose=False)
t_sf = time.time() - t0
t0 = time.time()
raw_sfa = raw_pd.copy().notch_filter(None, method="spectrum_fit", picks="eeg", verbose=False)   # frequencies detected automatically
t_sfa = time.time() - t0
print(f"spectrum_fit on {len(ch_pd)} channels x {raw_pd.times[-1]:.0f} s: {t_sf:.1f} s with the frequencies given, {t_sfa:.1f} s with automatic detection")

x_pd = helpers_l1.to_uV(raw_pd)                                   # (n_channels, n_times) in uV, bads excluded
t0 = time.time()
x_zl, zl = helpers_l1.zapline_like(x_pd, SF, float(MAINS))
t_zl = time.time() - t0
print(f"ZapLine-style: period {zl['period_samples']:.2f} samples, {zl['n_harmonics']} harmonics below Nyquist in the bias bands, "
      f"{zl['n_components']} components; removed {zl['n_removed']} (scores: {', '.join(f'{s:.2f}' for s in zl['scores'][:zl['n_removed'] + 3])} ...); {t_zl:.1f} s")
raw_zl = raw_pd.copy()
raw_zl._data[mne.pick_types(raw_pd.info, eeg=True, exclude="bads")] = x_zl * 1e-6

fig, axes = plt.subplots(1, 2, figsize=(12, 3.6))
n_show = 15
colors = ["tab:orange" if i < zl["n_removed"] else "tab:blue" for i in range(n_show)]
axes[0].bar(np.arange(1, n_show + 1), zl["scores"][:n_show], color=colors)
axes[0].set(xlabel="Component (sorted)", ylabel="Share of variance in the line bands", title=f"ZapLine-style scores; orange = removed ({zl['n_removed']}) (fraction, 0-1)")
axes[0].grid(alpha=0.3, axis="y")
removed_rms = zl["removed"].std(axis=1)
helpers.plot_topomap_values(raw_pd, 20 * np.log10(removed_rms / removed_rms.max()), unit="dB re max", title="Removed line signal (RMS per channel)", ax=axes[1])
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

after = {}
for label, r in (("FIR notch, MNE defaults", raw_fir), ("FIR notch, 4 Hz wide", raw_wide),
                 ("spectrum_fit, frequencies given", raw_sf), ("spectrum_fit, automatic frequencies", raw_sfa)):
    after[label] = helpers_l1.welch_raw(r, **WELCH)[1]
after[f"ZapLine-style, {zl['n_removed']} components (auto)"] = helpers_l1.welch_psd(x_zl, SF, **WELCH)[1]
for n in (1, 2):
    x_n, _ = helpers_l1.zapline_like(x_pd, SF, float(MAINS), n_remove=n)
    after[f"ZapLine-style, {n} component{'s' if n > 1 else ''}"] = helpers_l1.welch_psd(x_n, SF, **WELCH)[1]
results_pd = residual_table(after, f_pd, psd_pd, LINES)

show = {"original": (f_pd, psd_pd)}
show.update({k: (f_pd, v) for k, v in after.items() if k in ("FIR notch, MNE defaults", "spectrum_fit, frequencies given") or k.startswith("ZapLine-style,") and "(auto)" in k})
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
helpers_l1.plot_spectra(show, ax=axes[0], lines=LINES, title=f"{SUBJECT}, channel mean: three methods, full band")
helpers_l1.plot_spectra(show, ax=axes[1], xlim=(55, 65), title="Zoom on the 60 Hz line")
fig.tight_layout()

i_ch = ch_pd.index(CH)
traces = {"original": x_pd[i_ch], "FIR notch (defaults)": raw_fir.get_data(picks=CH)[0] * 1e6,
          "spectrum_fit": raw_sf.get_data(picks=CH)[0] * 1e6, "ZapLine-style": x_zl[i_ch]}
fig, axes = plt.subplots(len(traces), 1, figsize=(11, 8), sharex=True, sharey=True)
for ax, (label, y) in zip(axes, traces.items()):
    trace_panel(ax, y, SF, 100.0, 0.3, f"{CH}, {label} (uV)")
fig.suptitle(f"{SUBJECT} {CH}: 0.3 s of the raw trace before and after each method (uV)", y=1.0)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

raw_bci = helpers.load_spine("ds-eegbci", "S001", "R01")     # eyes open, ~1 min, 160 Hz
SF_B = raw_bci.info["sfreq"]
f_b, psd_b, ch_b = helpers_l1.welch_raw(raw_bci, **WELCH)
peaks_b = helpers_l1.line_peaks(f_b, psd_b, fmin=5, min_prominence_db=4)
print(f"ds-eegbci S001 R01 ({SF_B:g} Hz, Nyquist {SF_B / 2:g} Hz): sharp peaks >= 4 dB: "
      + (", ".join(f"{p['freq_hz']:.2f} Hz (+{p['prominence_db']:.1f} dB)" for p in peaks_b) or "none"))
for f_true in (120, 180):
    fa = helpers_l1.alias_frequency(f_true, SF_B)
    print(f"  {f_true} Hz would alias to {fa:g} Hz: PSD there is {helpers_l1.peak_prominence_db(f_b, psd_b, fa):+.1f} dB over its flanks "
          f"-> {'a sharp peak' if any(abs(p['freq_hz'] - fa) < 1 for p in peaks_b) else 'no sharp peak'} (this subject)")
raw_bci_n = raw_bci.copy().notch_filter([60.0], picks="eeg", verbose=False)
res_bci = residual_table({"FIR notch at 60 Hz only (MNE defaults)": helpers_l1.welch_raw(raw_bci_n, **WELCH)[1]}, f_b, psd_b, [60.0])
ax = helpers_l1.plot_spectra({"S001 R01, channel mean": (f_b, psd_b), "after a 60 Hz notch": (f_b, helpers_l1.welch_raw(raw_bci_n, **WELCH)[1])},
                             lines=[20, 40, 60], title=f"ds-eegbci S001 R01 at {SF_B:g} Hz: 60 Hz is in band; dotted 20/40 Hz = where 180/120 Hz would have folded")
plt.show()   # render the static figure(s) of this cell inline

raw_de, de_paths = helpers_l1.load_dortmund("sub-001", "1", "EyesClosed", "pre", return_paths=True)
SF_D = raw_de.info["sfreq"]
sidecar_de = helpers_l1.read_sidecar(de_paths["eeg.json"])
print("BIDS sidecar:", {k: sidecar_de.get(k) for k in ("PowerLineFrequency", "EEGReference", "SamplingFrequency", "SoftwareFilters", "RecordingDuration")})
eeg_idx = mne.pick_types(raw_de.info, eeg=True)
sd = raw_de.get_data(picks=eeg_idx).std(axis=1) * 1e6
raw_de.info["bads"] = [raw_de.ch_names[i] for i, s in zip(eeg_idx, sd) if s > 5 * np.median(sd)]
print(f"{len(eeg_idx)} EEG channels; marked bad (SD > 5 x median of {np.median(sd):.0f} uV): {raw_de.info['bads'] or 'none'}")
f_d, psd_d, ch_d = helpers_l1.welch_raw(raw_de, **WELCH)
peaks_d = helpers_l1.line_peaks(f_d, psd_d, fmin=5, min_prominence_db=6)
print("sharp peaks in the channel mean: " + ", ".join(f"{p['freq_hz']:.1f} Hz (+{p['prominence_db']:.0f} dB)" for p in peaks_d))
MAINS_D = int(sidecar_de.get("PowerLineFrequency", 50))
cand_d = [k * MAINS_D for k in range(1, int(SF_D / 2 // MAINS_D) + 1)]
LINES_D = [f for f in cand_d if helpers_l1.peak_prominence_db(f_d, psd_d, f) >= 3.0]
print(f"multiples of {MAINS_D} Hz at least 3 dB above their flanks: {LINES_D}  (of {cand_d})")

raw_de_fir = raw_de.copy().notch_filter(LINES_D, picks="eeg", verbose=False)
x_de = helpers_l1.to_uV(raw_de)
x_de_zl, zl_de = helpers_l1.zapline_like(x_de, SF_D, float(MAINS_D))
print(f"ZapLine-style on ds-dortmund: period {zl_de['period_samples']:.0f} samples (an integer at 1000 Hz), removed {zl_de['n_removed']} components")
after_d = {"FIR notch, MNE defaults": helpers_l1.welch_raw(raw_de_fir, **WELCH)[1],
           f"ZapLine-style, {zl_de['n_removed']} components (auto)": helpers_l1.welch_psd(x_de_zl, SF_D, **WELCH)[1]}
results_de = residual_table(after_d, f_d, psd_d, LINES_D)
fig, axes = plt.subplots(1, 2, figsize=(13, 4))
helpers_l1.plot_spectra({"original": (f_d, psd_d), **{k: (f_d, v) for k, v in after_d.items()}}, ax=axes[0], lines=cand_d[:7],
                        title=f"ds-dortmund sub-001 EC pre, channel mean, {SF_D:g} Hz: 50 Hz and harmonics; 250 Hz online low-pass")
helpers_l1.plot_spectra({"original": (f_d, psd_d), **{k: (f_d, v) for k, v in after_d.items()}}, ax=axes[1], xlim=(45, 55), title="Zoom on 50 Hz")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

raw_ar, ar_path = helpers_l1.load_eegmat("Subject00", "rest", return_paths=True)
f_a, psd_a, ch_a = helpers_l1.welch_raw(raw_ar, **WELCH)
print(f"ds-arithmetic Subject00 rest: 50 Hz is {helpers_l1.peak_prominence_db(f_a, psd_a, 50.0):+.1f} dB relative to its flanks (negative = a hole); "
      f"sharp peaks >= 6 dB above 15 Hz: {[round(p['freq_hz'], 1) for p in helpers_l1.line_peaks(f_a, psd_a, fmin=15, min_prominence_db=6)] or 'none'}")
raw_ar_n = raw_ar.copy().notch_filter([50.0], picks="eeg", verbose=False)
res_ar = residual_table({"software notch at 50 Hz on top of the hardware hole": helpers_l1.welch_raw(raw_ar_n, **WELCH)[1]}, f_a, psd_a, [50.0])
print("(excess is not meaningful here: the band held less power than its flanks before anything was done)")
ax = helpers_l1.plot_spectra({"as recorded (hardware notch)": (f_a, psd_a), "after a software notch at 50 Hz": (f_a, helpers_l1.welch_raw(raw_ar_n, **WELCH)[1])},
                             xlim=(30, 70), lines=[50], title="ds-arithmetic Subject00, channel mean: a hole, then a wider hole")
plt.show()   # render the static figure(s) of this cell inline

DELETE_DOWNLOADS = False
if DELETE_DOWNLOADS:
    helpers_l1.cleanup(list(pd_paths.values()) + list(de_paths.values()) + [ar_path])

key = results_pd["FIR notch, MNE defaults"]
print("nb-1-6-line-noise -- L1.6 exercise numbers (draft; TODO(confirm) at author review)")
print(f"Data: {DATASET} {SUBJECT} (OpenNeuro ds004584 v1.0.0, CC0), {len(ch_pd)} channels, {raw_pd.times[-1]:.0f} s at {SF:g} Hz")
print(f"Lines found and removed: {LINES} Hz (sharp peaks at multiples of {MAINS} Hz below the {SF / 2:g} Hz Nyquist frequency)")
print(f"Metric: line-band power after / before, Welch {WELCH['seg_s']:g}-s {WELCH['window']} segments with {int(100 * WELCH['overlap'])} % overlap over the whole "
      f"recording, mean over the {len(ch_pd)} channels, +/- {key['half_width_hz']:g} Hz around each line")
print(f"  residual power ratio, FIR notch (raw.notch_filter, MNE defaults)      : {key['ratio']:.4f}   <-- proposed answer key")
print(f"  residual power ratio, spectrum_fit (frequencies given)              : {results_pd['spectrum_fit, frequencies given']['ratio']:.4f}")
zl_key = [k for k in results_pd if k.startswith("ZapLine-style,") and "(auto)" in k][0]
print(f"  residual power ratio, ZapLine-style ({zl['n_removed']} components)                : {results_pd[zl_key]['ratio']:.4f}")
print(f"  line-band power before: {key['band_power_before_uV2']:.2f} uV^2, of which background (flank median x width) {key['background_before_uV2']:.2f} uV^2; "
      f"the excess-over-background ratios are {key['ratio_excess']:+.3f} (notch), {results_pd['spectrum_fit, frequencies given']['ratio_excess']:+.3f} (spectrum_fit), "
      f"{results_pd[zl_key]['ratio_excess']:+.3f} (ZapLine-style)")
print(f"  suggested tolerance for the numeric exercise: answer {key['ratio']:.2f}, tolerance 0.03 (accepts any of the three methods)")
print(f"Second part (free-response): on ds-eegbci at {SF_B:g} Hz only 60 Hz is below the {SF_B / 2:g} Hz Nyquist frequency; 120 and 180 Hz are not in the file, "
      f"so 'every harmonic' cannot be removed at its true frequency. For information: a 60 Hz notch on S001 R01 leaves a ratio of "
      f"{res_bci['FIR notch at 60 Hz only (MNE defaults)']['ratio']:.4f} in the 60 Hz band.")
print(f"50 Hz example: ds-dortmund sub-001 EC pre, lines {LINES_D} Hz, FIR notch ratio {results_de['FIR notch, MNE defaults']['ratio']:.4f}")