# 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", "moabb", "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", "moabb==1.7.2", "pooch>=1.8"]
    subprocess.check_call(_cmd)

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

# 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

mne.viz.set_browser_backend("matplotlib", verbose=False)
mne.set_log_level("WARNING")
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers imported from notebooks/_shared; "
      "downloads go to MNE's default data directory unless EEG_COURSE_DATA is set")

from scipy import signal

SF = 160.0      # ds-eegbci sampling rate (catalog); designs are specified in Hz
L_FREQ = 1.0    # a 1 Hz high-pass, a common choice before an ICA fit (L2.4)

h_fir = mne.filter.create_filter(None, SF, l_freq=L_FREQ, h_freq=None, verbose=False)
sos_iir = signal.butter(4, L_FREQ, btype="highpass", fs=SF, output="sos")
print(f"FIR: {len(h_fir)} taps = {len(h_fir) / SF:.2f} s of impulse response; "
      f"IIR: {sos_iir.shape[0]} second-order sections ({2 * sos_iir.shape[0]} poles)")
fig = helpers.plot_filter_response({f"FIR, MNE defaults ({len(h_fir)} taps)": {"h": h_fir},
                                    "IIR Butterworth, order 4": {"sos": sos_iir}},
                                   SF, title=f"High-pass {L_FREQ:g} Hz at {SF:g} Hz", fmax=6, impulse_t_s=4)

w, H_fir = signal.freqz(h_fir, worN=8192, fs=SF)
_, H_iir = signal.sosfreqz(sos_iir, worN=8192, fs=SF)


def freq_at_gain(w, H, db):
    """Lowest frequency at which the gain reaches `db` dB (a high-pass rises from the left)."""
    gain = 20 * np.log10(np.maximum(np.abs(H), 1e-12))
    return float(w[np.argmax(gain >= db)])


print(f"FIR: -6 dB at {freq_at_gain(w, H_fir, -6):.2f} Hz, -3 dB at {freq_at_gain(w, H_fir, -3):.2f} Hz "
      "(MNE reports the -6 dB point as the cutoff and l_freq as the passband edge)")
print(f"IIR: -3 dB at {freq_at_gain(w, H_iir, -3):.2f} Hz (the Butterworth convention); "
      "applied forward and backward the response is squared, so that point becomes -6 dB")

designs = {}
for tb in (0.1, 0.25, 0.5, 1.0):
    h = mne.filter.create_filter(None, SF, l_freq=L_FREQ, h_freq=None, l_trans_bandwidth=tb, verbose=False)
    designs[f"transition {tb:g} Hz: {len(h)} taps ({len(h) / SF:.1f} s)"] = {"h": h}
    print(f"high-pass {L_FREQ:g} Hz, transition bandwidth {tb:5.2f} Hz -> {len(h):5d} taps = {len(h) / SF:5.1f} s")
fig = helpers.plot_filter_response(designs, SF, title="FIR high-pass 1 Hz: transition bandwidth", fmax=3, impulse_t_s=40)
orders = {f"Butterworth order {n}": {"sos": signal.butter(n, L_FREQ, btype="highpass", fs=SF, output="sos")}
          for n in (2, 4, 8)}
fig = helpers.plot_filter_response(orders, SF, title="IIR high-pass 1 Hz: order", fmax=6, impulse_t_s=4)
plt.show()   # render the static figure(s) of this cell inline

raw_pop = helpers.load_spine("ds-eegbci", "S011", "R01")   # one eyes-open run; electrode-pop candidate on O2
x = raw_pop.get_data(picks="O2")[0] * 1e6
t = raw_pop.times
d = raw_pop.get_data() * 1e6
ci, si = np.unravel_index(int(np.abs(np.diff(d, axis=1)).argmax()), (d.shape[0], d.shape[1] - 1))
t_pop = (si + 1) / SF
print(f"pop candidate: {raw_pop.ch_names[ci]} at {t_pop:.2f} s (largest sample-to-sample jump in the run)")

w_hz = np.linspace(0.05, 10, 2000)          # start above 0 Hz: a high-pass has a zero there
with warnings.catch_warnings():
    warnings.simplefilter("ignore")           # SciPy warns about that zero even when it is not evaluated
    w, gd_fir = signal.group_delay((h_fir, [1.0]), w=w_hz, fs=SF)
    _, gd_iir = signal.group_delay(signal.sos2tf(sos_iir), w=w_hz, fs=SF)
fig, ax = plt.subplots(figsize=(8, 3.5))
ax.plot(w, gd_fir / SF * 1000, label=f"FIR ({len(h_fir)} taps): constant {(len(h_fir) - 1) / 2 / SF * 1000:.0f} ms")
ax.plot(w, gd_iir / SF * 1000, label="IIR Butterworth order 4: frequency-dependent")
ax.set(xlim=(0, 10), ylim=(0, 2000), xlabel="Frequency (Hz)", ylabel="Group delay (ms)",
       title="Causal application: group delay of the two 1 Hz high-passes (ms)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)
plt.show()   # render the static figure(s) of this cell inline

outputs = {
    "original (uV)": x,
    "FIR causal (lfilter): the whole trace arrives half a filter later": signal.lfilter(h_fir, [1.0], x),
    "FIR zero-phase (MNE phase='zero'): delay compensated, ringing on both sides": mne.filter.filter_data(x, SF, L_FREQ, None, phase="zero", verbose=False),
    "IIR causal (sosfilt): frequency-dependent delay distorts the step": signal.sosfilt(sos_iir, x),
    "IIR zero-phase (sosfiltfilt): forward-backward, symmetric": signal.sosfiltfilt(sos_iir, x),
}
fig, axes = plt.subplots(len(outputs), 1, figsize=(11, 11), sharex=True)
for ax, (label, y) in zip(axes, outputs.items()):
    ax.plot(t, y, "k", lw=0.7)
    ax.axvspan(t_pop - 0.05, t_pop + 0.9, color="tab:orange", alpha=0.2)
    ax.axvline(t_pop, color="tab:orange", lw=0.8)
    ax.set(ylabel="uV", title=label, xlim=(t_pop - 2, t_pop + 4))
    ax.grid(alpha=0.3)
axes[-1].set_xlabel("Time (s)")
fig.suptitle(f"S011 R01, O2 around the pop at {t_pop:.2f} s: one 1 Hz high-pass applied four ways (uV)", y=1.0)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

h_narrow = mne.filter.create_filter(None, SF, l_freq=L_FREQ, h_freq=None, l_trans_bandwidth=0.1, verbose=False)
sos8 = signal.butter(8, L_FREQ, btype="highpass", fs=SF, output="sos")
ring = {
    f"FIR, transition 1 Hz ({len(h_fir)} taps), zero-phase": mne.filter.filter_data(x, SF, L_FREQ, None, phase="zero", verbose=False),
    f"FIR, transition 0.1 Hz ({len(h_narrow)} taps), zero-phase": mne.filter.filter_data(x, SF, L_FREQ, None, l_trans_bandwidth=0.1, phase="zero", verbose=False),
    "IIR Butterworth order 8, forward-backward": signal.sosfiltfilt(sos8, x),
}
fig, axes = plt.subplots(len(ring), 1, figsize=(12, 8), sharex=True, sharey=True)
for ax, (label, y) in zip(axes, ring.items()):
    ax.plot(t, x - np.median(x), color="0.7", lw=0.7, label="original (median removed)")
    ax.plot(t, y, "k", lw=0.8, label="filtered")
    ax.axvline(t_pop, color="tab:orange", lw=0.8)
    ax.set(ylabel="uV", title=label, xlim=(t_pop - 4, t_pop + 4.5))
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8, loc="upper right")
axes[-1].set_xlabel("Time (s)")
fig.suptitle("Ringing around the pop grows with the length of the impulse response (uV)", y=1.0)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    h = mne.filter.create_filter(raw_pop.get_data(picks="eeg"), SF, l_freq=0.01, h_freq=None, verbose=False)
print(f"0.01 Hz high-pass with MNE defaults at {SF:g} Hz: {len(h)} taps = {len(h) / SF:.0f} s, "
      f"on a {raw_pop.times[-1]:.0f}-s run")
for c in caught:
    print("MNE warns:", str(c.message))

from IPython.display import HTML, display

raw_p300 = helpers.load_spine("ds-brain-invaders", 1)
SF_P = raw_p300.info["sfreq"]
rep = helpers.first_look_report(raw_p300, "ds-brain-invaders")
display(HTML(helpers.report_html({k: rep[k] for k in ("description", "sfreq_hz", "n_channels", "channel_types", "eeg_channels",
                                                       "duration_s", "highpass_hz_in_header", "lowpass_hz_in_header",
                                                       "montage_attached", "annotation_counts")},
                                 "First look: bi2014a subject 1")))
EVENT_ID = helpers.BI2014A_EVENT_ID
events = mne.find_events(raw_p300, stim_channel="STI 014", shortest_event=1, verbose=False)
n_target = int((events[:, 2] == EVENT_ID["Target"]).sum())
n_nontarget = int((events[:, 2] == EVENT_ID["NonTarget"]).sum())
isi = np.diff(events[:, 0]) / SF_P
print(f"{len(events)} flashes: {n_target} Target, {n_nontarget} NonTarget; first at {events[0, 0] / SF_P:.1f} s, "
      f"last at {events[-1, 0] / SF_P:.1f} s of {raw_p300.times[-1]:.1f} s; median interval between flashes {np.median(isi):.3f} s")

fig, ax = plt.subplots(figsize=(10, 4))
helpers.plot_psd(raw_p300, ["Fp1", "P3", "Pz", "P4", "O1"], fmin=1, fmax=120, n_fft_s=4, ax=ax,
                 title="bi2014a subject 1, dry electrodes: PSD over the whole recording (dB re 1 uV^2/Hz)")
spec = raw_p300.compute_psd(method="welch", picks=["Pz"], fmin=1, fmax=120, n_fft=int(4 * SF_P), verbose=False)
psd, freqs = spec.get_data(return_freqs=True)
at = lambda f: float(psd[0, np.argmin(np.abs(freqs - f))] * 1e12)
print(f"Pz PSD: {at(10):.0f} uV^2/Hz at 10 Hz, {at(50):.0f} uV^2/Hz at 50 Hz (x{at(50) / at(10):.0f}); "
      "the mains line dominates the raw trace on these dry electrodes")

LP_HZ = 30.0
TMIN, TMAX, BASELINE = -0.2, 0.8, (-0.2, 0.0)
CH, WINDOW = "Pz", (0.3, 0.5)
PP_MAX_UV = 150.0

raw_lp = raw_p300.copy().filter(l_freq=None, h_freq=LP_HZ, picks="eeg", verbose=False)     # common low-pass, MNE default FIR
sel_raw = raw_lp.copy().filter(l_freq=1.0, h_freq=None, picks="eeg", verbose=False)         # selection copy only
sel_epochs = mne.Epochs(sel_raw, events, EVENT_ID, tmin=TMIN, tmax=TMAX, baseline=BASELINE, picks="eeg",
                        preload=True, verbose=False)
pp_ch = np.ptp(sel_epochs.get_data(picks=[CH]), axis=2)[:, 0] * 1e6
keep = np.where(pp_ch <= PP_MAX_UV)[0]
n_keep_t = int((sel_epochs.events[keep, 2] == EVENT_ID["Target"]).sum())
n_keep_n = int((sel_epochs.events[keep, 2] == EVENT_ID["NonTarget"]).sum())
print(f"trial selection on the 1-{LP_HZ:g} Hz copy (peak-to-peak at {CH} <= {PP_MAX_UV:g} uV): "
      f"{len(keep)} of {len(sel_epochs)} trials kept ({n_keep_t} Target, {n_keep_n} NonTarget); "
      f"{CH} peak-to-peak median {np.median(pp_ch):.0f} uV, 99th percentile {np.percentile(pp_ch, 99):.0f} uV")

CUTOFFS = [0.01, 0.1, 0.5, 1.0]
results = {}
for lf in CUTOFFS:
    h = mne.filter.create_filter(None, SF_P, l_freq=lf, h_freq=None, verbose=False)
    half = (len(h) - 1) / 2 / SF_P
    r = raw_lp.copy().filter(l_freq=lf, h_freq=None, picks="eeg", verbose=False)     # MNE defaults
    ep = mne.Epochs(r, events, EVENT_ID, tmin=TMIN, tmax=TMAX, baseline=BASELINE, picks="eeg", preload=True, verbose=False)
    assert len(ep) == len(sel_epochs)
    ep = ep[keep]
    ev_t, ev_n = ep["Target"].average(), ep["NonTarget"].average()
    diff = mne.combine_evoked([ev_t, ev_n], weights=[1, -1])
    onsets = events[keep, 0] / SF_P
    n_edge = int(((onsets + TMIN) < half).sum() + ((onsets + TMAX) > raw_p300.times[-1] - half).sum())
    results[lf] = dict(taps=len(h), half_s=half, n_edge=n_edge, target=ev_t, nontarget=ev_n, diff=diff,
                       amp={ch: helpers.erp_mean_amplitude(diff, ch, WINDOW) for ch in diff.ch_names})
    print(f"high-pass {lf:5.2f} Hz: {len(h):6d} taps = {2 * half:6.1f} s ({100 * 2 * half / raw_p300.times[-1]:4.1f} % of the "
          f"recording); kept trials within half a filter length of an edge: {n_edge:4d}; "
          f"{CH} T-N mean {WINDOW[0] * 1000:.0f}-{WINDOW[1] * 1000:.0f} ms = {results[lf]['amp'][CH]:+.2f} uV")

fig, ax = plt.subplots(figsize=(10, 4.5))
for lf in CUTOFFS:
    dw = results[lf]["diff"]
    ax.plot(dw.times * 1000, dw.data[dw.ch_names.index(CH)] * 1e6, lw=1.2, label=f"high-pass {lf:g} Hz")
ax.axvspan(WINDOW[0] * 1000, WINDOW[1] * 1000, color="tab:orange", alpha=0.2, label="a-priori window")
ax.axhline(0, color="gray", lw=0.6)
ax.axvline(0, color="gray", lw=0.6)
ax.set(xlabel="Time from flash (ms)", ylabel="Amplitude (uV)",
       title=f"Target minus NonTarget at {CH} for four high-pass cutoffs (uV, positive up; common {LP_HZ:g} Hz low-pass)")
ax.grid(alpha=0.3)
ax.legend(fontsize=8)

fig, axes = plt.subplots(2, 2, figsize=(12, 7), sharex=True, sharey=True)
for ax, lf in zip(axes.ravel(), CUTOFFS):
    for key, style in (("target", "k-"), ("nontarget", "k--")):
        e = results[lf][key]
        ax.plot(e.times * 1000, e.data[e.ch_names.index(CH)] * 1e6, style, lw=1, label=f"{key} (n = {e.nave})")
    ax.axvspan(WINDOW[0] * 1000, WINDOW[1] * 1000, color="tab:orange", alpha=0.2)
    ax.axhline(0, color="gray", lw=0.6)
    ax.set(title=f"high-pass {lf:g} Hz: {CH} (uV)")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8)
for ax in axes[1]:
    ax.set_xlabel("Time from flash (ms)")
for ax in axes[:, 0]:
    ax.set_ylabel("Amplitude (uV)")
fig.suptitle(f"Target and NonTarget averages at {CH} (uV, positive up)")
fig.tight_layout()

fig, axes = plt.subplots(1, 2, figsize=(9, 4))
for ax, lf in zip(axes, (0.1, 1.0)):
    vals16 = np.array([results[lf]["amp"][ch] for ch in results[lf]["diff"].ch_names])
    helpers.plot_topomap_values(raw_p300, vals16, unit="uV", ax=ax, vlim=(-8, 8),
                                title=f"T-N mean {WINDOW[0] * 1000:.0f}-{WINDOW[1] * 1000:.0f} ms, high-pass {lf:g} Hz")
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

vals = [results[lf]["amp"][CH] for lf in CUTOFFS]
if all(b < a for a, b in zip(vals, vals[1:])):
    trend = "decreases monotonically"
elif vals[-1] < vals[0]:
    trend = "decreases overall (not monotonically)"
else:
    trend = "does not decrease"
win_ms = f"{WINDOW[0] * 1000:.0f}-{WINDOW[1] * 1000:.0f} ms"

print("nb-1-5-filters -- L1.5 exercise numbers (draft; TODO(confirm) at author review; ERP CORE is the intended source, spec 13.22)")
print(f"Data: ds-brain-invaders bi2014a subject 1 (Zenodo DOI 10.5281/zenodo.3266223; CC BY 4.0); {len(keep)} trials "
      f"({n_keep_t} Target, {n_keep_n} NonTarget), the same trials for every cutoff")
print(f"Pipeline: MNE default FIR high-pass (firwin, Hamming, zero-phase, auto length) at the stated cutoff + common "
      f"{LP_HZ:g} Hz low-pass; epochs {TMIN} to {TMAX} s, baseline {BASELINE[0]} to {BASELINE[1]} s; recording reference kept")
print(f"Measure: mean amplitude of the Target minus NonTarget difference wave, {win_ms} (a-priori window), channel {CH}")
for lf, v in zip(CUTOFFS, vals):
    print(f"  high-pass {lf:5.2f} Hz: P3 mean amplitude = {v:+.2f} uV   ({CH}, {win_ms}, target minus non-target)")
print(f"Trend: the {CH} P3 mean amplitude {trend} from {vals[0]:+.2f} uV at 0.01 Hz to {vals[-1]:+.2f} uV at 1 Hz "
      "as the high-pass cutoff rises.")
print("Supplementary, same window and trials (not the exercise numbers): "
      + " | ".join(f"{ch}: " + ", ".join(f"{results[lf]['amp'][ch]:+.2f}" for lf in CUTOFFS) + " uV" for ch in ("P3", "P4")))
print("Exercise ids in L1.5: ex-1-5-p3-hp-0-01, ex-1-5-p3-hp-0-1, ex-1-5-p3-hp-0-5, ex-1-5-p3-hp-1 (tolerance TODO(confirm)).")