#!/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 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.

**The signal is generated here; the tool reads no audio file.** Both the training
signal and the test signal (the clip the vector codes) are written by `signal()`
below from a seed: 16 kHz, mono, signed 16-bit, a sequence of segments, each a tone
(one to three partials with an attack and a release), a chirp (a linear frequency
sweep), a noise burst (white noise, low-passed by a moving average or high-passed by
a difference) or a speech-like formant sweep (a harmonic series on a gliding,
vibrato-modulated pitch, shaped by three formant resonances that move between two
vowel targets), sometimes two overlaid, with short gaps. The generator is integer
arithmetic only — a SplitMix64 counter stream, a 32-bit phase accumulator, a
4096-entry sine table built by Bhaskara I's rational approximation — so a seed gives
the same samples on every machine and every Python. The training signal is stream 0
of the seed and the test signal stream 1: the codec is trained on one signal and
codes another of the same kind.

Pipeline (mirrors wai_sota_int_synth_export, in 1-D):
  1. generate the training and test signals from the seed;
  2. train g_a/g_s (Conv1d+GDN / ConvTranspose1d+IGDN) on the training signal (float;
     this is the encoder, float is fine here — only the DECODE must be integer), with
     torch seeded from the same seed, deterministic algorithms and a fixed thread
     count, so a run on one host and torch build gives the same weights every time;
  3. quantize the test signal's latent per channel (factorized prior, data-fit CDF +
     rANS, the byte-exact match of the Rust RansDecoder);
  4. PTQ g_s → integer (WAM1 binary), run an INTEGER reference decode (numpy,
     mirroring the Rust i64 arithmetic) → int16 PCM → FNV-1a hash.

  python tools/wai_int_audio_export.py <out-dir> [--seed S] [--channels C]
                                       [--samples N] [--steps K] [--train-seconds T]
  python tools/wai_int_audio_export.py --signal-only <out.s16> [--seed S] [--stream 0|1]
                                       [--samples N]

`--signal-only` writes one generated signal as raw s16le and needs nothing beyond the
Python standard library. The committed vector (`wai-rs/tests/vectors/int_audio/`) and
the site's demo clip are both exports of this tool; their commands are in that
directory's README.
"""
import argparse
import base64
import json
import math
import os
import struct
import sys

AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16   # same fixed-point as int_hyper
SR = 16000; T = 2048; K = 5; STR = 2; NLAY = 4   # codec shape (16x downsample)
LR = 1.5e-3; RATE_W = 8e-6
THREADS = 4                                      # fixed: CPU reductions depend on it

# ---- the generated signal (integer arithmetic only) ----------------------------------
M64 = (1 << 64) - 1
GOLDEN = 0x9E3779B97F4A7C15
TAB_BITS = 12
TAB_N = 1 << TAB_BITS


def _mix64(x):
    """SplitMix64's finalizer."""
    x &= M64
    x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & M64
    x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & M64
    return x ^ (x >> 31)


class Rng:
    """A SplitMix64 stream, keyed by (seed, stream)."""

    def __init__(self, seed, stream):
        self.s = _mix64(seed * GOLDEN + stream * 0xD1B54A32D192ED03)

    def next(self):
        self.s = (self.s + GOLDEN) & M64
        return _mix64(self.s)

    def int(self, lo, hi):
        """Uniform integer in [lo, hi]."""
        return lo + self.next() % (hi - lo + 1)

    def s16(self):
        """Uniform signed 16-bit value."""
        return (self.next() >> 48) - 32768


def _sine_table():
    """sin(2π i / TAB_N) × 32767, by Bhaskara I's approximation, rounded half up:
    sin(πp) ≈ 16 p (1 − p) / (5 − 4 p (1 − p)), p = u / half."""
    half = TAB_N // 2
    out = []
    for i in range(TAB_N):
        u = i % half
        a = u * (half - u)
        den = 5 * half * half - 4 * a
        v = (2 * 32767 * 16 * a + den) // (2 * den)
        out.append(v if i < half else -v)
    return out


SINE = _sine_table()


def _inc(f_q8):
    """32-bit phase increment per sample for a frequency in Hz × 256."""
    return (f_q8 << 24) // SR


def _osc(phase):
    return SINE[(phase & 0xFFFFFFFF) >> (32 - TAB_BITS)]


def _envelope(n, attack, release):
    """Linear attack and release, Q15, length n."""
    attack, release = max(1, min(attack, n // 2)), max(1, min(release, n // 2))
    return [min(32767, (32767 * (i + 1)) // attack, (32767 * (n - i)) // release) for i in range(n)]


def _root64(x):
    """2^(x/64) in Q16, floored: the largest y with y^64 ≤ 2^(x + 16·64). Exact."""
    lo, hi = 1 << 16, 1 << 17
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if mid ** 64 <= 1 << (x + 1024):
            lo = mid
        else:
            hi = mid - 1
    return lo


ROOT64 = [_root64(x) for x in range(64)]


def _log_freq(r, lo, hi):
    """A frequency in Hz × 256, log-uniform in [lo, hi] Hz: lo · 2^(k/64) for a random
    k, drawn again while it is above hi."""
    octaves = 0
    while (lo << octaves) < hi:
        octaves += 1
    while True:
        o, j = divmod(r.int(0, 64 * octaves), 64)
        f = ((lo << o) * ROOT64[j]) >> 8
        if f <= hi << 8:
            return f


# Formant targets (F1, F2, F3) in Hz for five vowels.
VOWELS = [(730, 1090, 2440), (270, 2290, 3010), (300, 870, 2240), (530, 1840, 2480), (570, 840, 2410)]


def _tone(r, n):
    out = [0] * n
    for _ in range(r.int(1, 3)):
        inc, amp, ph = _inc(_log_freq(r, 80, 3500)), r.int(4000, 16000), r.next() & 0xFFFFFFFF
        for i in range(n):
            out[i] += (amp * _osc(ph)) >> 15
            ph += inc
    env = _envelope(n, r.int(SR // 500, SR // 30), r.int(SR // 50, SR // 6))
    return [(v * e) >> 15 for v, e in zip(out, env)]


def _chirp(r, n):
    i0, i1 = _inc(_log_freq(r, 100, 6000)), _inc(_log_freq(r, 100, 6000))
    amp, ph, out = r.int(8000, 20000), r.next() & 0xFFFFFFFF, []
    for i in range(n):
        out.append((amp * _osc(ph)) >> 15)
        ph += i0 + (i1 - i0) * i // n
    env = _envelope(n, r.int(SR // 500, SR // 50), r.int(SR // 100, SR // 20))
    return [(v * e) >> 15 for v, e in zip(out, env)]


def _noise(r, n):
    w = [r.s16() for _ in range(n + 16)]
    kind = r.int(0, 2)
    if kind == 0:                                       # white
        x = w[:n]
    elif kind == 1:                                     # low-passed: moving average
        L = [2, 3, 5, 8, 13][r.int(0, 4)]
        x = [sum(w[i:i + L]) // L for i in range(n)]
    else:                                               # high-passed: first difference
        x = [(w[i + 1] - w[i]) >> 1 for i in range(n)]
    env = _envelope(n, r.int(SR // 1000, SR // 200), r.int(SR // 50, SR // 5))
    amp = r.int(4000, 14000)
    return [(((v * amp) >> 15) * e) >> 15 for v, e in zip(x, env)]


def _formant(r, n):
    """A harmonic series on a gliding pitch with vibrato, each harmonic weighted by
    three resonances 1 / (1 + ((f − F)/B)^2) whose centres move between two vowels."""
    f0a = r.int(90, 240) << 8
    f0b = (f0a * r.int(80, 125)) // 100
    va, vb = VOWELS[r.int(0, 4)], VOWELS[r.int(0, 4)]
    bw = (r.int(60, 110), r.int(90, 160), r.int(120, 220))
    vib_inc, vib_depth = _inc(r.int(4, 7) << 8), r.int(5, 25)      # depth in ‰
    nh = max(1, (4800 << 8) // max(f0a, f0b))
    ph0, vph, out = r.next() & 0xFFFFFFFF, r.next() & 0xFFFFFFFF, []
    amps = [0] * nh
    for i in range(n):
        f0 = f0a + (f0b - f0a) * i // n
        f0 += (f0 * vib_depth * _osc(vph)) // (1000 * 32767)
        vph += vib_inc
        if i % 32 == 0:                                 # the weights at block rate
            F = [va[j] + (vb[j] - va[j]) * i // n for j in range(3)]
            for k in range(nh):
                fk = ((k + 1) * f0) >> 8
                g = 0
                for j in range(3):
                    d = ((fk - F[j]) << 10) // bw[j]
                    g += (1 << 20) // (1024 + ((d * d) >> 10))   # Q10 resonance gain
                amps[k] = (g * 1024) // (1024 + 256 * k)          # spectral tilt
        ph0 += _inc(f0)
        acc = 0
        for k in range(nh):
            acc += amps[k] * _osc((k + 1) * ph0)
        out.append(acc >> 15)
    env = _envelope(n, r.int(SR // 200, SR // 20), r.int(SR // 50, SR // 10))
    return [(v * e) >> 15 for v, e in zip(out, env)]


KINDS = [_tone, _chirp, _noise, _formant]


def signal(seed, stream, n):
    """n samples of the generated signal (a list of ints in the int16 range)."""
    r = Rng(seed, stream)
    out = []
    while len(out) < n:
        m = r.int(SR * 8 // 100, SR // 2)               # a segment: 80–500 ms
        seg = KINDS[r.int(0, 3)](r, m)
        if r.int(0, 3) == 0:                            # sometimes a second, quieter kind
            over = KINDS[r.int(0, 3)](r, m)
            seg = [a + (b >> 2) for a, b in zip(seg, over)]
        peak = max(1, max(abs(v) for v in seg))
        target = r.int(6000, 28000)
        seg = [(v * target) // peak for v in seg]
        out += seg + [0] * r.int(0, SR // 25)           # a gap: 0–40 ms
    return [max(-32768, min(32767, v)) for v in out[:n]]


def s16le(x):
    return struct.pack(f'<{len(x)}h', *x)


# ---- learned model (float; the encoder) -------------------------------------
def build_codec(C):
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    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

    return Codec()


# ---- 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 does. xi [cin,T]; wi [cin,cout,k]."""
    import numpy as np
    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."""
    import numpy as np
    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):
    import numpy as np
    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."""
    import numpy as np
    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
    import numpy as np
    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):
    import numpy as np
    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(outdir, seed, C, nsamp, steps, train_seconds):
    import numpy as np
    import torch
    import torch.nn.functional as F
    from compressai.ans import RansEncoder

    assert nsamp % T == 0, f'--samples must be a multiple of {T}'
    os.makedirs(outdir, exist_ok=True)
    torch.manual_seed(seed)
    torch.use_deterministic_algorithms(True)
    torch.set_num_threads(THREADS)

    ntrain = (train_seconds * SR // T) * T
    train = np.array(signal(seed, 0, ntrain), np.float32) / 32768.0
    test_i16 = signal(seed, 1, nsamp)
    raw = np.array(test_i16, np.float32) / 32768.0
    frames = torch.from_numpy(train.reshape(-1, 1, T))           # [Nf,1,T]

    m = build_codec(C)
    opt = torch.optim.Adam(m.parameters(), LR)
    # 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) + RATE_W * 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 on the training signal {float(snr_f):.1f} dB)')

    # the test signal's latent → per-channel factorized symbols (s = round(y - median))
    with torch.no_grad():
        xfull = torch.from_numpy(raw.reshape(1, 1, nsamp))
        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
    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')
    open(f'{outdir}/audio_source.s16', 'wb').write(s16le(test_i16))
    print(f'learned 1-D audio codec (seed {seed}): {len(pcm)} samples @ {SR}Hz, C={C} L={L}, '
          f'model {len(mb)}B, rANS {len(bits)}B  ->  integer-decode SNR {snr:.1f} dB, pcm_hash {pcm_hash}')


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description='Train and export the integer learned audio codec on a '
                                             'generated signal (no audio input).')
    ap.add_argument('outdir', nargs='?', help='where audio_model.bin, audio_bitstream.json and audio.json go')
    ap.add_argument('--seed', type=int, default=1)
    ap.add_argument('--channels', type=int, default=32, help='latent channels C')
    ap.add_argument('--samples', type=int, default=32768, help='test-signal (coded clip) length, a multiple of 2048')
    ap.add_argument('--steps', type=int, default=3500, help='training steps')
    ap.add_argument('--train-seconds', type=int, default=8, help='training-signal length in seconds')
    ap.add_argument('--signal-only', metavar='OUT.s16', help='write one generated signal and exit')
    ap.add_argument('--stream', type=int, default=1, help='with --signal-only: 0 training, 1 test')
    a = ap.parse_args()
    if a.signal_only:
        n = a.samples if a.stream == 1 else (a.train_seconds * SR // T) * T
        open(a.signal_only, 'wb').write(s16le(signal(a.seed, a.stream, n)))
        sys.exit(0)
    if not a.outdir:
        ap.error('outdir is required unless --signal-only is given')
    main(a.outdir, a.seed, a.channels, a.samples, a.steps, a.train_seconds)