#!/usr/bin/env python3
"""
Generate "Noisy Residual" and "Denoised Residual" curves as SVGs.

- Produces TWO separate SVG files:
    noisy_residual.svg
    denoised_residual.svg

- Curves are synthetic but shaped like residual noise + denoised residual.
- Uses only matplotlib + numpy.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt


@dataclass
class CurveParams:
    seconds: float = 12.0          # length of the signal
    fs: int = 250                  # samples per second
    seed: int = 7                  # RNG seed for reproducibility
    base_amp: float = 0.12         # smooth baseline amplitude
    noise_amp: float = 0.55        # high-frequency noise amplitude
    burst_amp: float = 1.2         # occasional spike amplitude
    burst_rate_hz: float = 0.35    # average spike frequency
    denoise_smooth_ms: float = 120 # smoothing window for "denoised" (ms)


def gaussian_smooth(x: np.ndarray, sigma_samples: float) -> np.ndarray:
    """Gaussian smoothing using explicit kernel convolution (no SciPy dependency)."""
    if sigma_samples <= 0:
        return x.copy()

    radius = int(np.ceil(4 * sigma_samples))
    k = np.arange(-radius, radius + 1, dtype=float)
    kernel = np.exp(-(k**2) / (2 * sigma_samples**2))
    kernel /= kernel.sum()
    return np.convolve(x, kernel, mode="same")


def make_residual(params: CurveParams) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Create synthetic residual:
    - baseline: smooth wavy trend + slight drift
    - noise: band-limited-ish high-frequency noise
    - bursts: sparse spikes / impulse-like events
    Returns: (t, noisy, denoised)
    """
    rng = np.random.default_rng(params.seed)
    n = int(params.seconds * params.fs)
    t = np.linspace(0, params.seconds, n, endpoint=False)

    # Smooth baseline (small): combination of sinusoids + small random drift
    baseline = (
        0.7 * np.sin(2 * np.pi * 0.35 * t + 0.2)
        + 0.35 * np.sin(2 * np.pi * 0.9 * t + 1.2)
        + 0.25 * np.sin(2 * np.pi * 0.15 * t + 2.0)
    )
    baseline *= params.base_amp
    drift = np.cumsum(rng.normal(0, 1, size=n))
    drift = drift / (np.max(np.abs(drift)) + 1e-9) * (params.base_amp * 0.25)
    baseline = baseline + drift

    # High-frequency noise: whitened then lightly smoothed to look "oscillatory"
    raw = rng.normal(0, 1, size=n)
    hf = raw - gaussian_smooth(raw, sigma_samples=params.fs * 0.03)  # remove slow part
    hf = hf / (np.std(hf) + 1e-9)
    hf *= params.noise_amp

    # Bursts/spikes: Poisson process impulses convolved with short kernel
    expected_bursts = params.burst_rate_hz * params.seconds
    k_bursts = rng.poisson(expected_bursts)
    impulses = np.zeros(n)
    if k_bursts > 0:
        idx = rng.integers(0, n, size=k_bursts)
        impulses[idx] = rng.normal(loc=1.0, scale=0.4, size=k_bursts)

    # Shape impulses into spikes (asymmetric bump)
    spike_kernel_len = int(params.fs * 0.06)  # ~60ms
    spike_kernel_len = max(spike_kernel_len, 7)
    spike_t = np.arange(spike_kernel_len)
    spike_kernel = np.exp(-spike_t / (params.fs * 0.012))  # fast decay
    spike_kernel *= np.hanning(spike_kernel_len)  # taper
    spike_kernel /= (spike_kernel.max() + 1e-9)

    bursts = np.convolve(impulses, spike_kernel, mode="same")
    bursts *= params.burst_amp

    noisy = baseline + hf + bursts

    # "Denoised": remove high-frequency using Gaussian smoothing,
    # but keep spike structures partially.
    smooth_sigma = (params.denoise_smooth_ms / 1000.0) * params.fs / 3.0
    denoised = gaussian_smooth(noisy, sigma_samples=smooth_sigma)

    return t, noisy, denoised


def save_curve_svg(
    t: np.ndarray,
    y: np.ndarray,
    out_path: Path,
    *,
    width_in: float = 5.4,
    height_in: float = 1.6,
    lw: float = 2.2,
    pad: float = 0.03,
) -> None:
    """
    Save a clean, figure-only SVG suitable for embedding in diagrams.
    - No axes, ticks, labels.
    - Tight bounding box.
    """
    fig = plt.figure(figsize=(width_in, height_in), dpi=200)
    ax = fig.add_axes([pad, pad, 1 - 2 * pad, 1 - 2 * pad])

    ax.plot(t, y, linewidth=lw)

    # Make it "icon-like" for diagrams: no axes or frames
    ax.set_axis_off()

    # Ensure bounds include a little padding
    ymin, ymax = np.min(y), np.max(y)
    ypad = 0.08 * (ymax - ymin + 1e-9)
    ax.set_xlim(t[0], t[-1])
    ax.set_ylim(ymin - ypad, ymax + ypad)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, format="svg", bbox_inches="tight", pad_inches=0.0)
    plt.close(fig)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--outdir", type=Path, default=Path("."), help="Output directory")
    ap.add_argument("--seed", type=int, default=7, help="RNG seed")
    ap.add_argument("--seconds", type=float, default=12.0, help="Signal length (s)")
    ap.add_argument("--fs", type=int, default=250, help="Sampling rate (Hz)")
    ap.add_argument("--prefix", type=str, default="", help="Filename prefix (optional)")
    args = ap.parse_args()

    params = CurveParams(seconds=args.seconds, fs=args.fs, seed=args.seed)
    t, noisy, denoised = make_residual(params)

    noisy_path = args.outdir / f"{args.prefix}noisy_residual.svg"
    den_path = args.outdir / f"{args.prefix}denoised_residual.svg"

    save_curve_svg(t, noisy, noisy_path)
    save_curve_svg(t, denoised, den_path)

    print(f"Wrote:\n  {noisy_path}\n  {den_path}")


if __name__ == "__main__":
    main()