synthetic.py

# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 The Regents of the University of California, Mazin Lab

"""Data-driven synthetic MKID pulse generator for continuous energy training.

Builds IQ pulse templates from real calibration data, then generates
physically realistic synthetic pulses at arbitrary energies via linear
interpolation of templates and energy-scaled residual resampling.
Designed to replace discrete-label training which causes classification
behavior.

Usage:
    from synthetic import SyntheticPulseGenerator

    iq, energies, wl_ids, wl_labels, opt_R, peak_heights = load_ptsi_data()
    gen = SyntheticPulseGenerator(iq, wl_ids, wl_labels)
    gen.generate_report("synthetic_diagnostics.pdf", iq, wl_ids)

    # In training loop:
    iq_batch, e_batch = gen.generate_batch(n_pulses=256, rng=rng)
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path
from typing import Tuple

from backend import mx, to_numpy
import numpy as np
from scipy.interpolate import interp1d


# ---------------------------------------------------------------------------
# Pulse alignment
# ---------------------------------------------------------------------------


def align_pulses(
    traces: np.ndarray,
    max_shift: int = 5,
    n_iter: int = 3,
    subsample: bool = True,
) -> Tuple[np.ndarray, np.ndarray]:
    """Align IQ pulses using iterative FFT cross-correlation.

    Individual pulses have slight timing jitter from noise-induced trigger
    uncertainty.  This function aligns them to a common reference so that
    the mean template is sharp rather than smeared.

    Parameters
    ----------
    traces : ndarray, shape (n_pulses, seq_len, 2)
        IQ traces for a single wavelength.
    max_shift : int
        Maximum allowed shift in samples (both directions).
    n_iter : int
        Number of alignment iterations (template improves each round).
    subsample : bool
        If True, use parabolic interpolation for sub-sample precision
        and cubic spline shifting (no Gibbs ringing).

    Returns
    -------
    aligned : ndarray, shape (n_pulses, seq_len, 2)
    shifts : ndarray, shape (n_pulses,), dtype float
        Applied shifts (positive = original was late, shifted earlier).
        Float if subsample=True, int otherwise.
    """
    from scipy.ndimage import shift as ndimage_shift

    seq_len = traces.shape[1]
    aligned = traces.copy()
    cumulative_shifts = np.zeros(len(traces), dtype=np.float64)

    for _ in range(n_iter):
        template = aligned.mean(axis=0)  # (seq_len, 2)

        # Amplitude envelopes for cross-correlation
        t_amp = np.sqrt(template[:, 0] ** 2 + template[:, 1] ** 2)
        x_amp = np.sqrt(aligned[:, :, 0] ** 2 + aligned[:, :, 1] ** 2)

        # FFT-based circular cross-correlation (vectorized)
        T = np.fft.rfft(t_amp)  # (n_freq,)
        X = np.fft.rfft(x_amp, axis=1)  # (n_pulses, n_freq)
        corr = np.fft.irfft(X * np.conj(T), n=seq_len, axis=1)

        # Find peak within +/-max_shift of zero lag
        # Zero lag is index 0; negative lags wrap to end of array
        idx_neg = np.arange(-max_shift, 0) % seq_len
        idx_pos = np.arange(0, max_shift + 1)
        search_idx = np.concatenate([idx_neg, idx_pos])
        search = corr[:, search_idx]  # (n_pulses, 2*max_shift+1)
        peak_pos = np.argmax(search, axis=1)  # index into search window
        int_shifts = peak_pos - max_shift  # integer shifts

        if subsample:
            # Parabolic interpolation for sub-sample precision
            sub_offsets = np.zeros(len(traces), dtype=np.float64)
            for i, p in enumerate(peak_pos):
                if 0 < p < len(search_idx) - 1:
                    y_m1 = search[i, p - 1]
                    y_0 = search[i, p]
                    y_p1 = search[i, p + 1]
                    denom = y_m1 - 2 * y_0 + y_p1
                    if abs(denom) > 1e-12:
                        sub_offsets[i] = 0.5 * (y_m1 - y_p1) / denom

            shifts = int_shifts.astype(np.float64) + sub_offsets

            # Linear interpolation shift via scipy.ndimage.shift.
            # order=1 matches order=3 in practice for sub-sample shifts
            # and avoids any possibility of ringing artifacts.
            for i, s in enumerate(shifts):
                if abs(s) > 0.01:
                    for ch in range(2):
                        aligned[i, :, ch] = ndimage_shift(
                            aligned[i, :, ch], -s, order=1, mode='nearest',
                        )
        else:
            shifts = int_shifts
            for i, s in enumerate(shifts):
                if s != 0:
                    aligned[i] = np.roll(aligned[i], -s, axis=0)

        cumulative_shifts += shifts

    return aligned, cumulative_shifts


def align_pulses_to_template(
    traces: np.ndarray,
    template: np.ndarray,
    max_shift: int = 5,
    subsample: bool = True,
) -> Tuple[np.ndarray, np.ndarray]:
    """Align IQ pulses to a fixed template without refitting it.

    This is used when a validation split is held out before preprocessing:
    the template can be estimated from training pulses only, then applied to
    both train and validation pulses without letting validation pulses shape
    the reference.
    """
    from scipy.ndimage import shift as ndimage_shift

    seq_len = traces.shape[1]
    aligned = traces.copy()

    t_amp = np.sqrt(template[:, 0] ** 2 + template[:, 1] ** 2)
    x_amp = np.sqrt(aligned[:, :, 0] ** 2 + aligned[:, :, 1] ** 2)

    T = np.fft.rfft(t_amp)
    X = np.fft.rfft(x_amp, axis=1)
    corr = np.fft.irfft(X * np.conj(T), n=seq_len, axis=1)

    idx_neg = np.arange(-max_shift, 0) % seq_len
    idx_pos = np.arange(0, max_shift + 1)
    search_idx = np.concatenate([idx_neg, idx_pos])
    search = corr[:, search_idx]
    peak_pos = np.argmax(search, axis=1)
    int_shifts = peak_pos - max_shift

    if subsample:
        sub_offsets = np.zeros(len(traces), dtype=np.float64)
        for i, p in enumerate(peak_pos):
            if 0 < p < len(search_idx) - 1:
                y_m1 = search[i, p - 1]
                y_0 = search[i, p]
                y_p1 = search[i, p + 1]
                denom = y_m1 - 2 * y_0 + y_p1
                if abs(denom) > 1e-12:
                    sub_offsets[i] = 0.5 * (y_m1 - y_p1) / denom
        shifts = int_shifts.astype(np.float64) + sub_offsets
        for i, s in enumerate(shifts):
            if abs(s) > 0.01:
                for ch in range(2):
                    aligned[i, :, ch] = ndimage_shift(
                        aligned[i, :, ch], -s, order=1, mode="nearest",
                    )
    else:
        shifts = int_shifts
        for i, s in enumerate(shifts):
            if s != 0:
                aligned[i] = np.roll(aligned[i], -s, axis=0)

    return aligned, shifts


# ---------------------------------------------------------------------------
# CSD estimation (used only in diagnostic reports)
# ---------------------------------------------------------------------------


def _interpolate_csd_freq(csd_pre, freq_pre, csds, order_idx, freq_full):
    """Interpolate CSD from pre-trigger freq grid to full freq grid.

    Uses log interpolation for the diagonal PSDs (which can span many
    decades) and linear interpolation for the complex off-diagonal.

    We drop DC (k=0) from the pre-pulse grid before interpolating: after
    per-pulse detrending, csd_pre[0] is structurally zero (the mean was
    subtracted), so log(0) → -inf → 1e-30 floor → artificial low-freq
    dip. Instead, we treat the first resolved bin (freq_pre[1]) as the
    best estimate for DC and for any freq_full bin finer than 1/pre_len,
    using np.interp's default endpoint-clamping behaviour.
    """
    freq_pre_use = freq_pre[1:]
    for i in range(2):
        for j in range(2):
            vals = csd_pre[1:, i, j]
            if i == j:
                log_psd = np.log(np.maximum(vals.real, 1e-30))
                csds[order_idx, :, i, j] = np.exp(
                    np.interp(freq_full, freq_pre_use, log_psd)
                )
            else:
                csds[order_idx, :, i, j] = (
                    np.interp(freq_full, freq_pre_use, vals.real)
                    + 1j * np.interp(freq_full, freq_pre_use, vals.imag)
                )


def _compute_csd_from_raw(
    data_dir: str,
    wl_labels_sorted: list[str],
    n_freq: int,
    freq_full: np.ndarray,
    downsample: int = 1,
) -> np.ndarray:
    """Compute noise CSD from unwindowed raw traces (long pre-trigger).

    Returns
    -------
    csds : ndarray, shape (n_wl, n_freq, 2, 2), complex
    """
    from mkid_ssm import estimate_pulse_params

    n_wl = len(wl_labels_sorted)
    all_i, all_q = [], []
    wl_data = {}
    for wl in wl_labels_sorted:
        path = os.path.join(data_dir, f"{wl}.npz")
        d = np.load(path)
        mask = d["mask"]
        i_raw = d["i_trace"][mask]
        q_raw = d["q_trace"][mask]
        all_i.append(i_raw)
        all_q.append(q_raw)
        wl_data[wl] = (i_raw, q_raw)

    i_cat = np.concatenate(all_i, axis=0)
    q_cat = np.concatenate(all_q, axis=0)
    trigger, _ = estimate_pulse_params(i_cat, q_cat)

    n_pre = max(trigger - 10, 50)
    n_bl = min(200, n_pre)
    bl_i = i_cat[:, :n_bl].mean()
    bl_q = q_cat[:, :n_bl].mean()
    iq_radius = np.sqrt(bl_i ** 2 + bl_q ** 2)

    pre_len = -(-n_pre // downsample) if downsample > 1 else n_pre
    n_freq_pre = pre_len // 2 + 1
    freq_pre = np.fft.rfftfreq(pre_len)

    csds = np.zeros((n_wl, n_freq, 2, 2), dtype=complex)
    for order_idx, wl in enumerate(wl_labels_sorted):
        i_raw, q_raw = wl_data[wl]
        i_pre = i_raw[:, :n_pre] / iq_radius
        q_pre = q_raw[:, :n_pre] / iq_radius
        if downsample > 1:
            from scipy.signal import decimate
            i_pre = decimate(i_pre, downsample, axis=1, ftype='fir')
            q_pre = decimate(q_pre, downsample, axis=1, ftype='fir')

        noise = np.stack([i_pre, q_pre], axis=-1)
        noise -= noise.mean(axis=1, keepdims=True)

        R = np.fft.rfft(noise, axis=1)
        csd_pre = np.zeros((n_freq_pre, 2, 2), dtype=complex)
        for i in range(2):
            for j in range(2):
                csd_pre[:, i, j] = (
                    np.mean(R[:, :, i] * np.conj(R[:, :, j]), axis=0)
                    / pre_len
                )
        _interpolate_csd_freq(csd_pre, freq_pre, csds, order_idx, freq_full)

    return csds


# ---------------------------------------------------------------------------
# SyntheticPulseGenerator
# ---------------------------------------------------------------------------

# Constants from mkid_ssm.py
_SOURCE_NM = {
    "808nm": 808.0, "920nm": 920.0, "980nm": 980.0,
    "1120nm": 1120.0, "1310nm": 1310.0,
    "1110nm": 1110.0, "979nm": 978.6, "917nm": 916.8,
    "814nm": 813.7, "663nm": 663.1, "546nm": 546.0,
    "406nm": 405.9, "313nm": 313.0, "254nm": 254.0,
}


class SyntheticPulseGenerator:
    """Data-driven synthetic MKID pulse generator with continuous energy labels.

    Builds mean IQ pulse templates from real calibration data, estimates a
    per-wavelength 2x2 noise cross-spectral density from the pre-pulse
    baseline region, and generates synthetic pulses on-the-fly as

        synth_pulse(E) = template(E) + colored_Gaussian_noise_from_CSD

    Noise is synthesized fresh from the Cholesky-factored CSD every call,
    so no per-pulse residual (and no per-pulse amplitude fluctuation) from
    the real dataset is reused -- the model cannot memorize individual
    validation pulses through their noise fingerprint.

    Parameters
    ----------
    iq_data : mx.array or ndarray, shape (n_pulses, seq_len, 2)
        Normalized IQ traces from the existing data pipeline. MUST contain
        only training pulses -- validation pulses should be held out before
        this generator is constructed, so the pre-pulse CSD is estimated
        only from training data.
    wl_ids : ndarray, shape (n_pulses,)
        Wavelength ID (0-indexed) for each pulse.
    wl_labels : list of str
        Wavelength label strings.
    energy_margin : float
        Fractional margin beyond calibration range for energy sampling.
    sigma_clip : float or None
        If set, use iterative sigma-clipping when computing mean templates.
    pre_pulse_samples : int or None
        Number of pre-trigger samples to use for noise CSD estimation.
        If None, auto-detected from the highest-energy template's rising
        edge with a 4-sample safety back-off.
    """

    def __init__(
        self,
        iq_data: mx.array | np.ndarray,
        wl_ids: np.ndarray,
        wl_labels: list[str],
        energy_margin: float = 0.025,
        sigma_clip: float | None = 3.0,
        pre_pulse_samples: int | None = None,
    ):
        iq_np = to_numpy(iq_data) if not isinstance(iq_data, np.ndarray) else iq_data
        wl_ids = np.asarray(wl_ids)
        n_wl = len(wl_labels)
        seq_len = iq_np.shape[1]

        # Calibration energies in eV (sorted ascending by energy = descending by wavelength)
        cal_energies = np.array([
            1239.84 / _SOURCE_NM.get(wl, float(wl.replace("nm", "")))
            for wl in wl_labels
        ])

        # Sort by energy for monotonic spline interpolation
        sort_idx = np.argsort(cal_energies)
        cal_energies_sorted = cal_energies[sort_idx]
        wl_labels_sorted = [wl_labels[i] for i in sort_idx]
        wl_id_map = {wl_labels[i]: i for i in range(n_wl)}

        # --- Step 1: Compute templates via (optionally sigma-clipped) mean ---
        templates = np.zeros((n_wl, seq_len, 2), dtype=np.float64)

        for order_idx, sorted_i in enumerate(sort_idx):
            wl = wl_labels[sorted_i]
            wl_id = wl_id_map[wl]
            mask = wl_ids == wl_id
            traces = iq_np[mask]

            if sigma_clip is not None:
                # Vectorized iterative sigma-clipping (2 passes)
                data = traces.astype(np.float64)
                for _ in range(2):
                    mean = np.nanmean(data, axis=0, keepdims=True)
                    std = np.nanstd(data, axis=0, keepdims=True)
                    outlier = np.abs(data - mean) > sigma_clip * std
                    data = np.where(outlier, np.nan, data)
                tmpl = np.nanmean(data, axis=0)
            else:
                tmpl = traces.mean(axis=0)

            templates[order_idx] = tmpl

        # --- Step 2: Auto-detect pre-pulse length from highest-energy template ---
        # The highest-E template has the sharpest rising edge, so first-
        # threshold-crossing gives the most conservative trigger estimate.
        if pre_pulse_samples is None:
            tmpl_hi = templates[-1]
            amp = np.sqrt(tmpl_hi[:, 0] ** 2 + tmpl_hi[:, 1] ** 2)
            peak_idx = int(np.argmax(amp))
            bl_end = max(4, peak_idx // 4)
            baseline = amp[:bl_end].mean()
            peak_amp = amp[peak_idx]
            # 5% above baseline-to-peak is sensitive enough to catch the rise
            thr = baseline + 0.05 * (peak_amp - baseline)
            above = np.where(amp > thr)[0]
            trigger_idx = int(above[0]) if len(above) > 0 else peak_idx
            # Back off 4 samples so no rising edge bleeds into the noise region
            pre_pulse_samples = max(8, trigger_idx - 4)

        self._pre_pulse_samples = int(pre_pulse_samples)

        # --- Step 3: Estimate per-wavelength noise CSD from the pre-pulse
        # region of the windowed traces. This is pure detector noise (no
        # photon), and is what we sample fresh from to generate synthetic
        # noise -- no per-pulse residual reuse, so nothing from a real pulse
        # leaks into training.
        n_freq_full = seq_len // 2 + 1
        freq_full = np.fft.rfftfreq(seq_len)

        pre_len = self._pre_pulse_samples
        n_freq_pre = pre_len // 2 + 1
        freq_pre = np.fft.rfftfreq(pre_len)

        csds = np.zeros((n_wl, n_freq_full, 2, 2), dtype=complex)

        for order_idx, sorted_i in enumerate(sort_idx):
            wl_id = wl_id_map[wl_labels[sorted_i]]
            mask = wl_ids == wl_id
            pre = iq_np[mask, :pre_len, :].astype(np.float64)
            # Per-pulse DC detrend (doesn't affect non-DC PSD bins by DFT identity)
            pre = pre - pre.mean(axis=1, keepdims=True)

            R = np.fft.rfft(pre, axis=1)  # (n, n_freq_pre, 2)
            csd_pre = np.zeros((n_freq_pre, 2, 2), dtype=complex)
            for i in range(2):
                for j in range(2):
                    csd_pre[:, i, j] = (
                        np.mean(R[:, :, i] * np.conj(R[:, :, j]), axis=0)
                        / pre_len
                    )
            _interpolate_csd_freq(csd_pre, freq_pre, csds, order_idx, freq_full)

            # Rescale the interpolated CSD so that Parseval over the full
            # seq_len grid recovers the empirical per-channel variance of
            # the pre-pulse region. This corrects for the discretization
            # bias that arises because a coarse rfft-grid Parseval (pre_len
            # bins) and a fine rfft-grid Parseval (seq_len bins) give
            # systematically different integrals of the same smooth PSD.
            # The spectral SHAPE is still the interpolated estimate; only
            # the overall per-channel amplitude is anchored to ground truth.
            target_var = np.array([
                pre[:, :, 0].var(),
                pre[:, :, 1].var(),
            ])
            psd_full = csds[order_idx]  # (n_freq_full, 2, 2)
            if seq_len % 2 == 0:
                implied_var = np.array([
                    (psd_full[0, 0, 0].real
                     + 2 * psd_full[1:-1, 0, 0].real.sum()
                     + psd_full[-1, 0, 0].real) / seq_len,
                    (psd_full[0, 1, 1].real
                     + 2 * psd_full[1:-1, 1, 1].real.sum()
                     + psd_full[-1, 1, 1].real) / seq_len,
                ])
            else:
                implied_var = np.array([
                    (psd_full[0, 0, 0].real
                     + 2 * psd_full[1:, 0, 0].real.sum()) / seq_len,
                    (psd_full[0, 1, 1].real
                     + 2 * psd_full[1:, 1, 1].real.sum()) / seq_len,
                ])
            scale = np.where(implied_var > 0, target_var / implied_var, 1.0)
            # Apply scale to diagonals; for off-diagonals use geometric mean
            # (the 2x2 CSD then remains positive semi-definite after scaling).
            sII, sQQ = float(scale[0]), float(scale[1])
            sIQ = float(np.sqrt(sII * sQQ))
            csds[order_idx, :, 0, 0] *= sII
            csds[order_idx, :, 1, 1] *= sQQ
            csds[order_idx, :, 0, 1] *= sIQ
            csds[order_idx, :, 1, 0] *= sIQ

        # Enforce positive-semi-definite CSD at every frequency bin and
        # factor via Cholesky. The empirical estimate can violate
        # Cauchy-Schwarz |S_IQ|^2 <= S_II * S_QQ at a small fraction of
        # bins due to finite samples + per-component interpolation. In
        # those bins we shrink the off-diagonal toward zero (preserving
        # the diagonal, i.e. per-channel variance) until the matrix is
        # PSD. A tiny diagonal jitter handles pure numerical round-off.
        chol = np.zeros_like(csds)
        for w in range(n_wl):
            for f in range(n_freq_full):
                m = 0.5 * (csds[w, f] + csds[w, f].conj().T)
                s_ii = float(m[0, 0].real)
                s_qq = float(m[1, 1].real)
                s_iq = complex(m[0, 1])

                max_cross_abs = np.sqrt(max(0.0, s_ii * s_qq))
                cross_abs = abs(s_iq)
                if cross_abs > max_cross_abs and cross_abs > 0:
                    # Shrink to 99.9% of the Cauchy-Schwarz bound so the
                    # Hermitian determinant is strictly positive.
                    s_iq = s_iq * (0.999 * max_cross_abs / cross_abs)

                diag_max = max(s_ii, s_qq, 1e-30)
                jitter = 1e-12 * diag_max
                m_reg = np.array(
                    [[s_ii + jitter, s_iq],
                     [np.conj(s_iq), s_qq + jitter]],
                    dtype=complex,
                )
                chol[w, f] = np.linalg.cholesky(m_reg)
                # Keep the stored CSD consistent with the factored matrix
                # so rescaling and downstream diagnostics see the same
                # PSD that noise is actually drawn from.
                csds[w, f] = m_reg

        # --- Step 4: Build template interpolator ---
        self.template_spline = interp1d(
            cal_energies_sorted, templates, axis=0, kind="linear",
            fill_value="extrapolate",
        )

        # Store metadata
        self.cal_energies = cal_energies_sorted
        self.wl_labels_sorted = wl_labels_sorted
        self.E_min = float(cal_energies_sorted[0]) * (1 - energy_margin)
        self.E_max = float(cal_energies_sorted[-1]) * (1 + energy_margin)
        self.seq_len = seq_len
        self.n_wl = n_wl
        self.n_freq = n_freq_full
        self.templates = templates
        self._csd = csds
        self._csd_chol = chol
        # Map from sorted order index → original wl_id in the input data
        self._sorted_to_orig_wl_id = {
            order_idx: wl_id_map[wl_labels[sorted_i]]
            for order_idx, sorted_i in enumerate(sort_idx)
        }

    def generate_batch(
        self,
        n_pulses: int,
        rng: np.random.Generator,
    ) -> Tuple[mx.array, mx.array]:
        """Generate synthetic pulses at random continuous energies.

        Pulse = interpolated template(E) + fresh colored Gaussian noise drawn
        from the nearest calibration wavelength's noise CSD. The noise is
        synthesized per call from Cholesky factors so no per-pulse residual
        (and therefore no per-pulse amplitude fluctuation) from the real
        dataset is reused -- this avoids the leakage that the previous
        residual-resampling path introduced.

        Parameters
        ----------
        n_pulses : int
            Number of pulses to generate.
        rng : numpy Generator
            Random number generator for reproducibility.

        Returns
        -------
        iq : mx.array, shape (n_pulses, seq_len, 2)
        energies : mx.array, shape (n_pulses, 1)
        """
        energies = rng.uniform(self.E_min, self.E_max, n_pulses)
        templates = self.template_spline(energies)  # (n_pulses, seq_len, 2)

        # Nearest calibration wavelength (in energy) for each pulse's CSD.
        idx_cal = np.argmin(
            np.abs(self.cal_energies[None, :] - energies[:, None]),
            axis=1,
        )
        noise = self._sample_csd_noise(idx_cal, rng)

        iq = (templates + noise).astype(np.float32)
        return mx.array(iq), mx.array(energies.reshape(-1, 1).astype(np.float32))

    def _sample_csd_noise(
        self,
        wl_idx: np.ndarray,
        rng: np.random.Generator,
    ) -> np.ndarray:
        """Draw colored Gaussian noise from per-wavelength CSD Cholesky factors.

        Parameters
        ----------
        wl_idx : ndarray, shape (n,), int
            For each of n output samples, the calibration wavelength index
            (sorted order) whose CSD should be used.
        rng : numpy Generator

        Returns
        -------
        noise : ndarray, shape (n, seq_len, 2), float32
        """
        n = len(wl_idx)
        n_freq = self.n_freq
        seq_len = self.seq_len

        # White complex Gaussian in the rfft frequency domain.
        # At DC and Nyquist (if seq_len even), X must be real for irfft
        # to yield a real time-domain signal.
        zr = rng.standard_normal((n, n_freq, 2))
        zi = rng.standard_normal((n, n_freq, 2))
        zi[:, 0, :] = 0.0  # DC is real
        if seq_len % 2 == 0:
            zi[:, -1, :] = 0.0  # Nyquist is real
        z = zr + 1j * zi

        # Normalize so E[|z[k]|^2] = 1 per channel for every k.
        # At DC/Nyquist, z = Zr with Var 1, so |z|^2 = 1 already.
        # At interior freqs, |z|^2 has expectation 2, so scale by 1/sqrt(2).
        scale_z = np.ones(n_freq, dtype=np.float64)
        if seq_len % 2 == 0:
            scale_z[1:-1] = 1.0 / np.sqrt(2)
        else:
            scale_z[1:] = 1.0 / np.sqrt(2)
        z *= scale_z[None, :, None]

        # Gather per-sample Cholesky factor and color the noise:
        #   X[k] = sqrt(seq_len) * L[k] @ z[k]
        # so E[X X^H] = seq_len * L L^H = seq_len * CSD, matching the
        # rfft normalization in which CSD[k] = mean(|X[k]|^2) / seq_len.
        L = self._csd_chol[wl_idx]  # (n, n_freq, 2, 2), complex
        x_freq = np.einsum("pfij,pfj->pfi", L, z) * np.sqrt(seq_len)

        noise = np.fft.irfft(x_freq, n=seq_len, axis=1)
        return noise.astype(np.float32)

    # --- Diagnostic report ---

    def generate_report(
        self,
        output_path: str,
        iq_data: mx.array | np.ndarray,
        wl_ids: np.ndarray,
        rng: np.random.Generator | None = None,
        data_dir: str | None = None,
        downsample: int = 1,
    ) -> None:
        """Generate a multi-page PDF validating the synthetic pulse generator.

        All per-wavelength plots use a 2x3 grid layout (6 panels per page).

        Parameters
        ----------
        output_path : str
            Path for the output PDF.
        iq_data : array, shape (n_pulses, seq_len, 2)
            Original real data (for comparison plots).
        wl_ids : ndarray, shape (n_pulses,)
            Wavelength IDs for the real data.
        rng : numpy Generator, optional
            RNG for synthetic pulse generation. If None, uses seed 42.
        data_dir : str or None
            Path to raw npz data files for noise CSD estimation.
            When None, CSD-based diagnostic plots are skipped.
        downsample : int
            Downsampling factor applied to the windowed data.
        """
        import matplotlib
        matplotlib.use("Agg")
        matplotlib.rcParams["text.usetex"] = False
        import matplotlib.pyplot as plt
        from reportlab.lib.pagesizes import letter
        from reportlab.lib.styles import getSampleStyleSheet
        from reportlab.lib.units import inch
        from reportlab.platypus import (
            Image, PageBreak, Paragraph, SimpleDocTemplate, Spacer,
        )

        style_path = Path(__file__).parent / "venom.mplstyle"
        if style_path.exists():
            plt.style.use(str(style_path))
        matplotlib.rcParams["text.usetex"] = False  # fallback if latex unavailable

        if rng is None:
            rng = np.random.default_rng(42)

        iq_np = to_numpy(iq_data) if not isinstance(iq_data, np.ndarray) else iq_data
        tmpdir = tempfile.mkdtemp()
        styles = getSampleStyleSheet()
        title_style = styles["Title"]
        heading_style = styles["Heading2"]
        body_style = styles["BodyText"]
        elements = []
        t_us = np.arange(self.seq_len)
        freq = np.fft.rfftfreq(self.seq_len)

        # Helper: adaptive grid with n_wl subplots, leaving extras empty
        def _grid(figsize=(10, 6.5)):
            ncols = 3 if self.n_wl <= 6 else 5
            nrows = -(-self.n_wl // ncols)  # ceil division
            w = figsize[0] * ncols / 3  # scale width for wider grids
            fig, axes = plt.subplots(nrows, ncols, figsize=(w, figsize[1]))
            axes_flat = axes.flatten()
            for j in range(self.n_wl, nrows * ncols):
                axes_flat[j].set_visible(False)
            return fig, axes_flat

        elements.append(Paragraph("Synthetic Pulse Generator Report", title_style))
        elements.append(Spacer(1, 0.3 * inch))
        elements.append(Paragraph(
            f"Calibration energies: {self.n_wl} wavelengths, "
            f"E range [{self.E_min:.3f}, {self.E_max:.3f}] eV, "
            f"seq_len={self.seq_len}",
            body_style,
        ))
        elements.append(Spacer(1, 0.2 * inch))

        # ---- Page 1: Mean templates (2x3) ----
        elements.append(Paragraph("Mean Templates per Wavelength", heading_style))
        fig, ax_flat = _grid()
        for idx in range(self.n_wl):
            ax = ax_flat[idx]
            tmpl = self.templates[idx]
            ax.plot(t_us, tmpl[:, 0], label="I", alpha=0.8)
            ax.plot(t_us, tmpl[:, 1], label="Q", alpha=0.8)
            amp = np.sqrt(tmpl[:, 0] ** 2 + tmpl[:, 1] ** 2)
            ax.plot(t_us, amp, "--", label="|IQ|", color="black", alpha=0.5)
            wl = self.wl_labels_sorted[idx]
            E = self.cal_energies[idx]
            ax.set_title(f"{wl} ({E:.3f} eV)")
            ax.set_xlabel("Sample")
            if idx == 0:
                ax.set_ylabel("Amplitude (normalized)")
                ax.legend(fontsize=8)
        plt.tight_layout()
        p1_path = os.path.join(tmpdir, "templates.png")
        fig.savefig(p1_path, dpi=150)
        plt.close(fig)
        elements.append(Image(p1_path, width=7.5 * inch, height=5.0 * inch))
        elements.append(Paragraph(
            "Mean IQ templates after iterative cross-correlation alignment. "
            "Dashed black is the amplitude envelope.",
            body_style,
        ))
        elements.append(PageBreak())

        # ---- Page 2: Interpolated templates (3x3) ----
        elements.append(Paragraph("Interpolated Templates", heading_style))
        n_interp = 9
        interp_E = np.linspace(self.E_min, self.E_max, n_interp)
        fig, axes = plt.subplots(3, 3, figsize=(10, 8))
        for i, E in enumerate(interp_E):
            ax = axes.flatten()[i]
            tmpl = self.template_spline(E)
            amp = np.sqrt(tmpl[:, 0] ** 2 + tmpl[:, 1] ** 2)
            ax.plot(t_us, tmpl[:, 0], alpha=0.7, linewidth=0.8)
            ax.plot(t_us, tmpl[:, 1], alpha=0.7, linewidth=0.8)
            ax.plot(t_us, amp, "--", color="black", alpha=0.5, linewidth=0.8)
            is_cal = np.min(np.abs(self.cal_energies - E)) < 0.001
            marker = " *" if is_cal else ""
            ax.set_title(f"{E:.3f} eV{marker}", fontsize=10)
            ax.set_xlabel("Sample", fontsize=8)
        plt.tight_layout()
        p2_path = os.path.join(tmpdir, "interp_templates.png")
        fig.savefig(p2_path, dpi=150)
        plt.close(fig)
        elements.append(Image(p2_path, width=7.5 * inch, height=6.0 * inch))
        elements.append(Paragraph(
            "Templates interpolated via linear interpolation. "
            "(*) marks calibration energies.",
            body_style,
        ))
        elements.append(PageBreak())

        # ---- Page 3: Noise power spectra (requires data_dir for CSD) ----
        # Compute CSD from raw pre-trigger noise if data_dir is provided
        n_freq = self.seq_len // 2 + 1
        csds = None
        if data_dir is not None:
            csds = _compute_csd_from_raw(
                data_dir, self.wl_labels_sorted,
                n_freq, np.fft.rfftfreq(self.seq_len), downsample,
            )

        if csds is not None:
            elements.append(Paragraph("Noise Power Spectra", heading_style))
            fig, axes = plt.subplots(1, 3, figsize=(12, 4))
            for idx in range(self.n_wl):
                wl = self.wl_labels_sorted[idx]
                axes[0].semilogy(freq[1:], csds[idx, 1:, 0, 0].real,
                                 label=wl, alpha=0.8)
                axes[1].semilogy(freq[1:], csds[idx, 1:, 1, 1].real,
                                 label=wl, alpha=0.8)
                axes[2].semilogy(freq[1:], np.abs(csds[idx, 1:, 0, 1]),
                                 label=wl, alpha=0.8)
            axes[0].set_title(r"$S_{II}(f)$")
            axes[1].set_title(r"$S_{QQ}(f)$")
            axes[2].set_title(r"$|S_{IQ}(f)|$")
            for ax in axes:
                ax.set_xlabel("Frequency (cycles/sample)")
                ax.legend(fontsize=8)
            plt.tight_layout()
            p3_path = os.path.join(tmpdir, "noise_psd.png")
            fig.savefig(p3_path, dpi=150)
            plt.close(fig)
            elements.append(Image(p3_path, width=7.5 * inch, height=3.0 * inch))
            elements.append(Paragraph(
                "Auto-spectra and cross-spectral magnitude from pre-trigger "
                "noise at each calibration wavelength.",
                body_style,
            ))
            elements.append(PageBreak())

        # ---- Page 4: Synthetic vs real comparison ----
        elements.append(Paragraph("Synthetic vs Real Pulse Comparison", heading_style))
        n_synth_examples = 5
        fig, ax_flat = _grid(figsize=(10, 7))
        for idx in range(self.n_wl):
            wl = self.wl_labels_sorted[idx]
            orig_wl_id = self._sorted_to_orig_wl_id[idx]
            mask = wl_ids == orig_wl_id
            real_traces = iq_np[mask]
            real_amps = np.sqrt(
                real_traces[:, :, 0] ** 2 + real_traces[:, :, 1] ** 2
            )  # (n_pulses, seq_len)

            ax = ax_flat[idx]

            # Real data: percentile band + mean
            p5 = np.percentile(real_amps, 5, axis=0)
            p95 = np.percentile(real_amps, 95, axis=0)
            mean_amp = real_amps.mean(axis=0)
            ax.fill_between(t_us, p5, p95, color="C0", alpha=0.15,
                            label="Real 5\u201395%")
            ax.plot(t_us, mean_amp, color="C0", linewidth=1.0,
                    alpha=0.8, label="Real mean")

            # Template
            tmpl = self.templates[idx]
            tmpl_amp = np.sqrt(tmpl[:, 0] ** 2 + tmpl[:, 1] ** 2)
            ax.plot(t_us, tmpl_amp, color="black", linewidth=1.0,
                    linestyle="--", alpha=0.7, label="Template")

            # Synthetic examples: CSD noise added to this wavelength's template
            wl_idx_arr = np.full(n_synth_examples, idx, dtype=np.int64)
            synth_noise = self._sample_csd_noise(wl_idx_arr, rng)
            synth_np = (tmpl[None, :, :] + synth_noise).astype(np.float32)
            for k in range(n_synth_examples):
                amp = np.sqrt(synth_np[k, :, 0] ** 2 + synth_np[k, :, 1] ** 2)
                label = "Synthetic" if k == 0 else None
                ax.plot(t_us, amp, alpha=0.3, linewidth=0.6, color="C1",
                        label=label)

            ax.set_title(f"{wl}")
            ax.set_xlabel("Sample")
            if idx == 0:
                ax.set_ylabel("|IQ|")

        # Legend on last visible panel
        ax_flat[self.n_wl - 1].legend(fontsize=6, loc="lower right")
        plt.tight_layout()
        p5_path = os.path.join(tmpdir, "synth_vs_real.png")
        fig.savefig(p5_path, dpi=150)
        plt.close(fig)
        elements.append(Image(p5_path, width=7.5 * inch, height=5.5 * inch))
        elements.append(Paragraph(
            "Blue band: 5th\u201395th percentile of real pulse amplitudes. "
            "Blue line: real mean. Black dashed: template. "
            "Orange: 5 individual synthetic pulses.",
            body_style,
        ))
        elements.append(PageBreak())

        # ---- Page 5: Noise PSD validation (2x3) ----
        elements.append(Paragraph("Noise PSD Validation", heading_style))
        n_val = 2000
        fig, ax_flat = _grid()
        for idx in range(self.n_wl):
            wl = self.wl_labels_sorted[idx]
            E_cal = self.cal_energies[idx]
            ax = ax_flat[idx]

            # PSD from real residuals
            orig_wl_id = self._sorted_to_orig_wl_id[idx]
            real_mask = wl_ids == orig_wl_id
            real_traces_wl = iq_np[real_mask]
            real_resid = real_traces_wl - self.templates[idx]
            real_R = np.fft.rfft(real_resid, axis=1)
            real_psd_ii = np.mean(np.abs(real_R[:, :, 0]) ** 2, axis=0) / self.seq_len
            real_psd_qq = np.mean(np.abs(real_R[:, :, 1]) ** 2, axis=0) / self.seq_len

            # Synthetic "residuals" at this wavelength: fresh CSD-generated
            # noise (no template, no dataset residual reuse). PSD of this
            # should match the real residual PSD if the noise model is
            # consistent with the pre-pulse baseline statistics.
            wl_idx_arr = np.full(n_val, idx, dtype=np.int64)
            synth_resid = self._sample_csd_noise(wl_idx_arr, rng)
            synth_R = np.fft.rfft(synth_resid, axis=1)
            synth_psd_ii = np.mean(np.abs(synth_R[:, :, 0]) ** 2, axis=0) / self.seq_len
            synth_psd_qq = np.mean(np.abs(synth_R[:, :, 1]) ** 2, axis=0) / self.seq_len

            ax.semilogy(freq[1:], real_psd_ii[1:], "b-", alpha=0.7,
                        label="Real $S_{II}$", linewidth=0.8)
            ax.semilogy(freq[1:], synth_psd_ii[1:], "b--", alpha=0.7,
                        label="Synth $S_{II}$", linewidth=0.8)
            ax.semilogy(freq[1:], real_psd_qq[1:], "r-", alpha=0.7,
                        label="Real $S_{QQ}$", linewidth=0.8)
            ax.semilogy(freq[1:], synth_psd_qq[1:], "r--", alpha=0.7,
                        label="Synth $S_{QQ}$", linewidth=0.8)
            ax.set_title(f"{wl}")
            ax.set_xlabel("Freq")
            if idx == 0:
                ax.legend(fontsize=7)
                ax.set_ylabel("PSD")

        plt.tight_layout()
        p6_path = os.path.join(tmpdir, "psd_validation.png")
        fig.savefig(p6_path, dpi=150)
        plt.close(fig)
        elements.append(Image(p6_path, width=7.5 * inch, height=5.0 * inch))
        elements.append(Paragraph(
            "Solid: PSD from real residuals. Dashed: PSD from synthetic residuals.",
            body_style,
        ))

        # ---- Build PDF ----
        doc = SimpleDocTemplate(
            output_path,
            pagesize=letter,
            leftMargin=0.5 * inch,
            rightMargin=0.5 * inch,
            topMargin=0.5 * inch,
            bottomMargin=0.5 * inch,
        )
        doc.build(elements)
        print(f"  Synthetic pulse report saved: {output_path}")