# 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

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"])
KEYS = ["description", "sfreq_hz", "nyquist_hz", "n_channels", "duration_s", "highpass_hz_in_header",
        "lowpass_hz_in_header", "montage_attached", "bads", "amplitude_uV"]
rep = helpers.first_look_report(raw_pd, DATASET)
display(HTML(helpers.report_html({k: rep[k] for k in KEYS}, f"First look: {DATASET} {SUBJECT}")))
print("BIDS sidecar:", {k: sidecar.get(k) for k in ("SamplingFrequency", "PowerLineFrequency", "EEGReference",
                                                    "EEGChannelCount", "RecordingDuration", "SoftwareFilters")})
facts = helpers_l1.DATASETS_L1[DATASET]
print(f"Catalog: {facts['sfreq']} Hz, {facts['n_channels']} channels, {facts['mains_hz']} Hz mains; {facts['hardware_filters']}; "
      f"license {facts['license']}")

f_pd, psd_pd, ch_pd = helpers_l1.welch_raw(raw_pd, seg_s=4.0)          # Welch, 4-s Hann segments, 50 % overlap, uV^2/Hz
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))]                                      # the channel with the strongest 60 Hz line (a stated rule)
peaks = helpers_l1.line_peaks(f_pd, psd_pd, fmin=5, min_prominence_db=6)
print(f"sharp peaks (>= 6 dB over their surroundings) in the channel-mean PSD: "
      + ", ".join(f"{p['freq_hz']:.1f} Hz (+{p['prominence_db']:.0f} dB)" for p in peaks))
print(f"strongest 60 Hz line: {CH} (+{prom60.max():.0f} dB over its flanks); median channel +{np.median(prom60):.0f} dB")
HARMONICS = [60, 120, 180, 240]
ax = helpers_l1.plot_spectra({f"mean of {len(ch_pd)} channels": (f_pd, psd_pd),
                              f"{CH} (strongest 60 Hz line)": (f_pd, psd_pd[ch_pd.index(CH)], dict(lw=0.7, alpha=0.8))},
                             title=f"{DATASET} {SUBJECT}: PSD up to the {SF / 2:g} Hz Nyquist frequency; dotted = 60 Hz and harmonics",
                             lines=HARMONICS)
plt.show()   # render the static figure(s) of this cell inline

FS_NEW = 100.0
print(f"decimating {SF:g} -> {FS_NEW:g} Hz (new Nyquist {FS_NEW / 2:g} Hz): f_alias = |f - round(f / fs) * fs|")
for f in (60, 70, 120, 180, 240):
    fa = helpers_l1.alias_frequency(f, FS_NEW)
    print(f"  {f:4d} Hz -> {fa:5.1f} Hz" + ("   (within the new band: unchanged)" if f <= FS_NEW / 2 else ""))
print(f"textbook case: 140 Hz sampled at 200 Hz -> {helpers_l1.alias_frequency(140, 200):g} Hz")

f_true = np.linspace(0, SF / 2, 2001)
fig, ax = plt.subplots(figsize=(9, 3.6))
ax.plot(f_true, helpers_l1.alias_frequency(f_true, FS_NEW), "k", lw=1)
for f, c in zip((60, 70, 120, 180, 240), ("tab:orange", "tab:blue", "tab:orange", "tab:orange", "tab:orange")):
    ax.plot(f, helpers_l1.alias_frequency(f, FS_NEW), "o", color=c)
    ax.annotate(f"{f} -> {helpers_l1.alias_frequency(f, FS_NEW):g} Hz", (f, helpers_l1.alias_frequency(f, FS_NEW)),
                textcoords="offset points", xytext=(4, 6), fontsize=8, color=c)
ax.set(xlabel="True frequency (Hz)", ylabel="Apparent frequency (Hz)",
       title=f"Folding at fs = {FS_NEW:g} Hz: apparent frequency of each true frequency (Hz); orange = real lines, blue = the 70 Hz marker")
ax.grid(alpha=0.3)
plt.show()   # render the static figure(s) of this cell inline

TONE_HZ = 70.0
t_fine = np.arange(0, 0.2, 1 / 5000)            # a dense grid standing in for the continuous signal
t_s = np.arange(0, 0.2, 1 / FS_NEW)             # the samples at 100 Hz
tone = lambda t: np.sin(2 * np.pi * TONE_HZ * t)
fig, ax = plt.subplots(figsize=(10, 3.4))
ax.plot(t_fine * 1000, tone(t_fine), color="0.6", lw=1, label=f"{TONE_HZ:g} Hz tone (continuous)")
ax.plot(t_s * 1000, tone(t_s), "o", color="tab:blue", ms=6, label=f"samples at {FS_NEW:g} Hz")
ax.plot(t_fine * 1000, -np.sin(2 * np.pi * helpers_l1.alias_frequency(TONE_HZ, FS_NEW) * t_fine), "--", color="tab:orange", lw=1.2,
        label=f"{helpers_l1.alias_frequency(TONE_HZ, FS_NEW):g} Hz sinusoid (sign reversed) through the same samples")
ax.set(xlabel="Time (ms)", ylabel="Amplitude (a.u.)", title=f"A {TONE_HZ:g} Hz tone sampled at {FS_NEW:g} Hz is a {helpers_l1.alias_frequency(TONE_HZ, FS_NEW):g} Hz tone (amplitude in arbitrary units)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8, loc="upper right")
plt.show()   # render the static figure(s) of this cell inline

from scipy import signal

TONE_UV = 20.0
Q = int(round(SF / FS_NEW))                                   # decimation factor 5
x = raw_pd.get_data(picks=CH)[0] * 1e6                         # uV
t = raw_pd.times
x_test = x + TONE_UV * np.sin(2 * np.pi * TONE_HZ * t)         # SYNTHETIC marker added to the real channel
raw_test = mne.io.RawArray(x_test[None] * 1e-6, mne.create_info([f"{CH}+70Hz"], SF, "eeg"), verbose=False)

versions = {
    f"original, {SF:g} Hz": (SF, x_test),
    "naive decimation x[::5] (no filter)": (FS_NEW, helpers_l1.naive_decimate(x_test, Q)),
    "raw.resample(100) (MNE: anti-alias low-pass, then decimate)": (FS_NEW, raw_test.copy().resample(FS_NEW, verbose=False).get_data()[0] * 1e6),
    "scipy.signal.decimate (FIR anti-alias, zero-phase)": (FS_NEW, signal.decimate(x_test, Q, ftype="fir", zero_phase=True)),
}
spectra = {}
for label, (fs, y) in versions.items():
    fr, p = helpers_l1.welch_psd(y, fs, seg_s=4.0)
    spectra[label] = (fr, p)
    pk = helpers_l1.line_peaks(fr, p, fmin=5, fmax=min(fs / 2, 50), min_prominence_db=6)
    print(f"{label:60s} sharp peaks 5-50 Hz: " + (", ".join(f"{q['freq_hz']:.2f} Hz (+{q['prominence_db']:.0f} dB)" for q in pk) or "none"))

fig, axes = plt.subplots(2, 1, figsize=(11, 8))
helpers_l1.plot_spectra({k: spectra[k] for k in list(spectra)[:1]}, ax=axes[0], lines=HARMONICS + [TONE_HZ],
                        title=f"{SUBJECT} {CH} + synthetic {TONE_HZ:g} Hz marker at {SF:g} Hz, 0-250 Hz")
helpers_l1.plot_spectra(spectra, ax=axes[1], xlim=(0, 50), lines=[20, 30, 40],
                        title=f"The same channel after decimation to {FS_NEW:g} Hz (dotted: predicted aliases 20, 30, 40 Hz)")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

rng = np.random.default_rng(1)
noise = rng.standard_normal(int(60 * SF)) * 10.0      # 60 s of white noise, 10 uV RMS, at 500 Hz
raw_noise = mne.io.RawArray(noise[None] * 1e-6, mne.create_info(["noise"], SF, "eeg"), verbose=False)
noise_versions = {
    f"white noise at {SF:g} Hz": (SF, noise),
    "naive decimation": (FS_NEW, helpers_l1.naive_decimate(noise, Q)),
    "raw.resample(100)": (FS_NEW, raw_noise.copy().resample(FS_NEW, verbose=False).get_data()[0] * 1e6),
    "scipy.signal.decimate": (FS_NEW, signal.decimate(noise, Q, ftype="fir", zero_phase=True)),
}
ns = {k: helpers_l1.welch_psd(y, fs, seg_s=2.0) for k, (fs, y) in noise_versions.items()}
ref_db = 10 * np.log10(np.mean(ns[f"white noise at {SF:g} Hz"][1]))
for label, (fr, p) in ns.items():
    db = 10 * np.log10(p)
    m = (fr >= 5) & (fr <= 40)
    line = f"{label:28s} floor 5-40 Hz: {db[m].mean() - ref_db:+5.1f} dB re the original"
    if fr[-1] <= 60:
        hi = fr >= 30                                   # look for the roll-off above 30 Hz only
        first_below = lambda drop: (f"at {fr[hi][np.argmax(db[hi] < ref_db - drop)]:.1f} Hz" if np.any(db[hi] < ref_db - drop) else "not reached below 50 Hz")
        line += f"; roll-off: -3 dB {first_below(3)}, -10 dB {first_below(10)}"
    print(line)
ax = helpers_l1.plot_spectra(ns, xlim=(0, 50), title="White noise resampled three ways: anti-alias roll-off vs folded noise")
plt.show()   # render the static figure(s) of this cell inline

raw_ar, ar_path = helpers_l1.load_eegmat("Subject00", "rest", return_paths=True)
rep_ar = helpers.first_look_report(raw_ar, "ds-arithmetic")
display(HTML(helpers.report_html({k: rep_ar[k] for k in ("description", "sfreq_hz", "nyquist_hz", "n_channels", "channel_types",
                                                          "duration_s", "highpass_hz_in_header", "lowpass_hz_in_header")},
                                 "First look: ds-arithmetic Subject00_1 (background EEG before the task)")))
print("Note: this file is", f"{rep_ar['duration_s']:.0f} s long; the catalog states that only artifact-free 60-s segments are released",
      "-- TODO(confirm) whether that statement applies to the task (_2) files only.")

f_ar, psd_ar, ch_ar = helpers_l1.welch_raw(raw_ar, seg_s=4.0)
m_ar = psd_ar.mean(axis=0)
ref = 10 * np.log10(m_ar[np.argmin(np.abs(f_ar - 20))])
db_ar = 10 * np.log10(m_ar)
for drop in (3, 10, 20):
    idx = np.where((f_ar > 20) & (db_ar <= ref - drop))[0]
    print(f"ds-arithmetic channel mean: -{drop:2d} dB relative to 20 Hz first reached at {f_ar[idx[0]]:.2f} Hz")
print(f"50 Hz relative to its flanks: {helpers_l1.peak_prominence_db(f_ar, m_ar, 50.0):+.1f} dB (a hole, not a peak)")
print(f"ds-iowapd {SUBJECT} channel mean, for comparison: 60 Hz {helpers_l1.peak_prominence_db(f_pd, psd_pd, 60.0):+.1f} dB over its flanks")

ax = helpers_l1.plot_spectra({f"ds-iowapd {SUBJECT}, mean of {len(ch_pd)} channels": (f_pd, psd_pd),
                              f"ds-arithmetic Subject00 rest, mean of {len(ch_ar)} channels": (f_ar, psd_ar)},
                             lines=[30, 50], title="Two nominal 500 Hz recordings: content up to 250 Hz vs a ~30 Hz hardware ceiling with a 50 Hz notch hole")
plt.show()   # render the static figure(s) of this cell inline

# Decimating ds-arithmetic naively to 100 Hz folds nothing worth seeing: there is almost no power above 50 Hz to fold.
x_ar = raw_ar.get_data(picks="O1")[0] * 1e6
ar_versions = {f"O1 at {SF:g} Hz": helpers_l1.welch_psd(x_ar, SF, seg_s=4.0),
               "O1, naive decimation to 100 Hz": helpers_l1.welch_psd(helpers_l1.naive_decimate(x_ar, Q), FS_NEW, seg_s=4.0)}
for label, (fr, p) in ar_versions.items():
    pk = helpers_l1.line_peaks(fr, p, fmin=15, fmax=50, min_prominence_db=6)
    print(f"{label:34s} sharp peaks 15-50 Hz: " + (", ".join(f"{q['freq_hz']:.2f} Hz" for q in pk) or "none"))
above = m_ar[f_ar > 50].sum() / m_ar.sum()
print(f"fraction of the channel-mean power above 50 Hz in ds-arithmetic: {100 * above:.3f} % "
      f"(ds-iowapd {SUBJECT}: {100 * psd_pd.mean(axis=0)[f_pd > 50].sum() / psd_pd.mean(axis=0).sum():.1f} %)")
ax = helpers_l1.plot_spectra(ar_versions, xlim=(0, 50), title="ds-arithmetic O1: naive decimation to 100 Hz changes almost nothing below 50 Hz")
plt.show()   # render the static figure(s) of this cell inline

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

fr_n, p_n = spectra["naive decimation x[::5] (no filter)"]
m30 = (fr_n >= 25) & (fr_n <= 35)
observed = float(fr_n[m30][np.argmax(p_n[m30])])
print("nb-1-1-resampling -- L1.1 numbers (draft; TODO(confirm) at author review)")
print(f"Data: {DATASET} {SUBJECT} (OpenNeuro ds004584 v1.0.0, CC0), channel {CH}, {SF:g} Hz, decimated to {FS_NEW:g} Hz; "
      f"synthetic {TONE_HZ:g} Hz marker of {TONE_UV:g} uV added and labelled as such")
print(f"alias frequency of the {TONE_HZ:g} Hz component when decimating {SF:g} -> {FS_NEW:g} Hz without an anti-alias filter: "
      f"{helpers_l1.alias_frequency(TONE_HZ, FS_NEW):g} Hz  (largest peak between 25 and 35 Hz in the decimated spectrum: {observed:.2f} Hz)")
print("real components that fold in the same step: "
      + ", ".join(f"{f} -> {helpers_l1.alias_frequency(f, FS_NEW):g} Hz" for f in HARMONICS))
print(f"textbook case (L1.1 exercise): a 140 Hz tone sampled at 200 Hz appears at {helpers_l1.alias_frequency(140, 200):g} Hz")
print(f"counter-example: ds-arithmetic Subject00 is nominally {raw_ar.info['sfreq']:g} Hz but its channel-mean PSD is 20 dB below its 20 Hz level by "
      f"{f_ar[np.where((f_ar > 20) & (db_ar <= ref - 20))[0][0]]:.1f} Hz and has a 50 Hz hardware notch hole")