# 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/)")

FS_DEMO, T_DEMO = 160.0, 2.0
N_DEMO = int(FS_DEMO * T_DEMO)
t_demo = np.arange(N_DEMO) / FS_DEMO
components = [(10.0, 20.0, 0.0), (4.0, 8.0, np.pi / 2), (25.0, 3.0, -np.pi / 4)]     # (f Hz, A uV, phi rad)
parts = [A * np.sin(2 * np.pi * f * t_demo + phi) for f, A, phi in components]
y_demo = np.sum(parts, axis=0)

X_demo = np.fft.rfft(y_demo)
f_demo = np.fft.rfftfreq(N_DEMO, 1 / FS_DEMO)
amp_demo = 2 * np.abs(X_demo) / N_DEMO          # amplitude of each sinusoidal component (uV); the 0 Hz and Nyquist bins count once
amp_demo[0] /= 2
amp_demo[-1] /= 2
print(f"N = {N_DEMO} samples, {len(f_demo)} bins from 0 to {f_demo[-1]:g} Hz, spacing {f_demo[1]:g} Hz")
for f, A, phi in components:
    k = int(round(f / f_demo[1]))
    phase = np.angle(X_demo[k]) + np.pi / 2       # sin = cos shifted by -pi/2
    print(f"  put in: {f:5.1f} Hz, {A:5.1f} uV, phase {phi:+.3f} rad  ->  bin {k:2d}: {f_demo[k]:5.1f} Hz, {amp_demo[k]:5.2f} uV, phase {phase:+.3f} rad")

fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))
for (f, A, phi), part in zip(components, parts):
    axes[0].plot(t_demo, part, lw=0.8, label=f"{f:g} Hz, {A:g} uV, phi = {phi:+.2f}")
axes[0].plot(t_demo, y_demo, "k", lw=1.2, label="sum")
axes[0].set(xlim=(0, 0.5), xlabel="Time (s)", ylabel="Amplitude (uV)", title="Three sinusoids and their sum, first 0.5 s (uV)")
axes[0].grid(alpha=0.3); axes[0].legend(fontsize=7)
axes[1].stem(f_demo, amp_demo, basefmt=" ", markerfmt="k.", linefmt="k-")
axes[1].set(xlim=(0, 40), xlabel="Frequency (Hz)", ylabel="Amplitude (uV)", title="Amplitude spectrum of the sum: one bin per component (uV)")
axes[1].grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

DATASET, SUBJECT, RUN, CH = "ds-eegbci", "S001", "R02", "O1"
raw = helpers.load_spine(DATASET, SUBJECT, RUN)              # eyes closed, ~1 min
SF = raw.info["sfreq"]
x_all = raw.get_data(picks=CH)[0] * 1e6                     # uV
T = 2.0
N = int(T * SF)


def alpha_power(seg):
    """8-12 Hz power of a Hann-windowed segment (relative units; used only to pick the epoch)."""
    Xs = np.fft.rfft(seg * np.hanning(len(seg)))
    fr = np.fft.rfftfreq(len(seg), 1 / SF)
    return float((np.abs(Xs[(fr >= 8) & (fr <= 12)]) ** 2).sum())


starts = np.arange(0.0, raw.times[-1] - 8.0, 1.0)            # leave room for the 8-s window of section 5
T0 = float(starts[np.argmax([alpha_power(x_all[int(s * SF):int(s * SF) + N]) for s in starts])])
x = x_all[int(T0 * SF):int(T0 * SF) + N]
t = np.arange(N) / SF
X = np.fft.rfft(x)
freqs = np.fft.rfftfreq(N, 1 / SF)
DF = freqs[1] - freqs[0]
amp = 2 * np.abs(X) / N
amp[0] /= 2
amp[-1] /= 2
print(f"{SUBJECT} {RUN} {CH}: epoch {T0:.0f}-{T0 + T:.0f} s (rule: the 2-s window with the most 8-12 Hz power), N = {N} samples at {SF:g} Hz")
print(f"rfft: {len(freqs)} bins, bin k <-> {DF:g} * k Hz, from 0 to {freqs[-1]:g} Hz (Nyquist); resolution fs / N = 1 / T = {DF:g} Hz")
order = np.argsort(amp[1:])[::-1] + 1                          # bins sorted by amplitude, DC excluded
print("largest components: " + ", ".join(f"{freqs[k]:g} Hz {amp[k]:.1f} uV (bin {k})" for k in order[:6]))

fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))
axes[0].plot(t + T0, x, "k", lw=0.8)
axes[0].set(xlabel="Time (s)", ylabel="Amplitude (uV)", title=f"{SUBJECT} {RUN} {CH}: the 2-s epoch (uV)")
axes[0].grid(alpha=0.3)
axes[1].plot(freqs, amp, "k", lw=0.8)
axes[1].plot(freqs, amp, ".", color="tab:orange", ms=3)
axes[1].set(xlim=(0, 80), xlabel="Frequency (Hz)", ylabel="Amplitude (uV)", title=f"Amplitude spectrum, one dot per bin ({DF:g} Hz apart) (uV)")
axes[1].grid(alpha=0.3)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

power = amp ** 2 / 2
power[0] = amp[0] ** 2                                  # a constant contributes its square
power[-1] = amp[-1] ** 2                                # the Nyquist bin counts once as well
ms_time = float(np.mean(x ** 2))
ms_freq = float(power.sum())
print(f"Parseval: mean square of the epoch = {ms_time:.2f} uV^2 in the time domain, {ms_freq:.2f} uV^2 summed over the bins")
k_a = order[0]
k_g = int(np.argmin(np.abs(freqs - 35.0)))
print(f"alpha peak bin {freqs[k_a]:g} Hz: {amp[k_a]:.1f} uV, {power[k_a]:.0f} uV^2, {10 * np.log10(power[k_a]):.1f} dB re 1 uV^2")
print(f"35 Hz bin:              {amp[k_g]:.2f} uV, {power[k_g]:.2f} uV^2, {10 * np.log10(power[k_g]):.1f} dB re 1 uV^2 "
      f"-> ratio {amp[k_a] / amp[k_g]:.0f}x in amplitude, {power[k_a] / power[k_g]:.0f}x in power, {10 * np.log10(power[k_a] / power[k_g]):.0f} dB")

fig, axes = plt.subplots(1, 3, figsize=(14, 3.6))
axes[0].plot(freqs, amp, "k", lw=0.8); axes[0].set(ylabel="Amplitude (uV)", title="Amplitude (uV)")
axes[1].plot(freqs, power, "k", lw=0.8); axes[1].set(ylabel="Power (uV^2)", title="Power (uV^2)")
axes[2].plot(freqs[1:], 10 * np.log10(power[1:]), "k", lw=0.8); axes[2].set(ylabel="Power (dB re 1 uV^2)", title="Power in decibels (dB re 1 uV^2)")
for ax in axes:
    ax.set(xlim=(0, 80), xlabel="Frequency (Hz)")
    ax.grid(alpha=0.3)
fig.suptitle(f"{SUBJECT} {RUN} {CH}, epoch {T0:.0f}-{T0 + T:.0f} s: one DFT, three scales", y=1.02)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

ks = [1, 2, 3, 5, 10, 20, 50, len(order)]
recon = {}
print(f"{'k':>4s} {'residual RMS (uV)':>18s} {'variance explained':>19s}")
for k in ks:
    keep = np.zeros_like(X)
    keep[0] = X[0]
    keep[order[:k]] = X[order[:k]]
    y = np.fft.irfft(keep, n=N)
    recon[k] = y
    resid = x - y
    print(f"{k:4d} {np.sqrt(np.mean(resid ** 2)):18.2f} {100 * (1 - np.var(resid) / np.var(x)):18.1f} %")

fig, axes = plt.subplots(4, 1, figsize=(11, 9), sharex=True, sharey=True)
for ax, k in zip(axes, (1, 5, 20, len(order))):
    ax.plot(t, x, color="0.7", lw=0.8, label="epoch")
    ax.plot(t, recon[k], "k", lw=0.9, label=f"k = {k} components")
    ax.set(ylabel="uV", title=f"k = {k}: residual RMS {np.sqrt(np.mean((x - recon[k]) ** 2)):.1f} uV")
    ax.grid(alpha=0.3); ax.legend(fontsize=8, loc="upper right")
axes[-1].set_xlabel("Time (s)")
fig.suptitle(f"{SUBJECT} {RUN} {CH}: reconstruction from the k largest DFT components (uV)", y=1.0)
fig.tight_layout()

fig, ax = plt.subplots(figsize=(6, 3.2))
ax.semilogx(ks, [np.sqrt(np.mean((x - recon[k]) ** 2)) for k in ks], "ko-")
ax.set(xlabel="k (components kept) [log]", ylabel="Residual RMS (uV)", title="Residual RMS versus number of components kept (uV)")
ax.grid(alpha=0.3, which="both")
plt.show()   # render the static figure(s) of this cell inline

lengths = [0.5, 1.0, 2.0, 4.0, 8.0]
fig, axes = plt.subplots(1, len(lengths), figsize=(15, 3.4), sharey=True)
for ax, Tn in zip(axes, lengths):
    n = int(Tn * SF)
    seg = x_all[int(T0 * SF):int(T0 * SF) + n]
    fr = np.fft.rfftfreq(n, 1 / SF)
    a = 2 * np.abs(np.fft.rfft(seg)) / n
    ax.plot(fr, a, "k", lw=0.8)
    ax.plot(fr, a, ".", color="tab:orange", ms=3)
    ax.set(xlim=(0, 30), title=f"T = {Tn:g} s: bins {fr[1]:g} Hz apart", xlabel="Frequency (Hz)")
    ax.grid(alpha=0.3)
    print(f"T = {Tn:4.1f} s: N = {n:5d}, resolution 1/T = {1 / Tn:.3f} Hz = fs/N = {SF / n:.3f} Hz, {len(fr)} bins")
axes[0].set_ylabel("Amplitude (uV)")
fig.suptitle(f"{SUBJECT} {RUN} {CH} from {T0:.0f} s: amplitude spectra of five epoch lengths (uV)", y=1.02)
fig.tight_layout()

fig, axes = plt.subplots(1, 3, figsize=(13, 3.2), sharey=True)
for ax, Tn in zip(axes, (0.5, 1.0, 2.0)):
    n = int(Tn * SF)
    tt = np.arange(n) / SF
    two = 10 * np.sin(2 * np.pi * 10 * tt) + 10 * np.sin(2 * np.pi * 11 * tt)
    fr = np.fft.rfftfreq(n, 1 / SF)
    a = 2 * np.abs(np.fft.rfft(two)) / n
    ax.stem(fr, a, basefmt=" ", markerfmt="k.", linefmt="k-")
    ax.set(xlim=(5, 16), title=f"10 + 11 Hz, T = {Tn:g} s ({fr[1]:g} Hz bins)", xlabel="Frequency (Hz)")
    ax.grid(alpha=0.3)
axes[0].set_ylabel("Amplitude (uV)")
fig.suptitle("Two synthetic tones 1 Hz apart: resolved only when the epoch is at least 1 s (uV)", y=1.02)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print("nb-1-2-fft -- L1.2 exercise numbers (draft; TODO(confirm) at author review)")
for fs in (160.0, 256.0):
    n = int(2.0 * fs)
    fr = np.fft.rfftfreq(n, 1 / fs)
    print(f"2-s epoch at fs = {fs:g} Hz: N = {n} samples, {len(fr)} rfft bins, frequency resolution fs / N = {fs / n:g} Hz (= 1 / T)")
fs, T_ex = 256.0, 2.0
n = int(T_ex * fs)
fr = np.fft.rfftfreq(n, 1 / fs)
k10 = int(np.argmin(np.abs(fr - 10.0)))
print(f"bin index of 10 Hz at fs = {fs:g} Hz for a {T_ex:g}-s epoch: k = {k10}  (f_k = k * fs / N = {fr[k10]:g} Hz; exact: {bool(fr[k10] == 10.0)})")
print(f"Exercise answers: frequency resolution of a 2-s epoch = {1 / T_ex:g} Hz; the 10 Hz bin at fs = 256 Hz, T = 2 s is k = {k10}")
print(f"Also from this notebook: {SUBJECT} {RUN} {CH} epoch {T0:.0f}-{T0 + T:.0f} s, largest component {freqs[order[0]]:g} Hz at {amp[order[0]]:.1f} uV; "
      f"k = 10 components explain {100 * (1 - np.var(x - recon[10]) / np.var(x)):.1f} % of the variance")