"""
Yb3+ spectroscopy and quantum-control testbed monitor
=====================================================

Purpose
-------
This program acquires two analog photodiode channels, visualises time-domain
signals and spectra, estimates noise/SNR, compares microwave ON/OFF periods,
and records candidate spectral features for later analysis.

Scientific scope
----------------
The software does NOT by itself demonstrate a qubit, ODMR, Rabi oscillations,
Ramsey fringes, T1 or T2. It labels repeatable signal changes as *candidate
features* until they are supported by calibration, controls and pulse-timing
measurements.

Important hardware note
-----------------------
The keyboard switches below annotate the experimental state. They do not turn
physical lasers or microwave equipment on/off unless explicit DAQ digital-output
or instrument-control code is added and validated separately.
"""

from __future__ import annotations

import csv
import json
import math
import sys
import time
from collections import deque
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

try:
    import nidaqmx
    from nidaqmx.constants import AcquisitionType, TerminalConfiguration
except ImportError:
    nidaqmx = None
    AcquisitionType = None
    TerminalConfiguration = None


# =============================================================================
# CONFIGURATION
# =============================================================================

@dataclass(frozen=True)
class Config:
    # NI-DAQ configuration
    device_name: str = "Dev1"
    pd1_channel: str = "ai0"
    pd2_channel: str = "ai4"
    terminal_mode: str = "RSE"  # RSE, NRSE or DIFF
    voltage_min: float = -5.0
    voltage_max: float = 5.0

    # Sampling
    # 50 kS/s gives a Nyquist frequency of 25 kHz, sufficient for the selected
    # 0.5-20 kHz exploratory band. The actual DAQ must support this rate for two
    # simultaneous channels.
    sample_rate_hz: float = 50_000.0
    samples_per_read: int = 10_000  # 0.2 s/block, 5 Hz spectral resolution
    animation_interval_ms: int = 220

    # Signal handling
    invert_second_channel: bool = False
    remove_linear_trend: bool = True
    spectrum_ema_alpha: float = 0.15
    state_spectrum_ema_alpha: float = 0.08

    # Exploratory frequency bands. These labels describe analysis windows only;
    # a peak in them is not proof of Rabi or Ramsey dynamics.
    high_frequency_band_min_hz: float = 500.0
    high_frequency_band_max_hz: float = 20_000.0
    low_frequency_band_min_hz: float = 10.0
    low_frequency_band_max_hz: float = 500.0

    # Candidate-feature detection
    candidate_snr_threshold: float = 6.0
    candidate_min_frequency_hz: float = 10.0
    excluded_mains_hz: float = 50.0
    excluded_mains_harmonics: int = 8
    excluded_half_width_hz: float = 2.5

    # Logging
    output_root: str = "yb_measurements"
    metrics_log_period_s: float = 1.0
    raw_snapshot_blocks: int = 1

    # Simulation fallback. Set False in the real experiment if failure to open
    # the DAQ should stop execution instead of using synthetic data.
    allow_simulation_fallback: bool = True

    @property
    def channels(self) -> Tuple[str, str]:
        return (
            f"{self.device_name}/{self.pd1_channel}",
            f"{self.device_name}/{self.pd2_channel}",
        )

    @property
    def block_duration_s(self) -> float:
        return self.samples_per_read / self.sample_rate_hz

    @property
    def frequency_resolution_hz(self) -> float:
        return self.sample_rate_hz / self.samples_per_read


CFG = Config()

if CFG.sample_rate_hz <= 2.0 * CFG.high_frequency_band_max_hz:
    raise ValueError(
        "sample_rate_hz must be greater than twice high_frequency_band_max_hz "
        "to satisfy the Nyquist criterion."
    )


# =============================================================================
# EXPERIMENT STATE AND SESSION FILES
# =============================================================================

state = {
    "pump_on": True,
    "yb_excitation_on": True,
    "mw_on": False,
    "pd1_enabled": True,
    "pd2_enabled": True,
    "hypothesis_channel": "A",
    "paused": False,
    "recording": True,
    "baseline_ready": False,
    "daq_mode": "starting",
    "last_event": "Program started",
}

session_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
session_dir = Path(CFG.output_root) / f"session_{session_timestamp}"
session_dir.mkdir(parents=True, exist_ok=True)
metrics_csv_path = session_dir / "metrics.csv"
events_csv_path = session_dir / "events.csv"
metadata_json_path = session_dir / "metadata.json"

metadata = {
    "created_utc": datetime.now(timezone.utc).isoformat(),
    "software": "Yb3+ spectroscopy testbed monitor v2",
    "configuration": asdict(CFG),
    "scientific_scope": (
        "Exploratory acquisition and candidate-feature detection. No displayed "
        "feature is by itself evidence of a physical qubit or coherent control."
    ),
}
metadata_json_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

with metrics_csv_path.open("w", newline="", encoding="utf-8") as f:
    writer = csv.writer(f)
    writer.writerow([
        "timestamp_utc",
        "elapsed_s",
        "daq_mode",
        "pump_on",
        "yb_excitation_on",
        "mw_on",
        "hypothesis_channel",
        "pd1_mean_v",
        "pd1_rms_ac_v",
        "pd1_peak_to_peak_v",
        "pd2_mean_v",
        "pd2_rms_ac_v",
        "pd2_peak_to_peak_v",
        "candidate_frequency_hz",
        "candidate_snr",
        "mw_on_off_delta_at_candidate_v",
        "classification",
    ])

with events_csv_path.open("w", newline="", encoding="utf-8") as f:
    writer = csv.writer(f)
    writer.writerow(["timestamp_utc", "elapsed_s", "event", "details"])


session_start_monotonic = time.monotonic()
last_metrics_log_monotonic = 0.0
last_block: Optional[Tuple[np.ndarray, np.ndarray]] = None


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def elapsed_s() -> float:
    return time.monotonic() - session_start_monotonic


def log_event(event: str, details: str = "") -> None:
    state["last_event"] = event
    with events_csv_path.open("a", newline="", encoding="utf-8") as f:
        csv.writer(f).writerow([utc_now_iso(), f"{elapsed_s():.3f}", event, details])


# =============================================================================
# DAQ / SIMULATION SOURCE
# =============================================================================

class DataSource:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.task = None
        self.sim_phase = 0.0
        self.sim_rng = np.random.default_rng(42)
        self.mode = "simulation"

        if nidaqmx is None:
            if not cfg.allow_simulation_fallback:
                raise RuntimeError("nidaqmx is not installed.")
            log_event("DAQ unavailable", "nidaqmx not installed; simulation enabled")
            return

        try:
            self.task = nidaqmx.Task()
            terminal_config = {
                "RSE": TerminalConfiguration.RSE,
                "NRSE": TerminalConfiguration.NRSE,
                "DIFF": TerminalConfiguration.DIFFERENTIAL,
            }[cfg.terminal_mode.upper()]

            for channel in cfg.channels:
                self.task.ai_channels.add_ai_voltage_chan(
                    channel,
                    terminal_config=terminal_config,
                    min_val=cfg.voltage_min,
                    max_val=cfg.voltage_max,
                )

            self.task.timing.cfg_samp_clk_timing(
                rate=cfg.sample_rate_hz,
                sample_mode=AcquisitionType.CONTINUOUS,
                samps_per_chan=cfg.samples_per_read,
            )
            self.task.start()
            self.mode = "NI-DAQ"
            log_event("DAQ started", f"channels={cfg.channels}, rate={cfg.sample_rate_hz}")
        except Exception as exc:
            if self.task is not None:
                try:
                    self.task.close()
                except Exception:
                    pass
                self.task = None

            if not cfg.allow_simulation_fallback:
                raise RuntimeError(f"Could not initialise NI-DAQ: {exc}") from exc

            self.mode = "simulation"
            log_event("DAQ initialisation failed", f"{exc}; simulation enabled")

    def read(self) -> Tuple[np.ndarray, np.ndarray]:
        if self.task is not None:
            data = self.task.read(
                number_of_samples_per_channel=self.cfg.samples_per_read,
                timeout=max(5.0, 3.0 * self.cfg.block_duration_s),
            )
            arr = np.asarray(data, dtype=float)
            if arr.shape != (2, self.cfg.samples_per_read):
                raise RuntimeError(f"Unexpected DAQ data shape: {arr.shape}")
            ch0, ch1 = arr[0], arr[1]
        else:
            ch0, ch1 = self._simulate()

        if self.cfg.invert_second_channel:
            ch1 = -ch1
        return ch0, ch1

    def _simulate(self) -> Tuple[np.ndarray, np.ndarray]:
        """Synthetic data for UI/software testing; never experimental evidence."""
        n = self.cfg.samples_per_read
        t = (np.arange(n) / self.cfg.sample_rate_hz) + self.sim_phase
        self.sim_phase += n / self.cfg.sample_rate_hz

        noise1 = self.sim_rng.normal(0.0, 0.010, n)
        noise2 = self.sim_rng.normal(0.0, 0.008, n)
        mains = 0.012 * np.sin(2.0 * np.pi * 50.0 * t)

        # Candidate frequency appears more strongly in the simulated MW-ON state.
        candidate_amp = 0.035 if state["mw_on"] else 0.006
        candidate = candidate_amp * np.sin(2.0 * np.pi * 3_200.0 * t)
        slow = 0.020 * np.sin(2.0 * np.pi * 73.0 * t)

        ch0 = 0.70 + mains + candidate + slow + noise1
        ch1 = 0.42 + 0.5 * mains - 0.4 * candidate + noise2
        return ch0, ch1

    def close(self) -> None:
        if self.task is not None:
            try:
                self.task.stop()
            finally:
                self.task.close()
                self.task = None


source = DataSource(CFG)
state["daq_mode"] = source.mode


# =============================================================================
# SIGNAL PROCESSING
# =============================================================================

freqs_full = np.fft.rfftfreq(CFG.samples_per_read, d=1.0 / CFG.sample_rate_hz)
hann_window = np.hanning(CFG.samples_per_read)
window_coherent_gain = np.mean(hann_window)

spectrum_ema: Optional[np.ndarray] = None
mw_on_spectrum_ema: Optional[np.ndarray] = None
mw_off_spectrum_ema: Optional[np.ndarray] = None
baseline_spectrum: Optional[np.ndarray] = None
baseline_noise_floor: Optional[float] = None

candidate_history = deque(maxlen=30)


def detrend_signal(signal: np.ndarray) -> np.ndarray:
    signal = np.asarray(signal, dtype=float)
    y = signal - np.mean(signal)
    if not CFG.remove_linear_trend:
        return y
    x = np.linspace(-1.0, 1.0, y.size)
    slope = np.dot(x, y) / max(np.dot(x, x), np.finfo(float).eps)
    return y - slope * x


def amplitude_spectrum(signal: np.ndarray) -> np.ndarray:
    """Single-sided Hann-windowed amplitude spectrum in volts."""
    y = detrend_signal(signal) * hann_window
    spec = np.abs(np.fft.rfft(y))
    scale = 2.0 / (CFG.samples_per_read * max(window_coherent_gain, 1e-12))
    spec *= scale
    spec[0] *= 0.5
    if CFG.samples_per_read % 2 == 0:
        spec[-1] *= 0.5
    return spec


def robust_noise_floor(values: np.ndarray) -> Tuple[float, float]:
    """Return robust centre and sigma estimate using median and MAD."""
    finite = values[np.isfinite(values)]
    if finite.size == 0:
        return 0.0, 1e-12
    median = float(np.median(finite))
    mad = float(np.median(np.abs(finite - median)))
    sigma = max(1.4826 * mad, 1e-12)
    return median, sigma


def mains_exclusion_mask(freqs: np.ndarray) -> np.ndarray:
    allowed = np.ones_like(freqs, dtype=bool)
    for harmonic in range(1, CFG.excluded_mains_harmonics + 1):
        centre = harmonic * CFG.excluded_mains_hz
        allowed &= np.abs(freqs - centre) > CFG.excluded_half_width_hz
    return allowed


def detect_candidate_feature(spectrum: np.ndarray) -> Tuple[float, float, str]:
    search_mask = (
        (freqs_full >= CFG.candidate_min_frequency_hz)
        & (freqs_full <= CFG.high_frequency_band_max_hz)
        & mains_exclusion_mask(freqs_full)
    )
    search_values = spectrum[search_mask]
    search_freqs = freqs_full[search_mask]

    if search_values.size == 0:
        return math.nan, 0.0, "no analyzable band"

    centre, sigma = robust_noise_floor(search_values)
    peak_index = int(np.argmax(search_values))
    peak_value = float(search_values[peak_index])
    peak_frequency = float(search_freqs[peak_index])
    snr = max(0.0, (peak_value - centre) / sigma)

    candidate_history.append(peak_frequency)
    stable = False
    if len(candidate_history) >= 10:
        recent = np.asarray(candidate_history, dtype=float)
        stable = float(np.std(recent)) <= 2.0 * CFG.frequency_resolution_hz

    if snr >= CFG.candidate_snr_threshold and stable:
        classification = "repeatable candidate feature"
    elif snr >= CFG.candidate_snr_threshold:
        classification = "candidate feature; frequency not yet stable"
    else:
        classification = "no candidate above threshold"

    return peak_frequency, snr, classification


def exponential_average(
    previous: Optional[np.ndarray], current: np.ndarray, alpha: float
) -> np.ndarray:
    if previous is None:
        return current.copy()
    return (1.0 - alpha) * previous + alpha * current


def signal_metrics(signal: np.ndarray) -> Tuple[float, float, float]:
    mean = float(np.mean(signal))
    ac = signal - mean
    rms_ac = float(np.sqrt(np.mean(ac * ac)))
    peak_to_peak = float(np.ptp(signal))
    return mean, rms_ac, peak_to_peak


def mw_delta_at_frequency(frequency_hz: float) -> float:
    if (
        mw_on_spectrum_ema is None
        or mw_off_spectrum_ema is None
        or not np.isfinite(frequency_hz)
    ):
        return math.nan
    idx = int(np.argmin(np.abs(freqs_full - frequency_hz)))
    return float(mw_on_spectrum_ema[idx] - mw_off_spectrum_ema[idx])


# =============================================================================
# FIGURE
# =============================================================================

plt.style.use("dark_background")
fig, axes = plt.subplots(3, 2, figsize=(14, 9))
fig.canvas.manager.set_window_title("Yb3+ Spectroscopy Testbed Monitor")

ax_pd1 = axes[0, 0]
ax_pd2 = axes[0, 1]
ax_spectrum = axes[1, 0]
ax_high = axes[1, 1]
ax_low = axes[2, 0]
ax_text = axes[2, 1]
ax_text.axis("off")

sample_time_ms = np.arange(CFG.samples_per_read) / CFG.sample_rate_hz * 1_000.0

line_pd1, = ax_pd1.plot([], [], label="PD1 measured voltage")
line_pd2, = ax_pd2.plot([], [], label="PD2 measured voltage")
ax_pd1.set_xlim(sample_time_ms[0], sample_time_ms[-1])
ax_pd2.set_xlim(sample_time_ms[0], sample_time_ms[-1])
for ax, title in (
    (ax_pd1, "PD1 time domain"),
    (ax_pd2, "PD2 time domain"),
):
    ax.set_xlabel("Time (ms)")
    ax.set_ylabel("Voltage (V)")
    ax.set_title(title)
    ax.legend(loc="upper right")

line_spectrum, = ax_spectrum.plot([], [], label="PD1 spectrum EMA")
line_baseline, = ax_spectrum.plot([], [], "--", alpha=0.75, label="Captured baseline")
line_mw_on, = ax_spectrum.plot([], [], alpha=0.55, label="MW ON average")
line_mw_off, = ax_spectrum.plot([], [], alpha=0.55, label="MW OFF average")
candidate_marker, = ax_spectrum.plot([], [], marker="o", linestyle="None", label="Candidate")
ax_spectrum.set_xlim(0.0, CFG.high_frequency_band_max_hz)
ax_spectrum.set_ylim(0.0, 1.0)
ax_spectrum.set_xlabel("Frequency (Hz)")
ax_spectrum.set_ylabel("Amplitude (V)")
ax_spectrum.set_title("PD1 spectrum and MW ON/OFF comparison")
ax_spectrum.legend(loc="upper right", fontsize=8)

line_high, = ax_high.plot([], [], label="Exploratory high-frequency band")
ax_high.set_xlim(CFG.high_frequency_band_min_hz, CFG.high_frequency_band_max_hz)
ax_high.set_ylim(0.0, 1.0)
ax_high.set_xlabel("Frequency (Hz)")
ax_high.set_ylabel("Amplitude (V)")
ax_high.set_title("High-frequency candidate band (not proof of Rabi dynamics)")
ax_high.legend(loc="upper right", fontsize=8)

line_low, = ax_low.plot([], [], label="Exploratory low-frequency band")
ax_low.set_xlim(CFG.low_frequency_band_min_hz, CFG.low_frequency_band_max_hz)
ax_low.set_ylim(0.0, 1.0)
ax_low.set_xlabel("Frequency (Hz)")
ax_low.set_ylabel("Amplitude (V)")
ax_low.set_title("Low-frequency candidate band (not proof of Ramsey dynamics)")
ax_low.legend(loc="upper right", fontsize=8)

status_text = ax_text.text(
    0.02, 0.98, "", transform=ax_text.transAxes,
    va="top", ha="left", fontsize=9.2, family="monospace"
)

help_text = (
    "Keyboard controls (state annotation unless instrument control is added):\n"
    " 1  Pump state ON/OFF       2  Yb excitation state ON/OFF\n"
    " 3  Microwave state ON/OFF  4/5  Show PD1/PD2 ON/OFF\n"
    " 6  Hypothesis channel A/B/C\n"
    " b  Capture current spectrum as baseline\n"
    " c  Clear MW ON/OFF averages and candidate history\n"
    " s  Save raw snapshot (.npz)\n"
    " e  Add manual event marker\n"
    " l  Metrics logging ON/OFF\n"
    " space  Pause/resume   r  Reset annotations   q/esc  Quit"
)
ax_text.text(
    0.02, 0.02, help_text, transform=ax_text.transAxes,
    va="bottom", ha="left", fontsize=8.5
)


# =============================================================================
# UI AND LOGGING
# =============================================================================

current_metrics = {
    "pd1_mean": math.nan,
    "pd1_rms": math.nan,
    "pd1_pp": math.nan,
    "pd2_mean": math.nan,
    "pd2_rms": math.nan,
    "pd2_pp": math.nan,
    "candidate_frequency": math.nan,
    "candidate_snr": 0.0,
    "mw_delta": math.nan,
    "classification": "not analysed",
}


def update_status_text() -> None:
    f_candidate = current_metrics["candidate_frequency"]
    f_text = "n/a" if not np.isfinite(f_candidate) else f"{f_candidate:9.2f} Hz"
    delta = current_metrics["mw_delta"]
    delta_text = "n/a" if not np.isfinite(delta) else f"{delta:+.3e} V"

    status = (
        f"SOURCE: {state['daq_mode']:<10}   "
        f"RUN: {'PAUSED' if state['paused'] else 'ACTIVE':<6}   "
        f"LOG: {'ON' if state['recording'] else 'OFF'}\n"
        f"Pump annotation: {'ON' if state['pump_on'] else 'OFF':<3}   "
        f"Yb excitation: {'ON' if state['yb_excitation_on'] else 'OFF':<3}   "
        f"MW annotation: {'ON' if state['mw_on'] else 'OFF':<3}\n"
        f"PD1 visible: {'YES' if state['pd1_enabled'] else 'NO ':<3}   "
        f"PD2 visible: {'YES' if state['pd2_enabled'] else 'NO ':<3}   "
        f"Hypothesis channel: {state['hypothesis_channel']}\n"
        f"Fs: {CFG.sample_rate_hz:,.0f} Hz   "
        f"N: {CFG.samples_per_read:,}   "
        f"df: {CFG.frequency_resolution_hz:.2f} Hz\n"
        f"PD1 mean/RMSac/p-p: {current_metrics['pd1_mean']:+.4f} / "
        f"{current_metrics['pd1_rms']:.4e} / {current_metrics['pd1_pp']:.4e} V\n"
        f"PD2 mean/RMSac/p-p: {current_metrics['pd2_mean']:+.4f} / "
        f"{current_metrics['pd2_rms']:.4e} / {current_metrics['pd2_pp']:.4e} V\n"
        f"Candidate: {f_text}   robust SNR: {current_metrics['candidate_snr']:.2f}\n"
        f"MW ON-OFF delta: {delta_text}\n"
        f"Assessment: {current_metrics['classification']}\n"
        f"Baseline: {'CAPTURED' if state['baseline_ready'] else 'not captured'}\n"
        f"Last event: {state['last_event']}\n\n"
        "CAUTION: candidate features require controls, calibration,\n"
        "repeatability and model comparison before interpretation."
    )
    status_text.set_text(status)


def autoscale_axis(ax, values: np.ndarray, symmetric: bool = False) -> None:
    finite = values[np.isfinite(values)]
    if finite.size == 0:
        return
    ymin = float(np.min(finite))
    ymax = float(np.max(finite))
    if math.isclose(ymin, ymax):
        span = max(abs(ymin) * 0.1, 1e-6)
        ymin -= span
        ymax += span
    margin = 0.10 * (ymax - ymin)
    if symmetric:
        max_abs = max(abs(ymin), abs(ymax), 1e-9)
        ax.set_ylim(-1.1 * max_abs, 1.1 * max_abs)
    else:
        ax.set_ylim(ymin - margin, ymax + margin)


def log_metrics_if_due() -> None:
    global last_metrics_log_monotonic
    now = time.monotonic()
    if not state["recording"] or now - last_metrics_log_monotonic < CFG.metrics_log_period_s:
        return
    last_metrics_log_monotonic = now

    with metrics_csv_path.open("a", newline="", encoding="utf-8") as f:
        csv.writer(f).writerow([
            utc_now_iso(),
            f"{elapsed_s():.3f}",
            state["daq_mode"],
            int(state["pump_on"]),
            int(state["yb_excitation_on"]),
            int(state["mw_on"]),
            state["hypothesis_channel"],
            f"{current_metrics['pd1_mean']:.9g}",
            f"{current_metrics['pd1_rms']:.9g}",
            f"{current_metrics['pd1_pp']:.9g}",
            f"{current_metrics['pd2_mean']:.9g}",
            f"{current_metrics['pd2_rms']:.9g}",
            f"{current_metrics['pd2_pp']:.9g}",
            f"{current_metrics['candidate_frequency']:.9g}",
            f"{current_metrics['candidate_snr']:.9g}",
            f"{current_metrics['mw_delta']:.9g}",
            current_metrics["classification"],
        ])


def save_raw_snapshot() -> None:
    if last_block is None:
        log_event("Snapshot skipped", "No acquired block available")
        return
    ch0, ch1 = last_block
    snapshot_path = session_dir / f"snapshot_{datetime.now().strftime('%H%M%S_%f')}.npz"
    np.savez_compressed(
        snapshot_path,
        pd1_v=ch0,
        pd2_v=ch1,
        sample_rate_hz=CFG.sample_rate_hz,
        timestamp_utc=utc_now_iso(),
        state_json=json.dumps(state),
        config_json=json.dumps(asdict(CFG)),
    )
    log_event("Raw snapshot saved", str(snapshot_path))


# =============================================================================
# ANIMATION
# =============================================================================

def animate(_frame):
    global last_block
    global spectrum_ema, mw_on_spectrum_ema, mw_off_spectrum_ema

    if state["paused"]:
        update_status_text()
        return (
            line_pd1, line_pd2, line_spectrum, line_baseline,
            line_mw_on, line_mw_off, candidate_marker,
            line_high, line_low, status_text,
        )

    try:
        ch0, ch1 = source.read()
    except Exception as exc:
        state["paused"] = True
        log_event("Acquisition error", repr(exc))
        update_status_text()
        return (
            line_pd1, line_pd2, line_spectrum, line_baseline,
            line_mw_on, line_mw_off, candidate_marker,
            line_high, line_low, status_text,
        )

    last_block = (ch0.copy(), ch1.copy())

    display_pd1 = ch0 if state["pd1_enabled"] else np.zeros_like(ch0)
    display_pd2 = ch1 if state["pd2_enabled"] else np.zeros_like(ch1)
    line_pd1.set_data(sample_time_ms, display_pd1)
    line_pd2.set_data(sample_time_ms, display_pd2)
    autoscale_axis(ax_pd1, display_pd1)
    autoscale_axis(ax_pd2, display_pd2)

    current_spectrum = amplitude_spectrum(ch0)
    spectrum_ema = exponential_average(
        spectrum_ema, current_spectrum, CFG.spectrum_ema_alpha
    )

    if state["mw_on"]:
        mw_on_spectrum_ema = exponential_average(
            mw_on_spectrum_ema, current_spectrum, CFG.state_spectrum_ema_alpha
        )
    else:
        mw_off_spectrum_ema = exponential_average(
            mw_off_spectrum_ema, current_spectrum, CFG.state_spectrum_ema_alpha
        )

    candidate_frequency, candidate_snr, classification = detect_candidate_feature(
        spectrum_ema
    )
    delta = mw_delta_at_frequency(candidate_frequency)

    current_metrics["pd1_mean"], current_metrics["pd1_rms"], current_metrics["pd1_pp"] = signal_metrics(ch0)
    current_metrics["pd2_mean"], current_metrics["pd2_rms"], current_metrics["pd2_pp"] = signal_metrics(ch1)
    current_metrics["candidate_frequency"] = candidate_frequency
    current_metrics["candidate_snr"] = candidate_snr
    current_metrics["mw_delta"] = delta

    if source.mode == "simulation":
        classification = "SIMULATION ONLY — " + classification
    elif baseline_spectrum is None:
        classification += "; capture baseline"
    elif mw_on_spectrum_ema is None or mw_off_spectrum_ema is None:
        classification += "; collect both MW states"
    current_metrics["classification"] = classification

    line_spectrum.set_data(freqs_full, spectrum_ema)
    line_baseline.set_data(
        freqs_full if baseline_spectrum is not None else [],
        baseline_spectrum if baseline_spectrum is not None else [],
    )
    line_mw_on.set_data(
        freqs_full if mw_on_spectrum_ema is not None else [],
        mw_on_spectrum_ema if mw_on_spectrum_ema is not None else [],
    )
    line_mw_off.set_data(
        freqs_full if mw_off_spectrum_ema is not None else [],
        mw_off_spectrum_ema if mw_off_spectrum_ema is not None else [],
    )

    if np.isfinite(candidate_frequency):
        candidate_idx = int(np.argmin(np.abs(freqs_full - candidate_frequency)))
        candidate_marker.set_data(
            [candidate_frequency], [spectrum_ema[candidate_idx]]
        )
    else:
        candidate_marker.set_data([], [])

    spectrum_mask = freqs_full <= CFG.high_frequency_band_max_hz
    visible_spectra = [spectrum_ema[spectrum_mask]]
    if baseline_spectrum is not None:
        visible_spectra.append(baseline_spectrum[spectrum_mask])
    if mw_on_spectrum_ema is not None:
        visible_spectra.append(mw_on_spectrum_ema[spectrum_mask])
    if mw_off_spectrum_ema is not None:
        visible_spectra.append(mw_off_spectrum_ema[spectrum_mask])
    combined = np.concatenate(visible_spectra)
    ax_spectrum.set_ylim(0.0, max(float(np.max(combined)) * 1.20, 1e-8))

    high_mask = (
        (freqs_full >= CFG.high_frequency_band_min_hz)
        & (freqs_full <= CFG.high_frequency_band_max_hz)
    )
    low_mask = (
        (freqs_full >= CFG.low_frequency_band_min_hz)
        & (freqs_full <= CFG.low_frequency_band_max_hz)
    )
    line_high.set_data(freqs_full[high_mask], spectrum_ema[high_mask])
    line_low.set_data(freqs_full[low_mask], spectrum_ema[low_mask])
    ax_high.set_ylim(0.0, max(float(np.max(spectrum_ema[high_mask])) * 1.20, 1e-8))
    ax_low.set_ylim(0.0, max(float(np.max(spectrum_ema[low_mask])) * 1.20, 1e-8))

    update_status_text()
    log_metrics_if_due()

    return (
        line_pd1, line_pd2, line_spectrum, line_baseline,
        line_mw_on, line_mw_off, candidate_marker,
        line_high, line_low, status_text,
    )


# =============================================================================
# KEYBOARD HANDLER
# =============================================================================

def on_key(event) -> None:
    global baseline_spectrum, baseline_noise_floor
    global mw_on_spectrum_ema, mw_off_spectrum_ema

    key = event.key
    if key == "1":
        state["pump_on"] = not state["pump_on"]
        log_event("Pump annotation changed", str(state["pump_on"]))
    elif key == "2":
        state["yb_excitation_on"] = not state["yb_excitation_on"]
        log_event("Yb excitation annotation changed", str(state["yb_excitation_on"]))
    elif key == "3":
        state["mw_on"] = not state["mw_on"]
        log_event("Microwave annotation changed", str(state["mw_on"]))
    elif key == "4":
        state["pd1_enabled"] = not state["pd1_enabled"]
    elif key == "5":
        state["pd2_enabled"] = not state["pd2_enabled"]
    elif key == "6":
        state["hypothesis_channel"] = {"A": "B", "B": "C", "C": "A"}[
            state["hypothesis_channel"]
        ]
        log_event("Hypothesis channel changed", state["hypothesis_channel"])
    elif key == " ":
        state["paused"] = not state["paused"]
        log_event("Pause changed", str(state["paused"]))
    elif key == "b":
        if spectrum_ema is not None:
            baseline_spectrum = spectrum_ema.copy()
            _, baseline_noise_floor = robust_noise_floor(
                baseline_spectrum[
                    (freqs_full >= CFG.candidate_min_frequency_hz)
                    & (freqs_full <= CFG.high_frequency_band_max_hz)
                ]
            )
            state["baseline_ready"] = True
            log_event("Baseline captured", f"noise_sigma={baseline_noise_floor:.6g}")
        else:
            log_event("Baseline skipped", "No spectrum available")
    elif key == "c":
        mw_on_spectrum_ema = None
        mw_off_spectrum_ema = None
        candidate_history.clear()
        log_event("Comparison averages cleared")
    elif key == "s":
        save_raw_snapshot()
    elif key == "e":
        log_event("Manual marker", "User pressed E")
    elif key == "l":
        state["recording"] = not state["recording"]
        log_event("Metrics logging changed", str(state["recording"]))
    elif key == "r":
        state.update({
            "pump_on": True,
            "yb_excitation_on": True,
            "mw_on": False,
            "pd1_enabled": True,
            "pd2_enabled": True,
            "hypothesis_channel": "A",
            "paused": False,
        })
        log_event("State annotations reset")
    elif key in ("q", "escape"):
        log_event("Quit requested")
        plt.close(fig)

    update_status_text()
    fig.canvas.draw_idle()


fig.canvas.mpl_connect("key_press_event", on_key)
update_status_text()
ani = FuncAnimation(
    fig,
    animate,
    interval=CFG.animation_interval_ms,
    blit=False,
    cache_frame_data=False,
)

try:
    plt.tight_layout()
    plt.show()
finally:
    source.close()
    log_event("Session closed", f"output={session_dir}")
    print(f"Session data saved in: {session_dir.resolve()}")