"""Export the REAL pretrained SNAC model to the ONNX `wai.neural.snac` decoder
(`wai-rs/src/neural/snac.rs`), via the model surgery needed to make it both
exportable AND deterministic.

Two upstream blockers, both patched here (monkey-patch, no fork):
  1. NoiseBlock injects `torch.randn` — that makes the decode STOCHASTIC (a codec
     must be reproducible) AND bakes a bad shape on ONNX export. We zero the
     noise: `x + 0*linear(x)` → deterministic, exportable. (Quality cost is the
     noise texture only.)
  2. Snake1d's `@torch.jit.script snake` does a reshape-roundtrip that ONNX
     mis-traces. We replace it with the identical 3D formula, no reshape.

Verified: snac_24khz exports clean and the ONNX matches the (now deterministic)
torch decode to ~9e-7; `wai-rs` then decodes the resulting envelope to ~9e-7 of
that reference (`snac_real_model_matches_reference`).

Run (needs torch + snac + onnxruntime):
  python tools/wai_snac_export.py --model hubertsiuzdak/snac_24khz \
    --out-dir wai-web/demo/models/snac_24khz
"""
from __future__ import annotations

import argparse
from pathlib import Path

import torch
import torch.nn as nn


def patch_for_export():
    import snac.layers
    # 1. deterministic, export-clean noise.
    def noise_fwd(self, x):
        return x + 0.0 * self.linear(x)
    snac.layers.NoiseBlock.forward = noise_fwd
    # 2. reshape-free snake.
    def snake1d_fwd(self, x):
        return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x) ** 2
    snac.layers.Snake1d.forward = snake1d_fwd


class _Dec(nn.Module):
    def __init__(s, m): super().__init__(); s.m = m
    def forward(s, *codes):
        return s.m.decoder(s.m.quantizer.from_codes(list(codes)))


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="hubertsiuzdak/snac_24khz")
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--opset", type=int, default=17)
    args = ap.parse_args()

    patch_for_export()
    from snac import SNAC
    m = SNAC.from_pretrained(args.model).eval()
    out = Path(args.out_dir); out.mkdir(parents=True, exist_ok=True)

    # Probe the codebook hierarchy (N levels, their strides) from a dummy encode.
    dummy_audio = torch.zeros(1, 1, m.sampling_rate)
    with torch.no_grad():
        codes = m.encode(dummy_audio)
    n = len(codes)
    names = [f"codes_{i}" for i in range(n)]
    torch.onnx.export(
        _Dec(m).eval(), tuple(codes), str(out / "decoder.onnx"),
        input_names=names, output_names=["audio_values"],
        dynamic_axes={**{f"codes_{i}": {1: f"t{i}"} for i in range(n)}, "audio_values": {2: "s"}},
        opset_version=args.opset, dynamo=False)
    print(f"exported {args.model}: {n} code levels @ {m.sampling_rate} Hz → {out}/decoder.onnx")
    print("decode is deterministic (zero-noise). Encode tool: pack <III> sr,n_samples,n_levels "
          "+ n×u32 lengths + per-level u16 codes (see snac.rs).")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())