"""
Honest test of the "parametric/16D gives 583x on a GW ringdown" claim.

Design (steelman for the parametric method):
- signal = sum of damped sinusoids (a real ringdown model, few QNM modes)
- add detector noise at several SNRs (clean -> realistic loud -> weak)
- quantize to int16 (real strain/ADC is finite precision) -> THIS is the data
- everything is LOSSLESS: each method must reproduce the int16 stream exactly,
  so its size = model params + entropy of the residual it leaves behind.
- the parametric coder is HANDED THE TRUE (A, tau, f, phi) — the best case it
  could ever achieve (no fit error). If it still can't beat LPC once noise is
  present, the compression was living entirely in the zero-noise self-fit.

Baselines on the same int16 array:
  raw 16 b/s, order-0 entropy, gzip -9, xz -9, and LPC (order-p linear predictor,
  residual entropy) — LPC is the 50-year-old codec that already does
  "few coefficients + residual" for decaying/oscillatory signals.
"""
import numpy as np, gzip, lzma, math

rng = np.random.default_rng(0)
fs = 4096.0
N = 2000
t = np.arange(N) / fs

# ---- true ringdown: 2 QNM modes (f Hz, tau s, amp, phase) ----
MODES = [(250.0, 0.004, 1.0, 0.0), (500.0, 0.002, 0.35, 1.1)]
def model(t):
    s = np.zeros_like(t)
    for f, tau, A, ph in MODES:
        s += A * np.exp(-t / tau) * np.cos(2 * np.pi * f * t + ph)
    return s
clean = model(t)
peak = np.max(np.abs(clean))

def ideal_bits_per_sample(residual_int):
    """entropy (bits/sample) an arithmetic coder would reach on integer residuals."""
    vals, counts = np.unique(residual_int, return_counts=True)
    p = counts / counts.sum()
    return float(-(p * np.log2(p)).sum())

def quantize_int16(x):
    scale = 30000.0 / peak
    return np.round(x * scale).astype(np.int64), scale

def size_gzip(int16):
    return len(gzip.compress(int16.astype('<i2').tobytes(), 9)) * 8 / len(int16)
def size_xz(int16):
    return len(lzma.compress(int16.astype('<i2').tobytes(), preset=9)) * 8 / len(int16)

def lpc_bits(x, p=8):
    """order-p LPC: least-squares predictor, residual entropy. +p coeff cost."""
    X = np.zeros((len(x) - p, p))
    for i in range(p):
        X[:, i] = x[p - 1 - i: len(x) - 1 - i]
    y = x[p:].astype(float)
    a, *_ = np.linalg.lstsq(X, y, rcond=None)
    pred = np.rint(X @ a).astype(np.int64)
    resid = x[p:] - pred
    warmup = x[:p]
    H = ideal_bits_per_sample(np.concatenate([warmup, resid]))
    coeff_bits = p * 32
    return H + coeff_bits / len(x)

def parametric_bits(x_int, scale):
    """BEST CASE: hand it the TRUE modes. residual = data - round(scale*model)."""
    m_int = np.round(model(t) * scale).astype(np.int64)
    resid = x_int - m_int
    H = ideal_bits_per_sample(resid)
    # store: per mode (f, tau, A, phi) as float32 = 4*32 bits, + scale
    param_bits = (len(MODES) * 4 + 1) * 32
    return H + param_bits / len(x_int)

print(f"ringdown: {N} samples @ {fs:.0f} Hz, {len(MODES)} true QNM modes, int16 quantized\n")
print(f"{'SNR regime':<22}{'raw':>6}{'order0':>8}{'gzip':>7}{'xz':>7}{'LPC-8':>8}{'PARAM*':>8}   note")
print("-" * 86)

regimes = [
    ("clean (sigma=0)",        0.0),
    ("60 dB (sigma=peak/1e3)", peak / 1000),
    ("40 dB (sigma=peak/1e2)", peak / 100),
    ("30 dB realistic loud",   peak / 31.6),
    ("20 dB weak",             peak / 10),
]
for name, sigma in regimes:
    noisy = clean + (rng.standard_normal(N) * sigma if sigma > 0 else 0.0)
    xi, scale = quantize_int16(noisy)
    raw = 16.0
    h0 = ideal_bits_per_sample(xi)
    g = size_gzip(xi); z = size_xz(xi)
    lpc = lpc_bits(xi, 8)
    par = parametric_bits(xi, scale)
    winner = min(("gzip", g), ("xz", z), ("LPC", lpc), ("PARAM", par), key=lambda kv: kv[1])
    print(f"{name:<22}{raw:6.1f}{h0:8.2f}{g:7.2f}{z:7.2f}{lpc:8.2f}{par:8.2f}   "
          f"ratio(raw/best)={raw/winner[1]:5.1f}x via {winner[0]}")

print("\n* PARAM is handed the TRUE mode parameters (upper bound on any 16D/braid fit).")
print("Lossless: every column reproduces the int16 stream exactly; size = params + residual entropy.")