# 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 scipy import signal

DATASET, SUBJECT, CH = "ds-eegbci", "S001", "O1"
raw_eo = helpers.load_spine(DATASET, SUBJECT, "R01")      # eyes open
raw_ec = helpers.load_spine(DATASET, SUBJECT, "R02")      # eyes closed
SF = raw_ec.info["sfreq"]
x_ec = raw_ec.get_data(picks=CH)[0] * 1e6                 # uV
x_eo = raw_eo.get_data(picks=CH)[0] * 1e6
DUR = raw_ec.times[-1] + 1 / SF


def scatter_db(f, p, band=(20.0, 30.0)):
    """SD (dB) of the PSD in `band` around a straight line fitted in log-frequency: a noisiness measure."""
    m = (f >= band[0]) & (f <= band[1])
    db = 10 * np.log10(p[m])
    fit = np.polyval(np.polyfit(np.log10(f[m]), db, 1), np.log10(f[m]))
    return float(np.std(db - fit))


estimates = {}
for label, seg in ((f"periodogram, first 10 s", x_ec[: int(10 * SF)]), (f"periodogram, all {DUR:.0f} s", x_ec)):
    f, p = signal.periodogram(seg, fs=SF, window="hann", detrend="constant", scaling="density")
    estimates[label] = (f[1:], p[1:])
f, p = helpers_l1.welch_psd(x_ec, SF, seg_s=2.0)
estimates[f"Welch, 2-s Hann segments, 50 % overlap, all {DUR:.0f} s"] = (f, p)
for label, (f, p) in estimates.items():
    print(f"{label:52s} bins {f[1] - f[0]:.4f} Hz apart; scatter in 20-30 Hz: {scatter_db(f, p):.1f} dB")
ax = helpers_l1.plot_spectra(estimates, xlim=(0, 80), title=f"{SUBJECT} R02 (eyes closed) {CH}: periodogram vs Welch", lw=0.7)
plt.show()   # render the static figure(s) of this cell inline

print(f"{'seg_s':>6s} {'nperseg':>8s} {'noverlap':>9s} {'df (Hz)':>8s} {'segments':>9s} {'scatter 20-30 Hz':>17s}")
trade = {}
for seg_s in (0.5, 1.0, 2.0, 4.0, 8.0):
    nperseg, noverlap = int(seg_s * SF), int(seg_s * SF) // 2
    f, p = helpers_l1.welch_psd(x_ec, SF, seg_s=seg_s, overlap=0.5)
    trade[f"{seg_s:g}-s segments ({helpers_l1.n_welch_segments(len(x_ec), nperseg, noverlap)} of them)"] = (f, p)
    print(f"{seg_s:6.1f} {nperseg:8d} {noverlap:9d} {1 / seg_s:8.3f} {helpers_l1.n_welch_segments(len(x_ec), nperseg, noverlap):9d} {scatter_db(f, p):17.2f}")
print()
print("2-s segments, three overlaps:")
for ov in (0.0, 0.5, 0.75):
    nperseg = int(2 * SF)
    f, p = helpers_l1.welch_psd(x_ec, SF, seg_s=2.0, overlap=ov)
    print(f"  overlap {100 * ov:3.0f} % (noverlap {int(ov * nperseg):3d}): {helpers_l1.n_welch_segments(len(x_ec), nperseg, int(ov * nperseg)):3d} segments; scatter 20-30 Hz {scatter_db(f, p):.2f} dB")
ax = helpers_l1.plot_spectra(trade, xlim=(0, 40), title=f"{SUBJECT} R02 {CH}: Welch with five segment lengths (Hann, 50 % overlap)", lw=0.8)
plt.show()   # render the static figure(s) of this cell inline

f, p_density = helpers_l1.welch_psd(x_ec, SF, seg_s=2.0)
_, p_spectrum = signal.welch(x_ec, fs=SF, window="hann", nperseg=int(2 * SF), noverlap=int(SF), scaling="spectrum")
df = f[1] - f[0]
alpha = (f >= 8) & (f <= 12)
broad = (f >= 1) & (f <= 40)
bp_density = p_density[alpha].sum() * df
bp_spectrum = p_spectrum[alpha].sum()
w_hann = signal.get_window("hann", int(2 * SF))
enbw_bins = len(w_hann) * (w_hann ** 2).sum() / w_hann.sum() ** 2          # equivalent noise bandwidth of the window, in bins
print(f"8-12 Hz band power at {CH}, eyes closed: {bp_density:.1f} uV^2 = sum of the density over the band x df ({df:g} Hz)")
print(f"summing scaling='spectrum' over the same band gives {bp_spectrum:.1f} uV^2, {bp_spectrum / bp_density:.2f}x more: a Hann window spreads each bin over "
      f"{enbw_bins:.2f} bins (its equivalent noise bandwidth), so the sum double-counts. Integrate the density; read peaks off the spectrum.")
print(f"relative alpha power (8-12 Hz / 1-40 Hz): {100 * bp_density / (p_density[broad].sum() * df):.1f} %")
print(f"mean PSD in 8-12 Hz: {p_density[alpha].mean():.1f} uV^2/Hz; RMS of the 8-12 Hz band: {np.sqrt(bp_density):.1f} uV")
fig, axes = plt.subplots(1, 3, figsize=(14, 3.6))
for ax, (db, log_y, ttl) in zip(axes, ((False, False, "linear axes"), (False, True, "log power axis"), (True, False, "decibels"))):
    helpers_l1.plot_spectra({f"{CH} eyes closed": (f, p_density)}, ax=ax, db=db, log_y=log_y, xlim=(0, 80), title=ttl, legend=False)
fig.suptitle(f"{SUBJECT} R02 {CH}: the same Welch PSD on three scales", y=1.03)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

f, p_eo = helpers_l1.welch_psd(x_eo, SF, seg_s=2.0)
_, p_ec = helpers_l1.welch_psd(x_ec, SF, seg_s=2.0)
peak_ec = helpers_l1.alpha_peak_frequency(f, p_ec)
peak_eo = helpers_l1.alpha_peak_frequency(f, p_eo)
ratio = (p_ec[alpha].sum() / p_eo[alpha].sum())
print(f"{SUBJECT} {CH}: alpha peak (largest PSD value in 7-13 Hz) at {peak_ec:g} Hz eyes closed, {peak_eo:g} Hz eyes open; "
      f"8-12 Hz power ratio closed/open = {ratio:.1f} ({10 * np.log10(ratio):.1f} dB)")
fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))
helpers_l1.plot_spectra({"eyes open (R01)": (f, p_eo), "eyes closed (R02)": (f, p_ec)}, ax=axes[0], xlim=(0, 80), lines=[peak_ec],
                        title=f"{SUBJECT} {CH}: Welch 2-s Hann, 50 % overlap")
helpers_l1.plot_spectra({"eyes open (R01)": (f, p_eo), "eyes closed (R02)": (f, p_ec)}, ax=axes[1], db=False, xlim=(4, 16), lines=[peak_ec],
                        title="Zoom on the alpha range, linear power")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

SUBJECTS = helpers_l1.eegbci_subjects(1, 10)          # S001-S010; the catalog's exclusions are not in this range
group = {"eo": [], "ec": []}
peaks = []
for s in SUBJECTS:
    r_eo = helpers.load_spine(DATASET, s, "R01")
    r_ec = helpers.load_spine(DATASET, s, "R02")
    assert r_eo.info["sfreq"] == SF and r_ec.info["sfreq"] == SF
    fg, p1 = helpers_l1.welch_psd(r_eo.get_data(picks=CH)[0] * 1e6, SF, seg_s=2.0)
    _, p2 = helpers_l1.welch_psd(r_ec.get_data(picks=CH)[0] * 1e6, SF, seg_s=2.0)
    group["eo"].append(p1)
    group["ec"].append(p2)
    peaks.append(helpers_l1.alpha_peak_frequency(fg, p2))
group = {k: np.array(v) for k, v in group.items()}
peaks = np.array(peaks)
alpha_g = (fg >= 8) & (fg <= 12)
print(f"{'subject':8s} {'EC peak (Hz)':>12s} {'EC 8-12 Hz (uV^2)':>18s} {'EO 8-12 Hz (uV^2)':>18s} {'ratio EC/EO':>11s}")
for s, pk, p2, p1 in zip(SUBJECTS, peaks, group["ec"], group["eo"]):
    print(f"{s:8s} {pk:12.1f} {p2[alpha_g].sum() * df:18.1f} {p1[alpha_g].sum() * df:18.1f} {p2[alpha_g].sum() / p1[alpha_g].sum():11.1f}")
lin_mean = group["ec"].mean(axis=0)
log_mean = 10 ** np.log10(group["ec"]).mean(axis=0)
print(f"\nEC alpha peak frequency: mean of the {len(SUBJECTS)} individual peaks {peaks.mean():.2f} Hz (SD {peaks.std(ddof=1):.2f}), median {np.median(peaks):.2f} Hz; "
      f"peak of the linear-mean spectrum {helpers_l1.alpha_peak_frequency(fg, lin_mean):g} Hz; peak of the log-mean spectrum {helpers_l1.alpha_peak_frequency(fg, log_mean):g} Hz")
m_alpha = (fg >= 9.5) & (fg <= 10.5)
print(f"at 9.5-10.5 Hz the linear mean is {10 * np.log10(lin_mean[m_alpha].mean() / log_mean[m_alpha].mean()):.1f} dB above the log mean "
      f"(the largest single subject sits {10 * np.log10(group['ec'][:, m_alpha].mean(axis=1).max() / log_mean[m_alpha].mean()):.1f} dB above it)")

fig, axes = plt.subplots(1, 2, figsize=(13, 4), sharey=True)
for ax, cond, ttl in zip(axes, ("ec", "eo"), ("eyes closed (R02)", "eyes open (R01)")):
    spectra = {f"{s}": (fg, pp, dict(lw=0.5, color="0.6")) for s, pp in zip(SUBJECTS, group[cond])}
    spectra["linear mean"] = (fg, group[cond].mean(axis=0), dict(lw=1.8, color="tab:orange"))
    spectra["log (geometric) mean"] = (fg, 10 ** np.log10(group[cond]).mean(axis=0), dict(lw=1.8, color="tab:blue"))
    helpers_l1.plot_spectra(spectra, ax=ax, xlim=(1, 40), title=f"{len(SUBJECTS)} subjects, {CH}, {ttl}", legend=False)
    ax.legend(handles=ax.get_lines()[-2:], fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

T, N = 2.0, int(2.0 * SF)
tt = np.arange(N) / SF
PAD = 8


def amp_spectrum_db(y, window, pad=PAD):
    """Amplitude spectrum (dB re 1 uV) of y with a window, zero-padded `pad`-fold; window gain compensated."""
    w = signal.get_window(window, len(y))
    Y = np.fft.rfft(y * w, n=pad * len(y))
    fr = np.fft.rfftfreq(pad * len(y), 1 / SF)
    return fr, 20 * np.log10(np.maximum(2 * np.abs(Y) / w.sum(), 1e-9))


fig, axes = plt.subplots(1, 2, figsize=(13, 3.8), sharey=True)
for ax, f0 in zip(axes, (10.0, 10.25)):
    tone = 10.0 * np.sin(2 * np.pi * f0 * tt)                       # 10 uV amplitude
    for window in ("boxcar", "hann"):
        fr, a = amp_spectrum_db(tone, window)
        ax.plot(fr, a, lw=0.9, label=f"{window} window")
        far = a[(fr >= f0 + 2) & (fr <= f0 + 4)].max()
        print(f"{f0:5.2f} Hz tone, {window:6s}: peak {a.max():5.1f} dB re 1 uV (true 20.0), highest level 2-4 Hz away {far:6.1f} dB")
    ax.set(xlim=(4, 16), ylim=(-70, 25), xlabel="Frequency (Hz)", ylabel="Amplitude (dB re 1 uV)",
           title=f"{f0:g} Hz tone in a {T:g}-s epoch ({'on' if f0 == 10.0 else 'between'} bins, {1 / T:g} Hz apart) (dB re 1 uV)")
    ax.grid(alpha=0.3); ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

windows = ["boxcar", "hann", "hamming", ("tukey", 0.25), "blackman"]
fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))
print(f"{'window':14s} {'-3 dB main-lobe width (bins)':>28s} {'highest side lobe (dB)':>22s}")
for window in windows:
    name = window if isinstance(window, str) else f"{window[0]}({window[1]})"
    w = signal.get_window(window, N)
    W = np.abs(np.fft.rfft(w, n=64 * N))
    W_db = 20 * np.log10(np.maximum(W / W.max(), 1e-12))
    bins = np.fft.rfftfreq(64 * N, 1 / SF) * T                      # frequency in units of bins (1/T)
    width = 2 * bins[np.argmax(W_db < -3)]                          # symmetric main lobe
    pk, _ = signal.find_peaks(W_db)
    side = W_db[pk].max() if len(pk) else np.nan
    print(f"{name:14s} {width:28.2f} {side:22.1f}")
    axes[0].plot(tt, w, lw=1, label=name)
    axes[1].plot(bins, W_db, lw=0.9, label=name)
axes[0].set(xlabel="Time (s)", ylabel="Window value (a.u.)", title=f"Window shapes over a {T:g}-s epoch (a.u.)")
axes[1].set(xlim=(0, 8), ylim=(-100, 3), xlabel="Frequency (bins of 1/T)", ylabel="Gain (dB re peak)", title="Window spectra: main lobe and side lobes (dB)")
for ax in axes:
    ax.grid(alpha=0.3); ax.legend(fontsize=8)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

F1, F2 = 10.0, 10.6
fig, axes = plt.subplots(1, 3, figsize=(14, 3.6), sharey=True)
for ax, (Tn, pad) in zip(axes, ((1.0, 1), (1.0, 16), (4.0, 1))):
    n = int(Tn * SF)
    tn = np.arange(n) / SF
    two = 10 * np.sin(2 * np.pi * F1 * tn) + 10 * np.sin(2 * np.pi * F2 * tn)
    w = signal.get_window("hann", n)
    Y = np.fft.rfft(two * w, n=pad * n)
    fr = np.fft.rfftfreq(pad * n, 1 / SF)
    a = 2 * np.abs(Y) / w.sum()
    ax.plot(fr, a, "k", lw=0.9)
    ax.plot(fr, a, ".", color="tab:orange", ms=3)
    pk, _ = signal.find_peaks(a[(fr > 7) & (fr < 14)], height=a.max() * 0.3)
    print(f"T = {Tn:g} s, {pad:2d}x zero-padding: {len(fr)} points, spacing {fr[1]:.4f} Hz, distinct peaks between 7 and 14 Hz: {len(pk)}")
    ax.set(xlim=(7, 14), xlabel="Frequency (Hz)", title=f"T = {Tn:g} s, {pad}x zero-padding: {len(pk)} peak{'s' if len(pk) != 1 else ''}")
    ax.grid(alpha=0.3)
axes[0].set_ylabel("Amplitude (uV)")
fig.suptitle(f"Two tones at {F1:g} and {F2:g} Hz (10 uV each), Hann window: padding smooths, epoch length resolves (uV)", y=1.03)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

three = {}
for window in ("boxcar", "hann", "hamming"):
    fw, pw = helpers_l1.welch_psd(x_ec, SF, seg_s=2.0, window=window)
    three[f"{'rectangular' if window == 'boxcar' else window} window"] = (fw, pw)
    m1 = (fw >= 14) & (fw <= 16)
    m2 = (fw >= 62) & (fw <= 65)
    print(f"{window:8s}: PSD 14-16 Hz {10 * np.log10(pw[m1].mean()):5.1f} dB, 62-65 Hz {10 * np.log10(pw[m2].mean()):5.1f} dB re 1 uV^2/Hz")
fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))
helpers_l1.plot_spectra(three, ax=axes[0], xlim=(0, 80), title=f"{SUBJECT} R02 {CH}: Welch 2-s segments, three windows")
helpers_l1.plot_spectra(three, ax=axes[1], xlim=(50, 70), title="Zoom on the 60 Hz line: the rectangular skirt")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

FS_EX, DUR_EX, DF_EX, MIN_SEG = 160.0, 60.0, 0.5, 20
n_ex = int(DUR_EX * FS_EX)
nperseg = int(round(FS_EX / DF_EX))
print("nb-1-3-psd -- L1.3 / L1.4 exercise numbers (draft; TODO(confirm) at author review)")
print(f"resolution {DF_EX:g} Hz at fs = {FS_EX:g} Hz -> nperseg = fs / df = {nperseg} samples ({nperseg / FS_EX:g} s)")
for noverlap in (0, nperseg // 2, 3 * nperseg // 4):
    n_seg = helpers_l1.n_welch_segments(n_ex, nperseg, noverlap)
    print(f"  noverlap = {noverlap:3d} ({100 * noverlap / nperseg:2.0f} %): {n_seg:3d} segments in {DUR_EX:g} s -> {'meets' if n_seg >= MIN_SEG else 'fails'} the >= {MIN_SEG} requirement")
print(f"Answer pair (the conventional choice): nperseg = {nperseg}, noverlap = {nperseg // 2} "
      f"({helpers_l1.n_welch_segments(n_ex, nperseg, nperseg // 2)} segments); the minimal pair nperseg = {nperseg}, noverlap = 0 gives "
      f"{helpers_l1.n_welch_segments(n_ex, nperseg, 0)} segments and also satisfies the constraint")
# check the count against SciPy by asking it for the individual segment periodograms
_, _, stft = signal.stft(np.zeros(n_ex), fs=FS_EX, window="hann", nperseg=nperseg, noverlap=nperseg // 2, boundary=None, padded=False)
print(f"SciPy check: stft with the same parameters yields {stft.shape[1]} segments")
print(f"Group-mean alpha peak frequency, eyes closed, {CH}, {len(SUBJECTS)} subjects ({SUBJECTS[0]}-{SUBJECTS[-1]}), Welch 2-s Hann 50 %, largest PSD in 7-13 Hz: "
      f"mean {peaks.mean():.2f} Hz, median {np.median(peaks):.2f} Hz, SD {peaks.std(ddof=1):.2f} Hz; peak of the log-mean spectrum {helpers_l1.alpha_peak_frequency(fg, log_mean):g} Hz")
print(f"Individual peaks: " + ", ".join(f"{s} {pk:g}" for s, pk in zip(SUBJECTS, peaks)))