#!/usr/bin/env python3
"""Train a small LEARNED 1-D audio transform codec and export its SYNTHESIS as an
INTEGER-exact decoder — the audio analog of `int_hyper`, extending "conformance is
a hash" to learned audio.

A neural audio codec (SoundStream/EnCodec-class) decodes through a FLOAT network,
so a ~1e-5 round-off makes the reconstructed waveform — and its hash — diverge
between machines: never receiptable. Here the synthesis is a stack of integer
1-D transposed convolutions + integer IGDN (per-out-channel int12 weights,
fixed-point activations, all i64), so the reconstructed PCM is byte-identical on
every machine (NeuralIntegerExact) — *sample-equivalence*, a portable PCM hash.

Pipeline (mirrors wai_sota_int_synth_export, in 1-D):
  1. train g_a/g_s (Conv1d+GDN / ConvTranspose1d+IGDN) on a real clip (float; this
     is the encoder, float is fine here — only the DECODE must be integer);
  2. quantize the latent per channel (factorized prior, data-fit CDF + rANS, the
     byte-exact match of the Rust RansDecoder);
  3. PTQ g_s → integer (WAM1 binary), run an INTEGER reference decode (numpy,
     mirroring the planned Rust i64 arithmetic) → int16 PCM → FNV-1a hash.

  ffmpeg -i <real>.aiff -ac 1 -ar 16000 -f s16le audio.s16
  python tools/wai_int_audio_export.py audio.s16 <out-dir>
"""
import sys, os, json, math, base64
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, warnings
warnings.filterwarnings('ignore')
from compressai.ans import RansEncoder

AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16   # same fixed-point as int_hyper
SR = 16000; T = 2048; C = 32; K = 5; STR = 2; NLAY = 4   # codec shape (16x downsample)
STEPS = 3500; MAXS = 0   # train steps; MAXS>0 caps #samples (small committed vector)


# ---- learned model (float; the encoder) -------------------------------------
class GDN1d(nn.Module):
    """1-D (inverse) generalized divisive normalization with a full C×C gamma.
    Forward divides by sqrt(beta + gamma·x²); inverse multiplies — the exact form
    crate::int_transform::igdn reproduces in integer (out = x·isqrt(norm))."""
    def __init__(self, ch, inverse):
        super().__init__()
        self.inverse = inverse
        self.rg = nn.Parameter(torch.randn(ch, ch) * 0.01)   # raw → softplus → γ≥0
        self.rb = nn.Parameter(torch.zeros(ch))              # raw → softplus → β≈1

    def gb(self):
        return F.softplus(self.rg), F.softplus(self.rb + 1.0)

    def forward(self, x):                                    # x [N,C,T]
        g, b = self.gb()
        norm = b[None, :, None] + torch.einsum('ij,njt->nit', g, x * x)
        return x * torch.sqrt(norm) if self.inverse else x / torch.sqrt(norm)


class Codec(nn.Module):
    def __init__(self):
        super().__init__()
        ga, gs = [], []
        for i in range(NLAY):
            ga += [nn.Conv1d(1 if i == 0 else C, C, K, STR, K // 2), GDN1d(C, False)]
        for i in range(NLAY):
            last = i == NLAY - 1
            gs += [nn.ConvTranspose1d(C, 1 if last else C, K, STR, K // 2, output_padding=STR - 1)]
            if not last:
                gs += [GDN1d(C, True)]
        self.g_a, self.g_s = nn.Sequential(*ga), nn.Sequential(*gs)

    def forward(self, x):
        y = self.g_a(x)
        y_hat = y + (torch.rand_like(y) - 0.5)          # STE quantization noise
        return self.g_s(y_hat), y


# ---- integer reference decode (numpy; mirrors the Rust i64 arithmetic) -------
def conv_t_1d_int(xi, wi, mult, bias, stride, pad, opad):
    """Integer 1-D transposed conv + per-out-channel (mult>>SHIFT)+bias — byte-for-
    byte what Rust int_audio::conv_t_1d will do. xi [cin,T]; wi [cin,cout,k]."""
    cin, cout, k = wi.shape; Tin = xi.shape[1]
    To = (Tin - 1) * stride + k + opad - 2 * pad
    Tp = To + 2 * pad
    raw = np.zeros((cout, Tp), dtype=object)
    for kx in range(k):
        contrib = np.tensordot(wi[:, :, kx], xi, axes=([0], [0]))      # [cout,Tin]
        raw[:, kx:kx + stride * Tin:stride] += contrib
    out = (raw[:, pad:pad + To] * mult[:, None]) >> SHIFT
    return out + bias[:, None]


def igdn_int_1d(x, gi, bi):
    """Integer IGDN, the EXACT arithmetic of crate::int_transform::igdn (per-term
    >>GF before the sum, isqrt at 2^AFRAC, (x·s)>>AFRAC). x [C,T] python-int."""
    Cc, Tt = x.shape
    x2 = x * x                                                          # 2^(2·AFRAC)
    norm = np.tile(np.array(bi, dtype=object)[:, None], (1, Tt))        # 2^(2·AFRAC)
    for j in range(Cc):
        norm = norm + ((gi[:, j:j + 1] * x2[j:j + 1, :]) >> GF)
    s = np.vectorize(lambda v: math.isqrt(int(v)) if v > 0 else 0, otypes=[object])(norm)
    return (x * s) >> AFRAC


def run_g_s_int(xi, layers):
    x = xi.astype(object)
    for L in layers:
        if L['type'] == 'convT':
            wi = np.array(L['weight'], dtype=object).reshape(L['in_ch'], L['out_ch'], L['k'])
            x = conv_t_1d_int(x, wi, np.array(L['mult'], dtype=object),
                              np.array(L['bias'], dtype=object), L['stride'], L['pad'], L['opad'])
        else:  # igdn
            c = L['c']
            gi = np.array(L['gamma'], dtype=object).reshape(c, c)
            x = igdn_int_1d(x, gi, np.array(L['beta'], dtype=object))
    return x                                                            # [1,T] @ 2^AFRAC


def to_pcm_i16(x_fixed):
    """[1,T] @ 2^AFRAC (≈[-1,1]) → int16 PCM, round-half-up, clamp. Matches Rust."""
    v = x_fixed[0]
    num = v * 32767 + (1 << (AFRAC - 1))
    pcm = num >> AFRAC                                                  # python>> floors (== i64 arith shift)
    return np.array([max(-32768, min(32767, int(p))) for p in pcm], dtype=np.int16)


# ---- PTQ + entropy helpers ---------------------------------------------------
def q_weight_pc(w):                          # ConvTranspose1d [in,out,k] → per-OUT int12
    s = np.abs(w).max(axis=(0, 2), keepdims=True) / (2 ** (WBITS - 1) - 1) + 1e-12
    wi = np.round(w / s).clip(-(2 ** (WBITS - 1)), 2 ** (WBITS - 1) - 1).astype(np.int64)
    return wi, s.reshape(-1)


def fit_cdf(sym, lo, hi):
    n = hi - lo + 1
    hist = np.bincount(np.asarray(sym, np.int64) - lo, minlength=n).astype(np.float64)
    total = 1 << 16
    f = np.maximum((hist / hist.sum() * (total - n)).astype(np.int64), 1)
    f[f.argmax()] += (total - 1) - int(f.sum())
    cdf = [0]
    for v in f:
        cdf.append(cdf[-1] + int(v))
    cdf.append(total)
    return cdf


def fnv1a(buf):
    h = 14695981039346656037
    for b in buf:
        h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
    return h


def main(src, outdir):
    os.makedirs(outdir, exist_ok=True)
    torch.manual_seed(0)
    raw = np.fromfile(src, np.int16).astype(np.float32) / 32768.0
    if MAXS:
        raw = raw[:MAXS]
    n = (len(raw) // T) * T
    raw = raw[:n]
    frames = torch.from_numpy(raw.reshape(-1, 1, T))             # [Nf,1,T]

    m = Codec()
    opt = torch.optim.Adam(m.parameters(), 1.5e-3)
    # Keep the rate term tiny: the latent must stay O(1+) so integer rounding (the
    # quantizer) preserves it — too much rate pressure collapses it below the grid.
    for step in range(STEPS):
        xr, y = m(frames)
        loss = F.mse_loss(xr, frames) + 8e-6 * y.abs().mean()
        opt.zero_grad(); loss.backward(); opt.step()
    m.eval()
    with torch.no_grad():
        xr_f, _ = m(frames)
        snr_f = 10 * np.log10((frames ** 2).sum() / ((xr_f - frames) ** 2).sum() + 1e-12)
    print(f'  (float recon SNR {float(snr_f):.1f} dB)')

    # full-signal latent → per-channel factorized symbols (s = round(y - median))
    with torch.no_grad():
        xfull = torch.from_numpy(raw.reshape(1, 1, n))
        y = m.g_a(xfull)[0].numpy()                              # [C, L]
    L = y.shape[1]
    median = np.round(np.median(y, axis=1))                      # integer-grid median
    sym = np.round(y - median[:, None]).astype(np.int64)        # [C, L] symbols
    median_q = np.round(median * (2 ** AFRAC)).astype(np.int64)
    xi = (sym * (2 ** AFRAC) + median_q[:, None]).astype(np.int64)   # fixed-point latent

    # PTQ g_s → WAM1 integer model
    layers = []
    for layer in m.g_s.children():
        t = type(layer).__name__
        if t == 'ConvTranspose1d':
            wi, sc = q_weight_pc(layer.weight.detach().numpy())
            layers.append({'type': 'convT', 'in_ch': int(wi.shape[0]), 'out_ch': int(wi.shape[1]),
                           'k': int(wi.shape[2]), 'stride': int(layer.stride[0]), 'pad': int(layer.padding[0]),
                           'opad': int(layer.output_padding[0]), 'weight': wi.reshape(-1).tolist(),
                           'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
                           'bias': np.round(layer.bias.detach().numpy() * (2 ** AFRAC)).astype(np.int64).tolist()})
        else:  # GDN1d (inverse)
            with torch.no_grad():
                g, b = layer.gb()
            gi = np.round(g.numpy() * (2 ** GF)).astype(np.int64)
            bi = np.round(b.numpy() * (2 ** (2 * AFRAC))).astype(np.int64)
            layers.append({'type': 'igdn', 'c': int(b.shape[0]), 'gamma': gi.reshape(-1).tolist(), 'beta': bi.tolist()})

    # INTEGER reference decode → PCM (what the Rust port reproduces)
    x_fixed = run_g_s_int(xi, layers)
    pcm = to_pcm_i16(x_fixed)
    snr = 10 * np.log10((raw[:len(pcm)] ** 2).sum() / (((pcm.astype(np.float64) / 32768) - raw[:len(pcm)]) ** 2).sum() + 1e-12)

    # factorized rANS over the symbols (channel-major), per-channel data-fit CDF
    lo, hi = int(sym.min()), int(sym.max())
    cdfs, syms_flat, idx = [], [], []
    for ch in range(C):
        cdfs.append(fit_cdf(sym[ch], lo, hi))
        syms_flat.extend(sym[ch].tolist())
        idx.extend([ch] * L)
    enc = RansEncoder()
    bits = enc.encode_with_indexes(syms_flat, idx, cdfs, [len(cdfs[0])] * C, [lo] * C)

    # WAM1 binary model
    import struct
    mb = bytearray(b'WAM1')
    mb += struct.pack('<IIII', AFRAC, GF, SHIFT, len(layers))
    for Lr in layers:
        if Lr['type'] == 'convT':
            mb += struct.pack('<BIIIIII', 0, Lr['in_ch'], Lr['out_ch'], Lr['k'], Lr['stride'], Lr['pad'], Lr['opad'])
            mb += np.array(Lr['weight'], dtype='<i2').tobytes()
            mb += np.array(Lr['mult'], dtype='<i8').tobytes() + np.array(Lr['bias'], dtype='<i8').tobytes()
        else:
            c = Lr['c']
            mb += struct.pack('<BI', 2, c)
            assert np.abs(np.array(Lr['gamma'])).max() < 2 ** 31, 'gamma > i32'
            mb += np.array(Lr['gamma'], dtype='<i4').tobytes() + np.array(Lr['beta'], dtype='<i8').tobytes()
    open(f'{outdir}/audio_model.bin', 'wb').write(mb)

    json.dump({'n_ch': C, 'lat_len': L, 't_out': int(len(pcm)), 'afrac': AFRAC,
               'lo': lo, 'cdfs': cdfs, 'median_q': median_q.tolist(),
               'rans_b64': base64.b64encode(bytes(bits)).decode()},
              open(f'{outdir}/audio_bitstream.json', 'w'))
    pcm_hash = format(fnv1a(pcm.tobytes()), '016x')
    json.dump({'pcm_fnv1a': pcm_hash, 'n_samples': int(len(pcm)), 'sample_rate': SR},
              open(f'{outdir}/audio.json', 'w'))
    pcm.tofile(f'{outdir}/audio_decoded.s16')
    print(f'learned 1-D audio codec: {len(pcm)} samples @ {SR}Hz, C={C} L={L}, model {len(mb)}B, '
          f'rANS {len(bits)}B  ->  integer-decode SNR {snr:.1f} dB, pcm_hash {pcm_hash}')


if __name__ == '__main__':
    # optional: C  MAXS  STEPS  (e.g. tiny committed vector: 16 8192 2500)
    if len(sys.argv) > 3:
        C = int(sys.argv[3])
    if len(sys.argv) > 4:
        MAXS = int(sys.argv[4])
    if len(sys.argv) > 5:
        STEPS = int(sys.argv[5])
    main(sys.argv[1], sys.argv[2])