#!/usr/bin/env python3
"""Train a learned INTEGER-exact Gaussian-splat ATTRIBUTE codec — "the int_hyper of
3D". A Gaussian splat is N Gaussians × 14 attributes (pos[3], scale[3], rot[4],
color[3], opacity) — already a fixed-point (Fx i64) tensor in WAI. A decode that runs
float arithmetic can round differently on different machines, and then the
reconstructed cloud — and its hash — differ between them. Here the synthesis is integer (the
int_hyper machinery, 14 channels instead of RGB's 3): the decoded attribute tensor is
byte-identical on every machine — **splat-equivalence**, a portable cloud hash. The
float rasterizer is sink-supplied presentation; the *attributes* are the contract.

Pipeline (mirrors the audio/world tools): arrange the N Gaussians as a [14, S, S]
attribute image (Morton-sorted by position for locality), train a small transform
codec (Conv2d+GDN / ConvTranspose2d+IGDN), factorized-rANS the latent, PTQ g_s → WHM1,
run an INTEGER reference decode (numpy, mirroring the Rust int_hyper conv/IGDN exactly)
→ the byte-exact attribute tensor → FNV-1a. Reuses the int_hyper 2D integer ops.

  python tools/wai_int_splat_export.py <out-dir>
"""
import sys, os, json, base64
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, warnings
warnings.filterwarnings('ignore')
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wai_sota_int_synth_export import q_weight_pc, q_weight_conv, conv_int, conv_t_raw, igdn_int, run_int_net, AFRAC, GF, SHIFT
from wai_int_world_export import whm1, ptq_conv, ptq_convT, ptq_igdn, GDN2d
from compressai.ans import RansEncoder

S, C, ATTR = 32, 32, 14            # grid side (N=S²=1024 gaussians) / latent channels / attributes
ONE = 1 << AFRAC
NQ = 6                              # attribute quantization step (coarser Fx → compressible, byte-exact)


def morton(xy):
    """Interleave low bits of (ix,iy) → a Z-order key for spatial locality."""
    def part(n):
        n &= 0xFFFF
        n = (n | (n << 8)) & 0x00FF00FF
        n = (n | (n << 4)) & 0x0F0F0F0F
        n = (n | (n << 2)) & 0x33333333
        n = (n | (n << 1)) & 0x55555555
        return n
    return part(xy[0]) | (part(xy[1]) << 1)


def make_cloud(n):
    """A synthetic but structured Gaussian cloud: points on a wavy shell, scales /
    colors / opacity varying smoothly with position — real spatial structure the
    transform can exploit. Returns [N, 14] float in natural ranges."""
    rng = np.random.RandomState(3)
    u, v = rng.rand(n), rng.rand(n)
    th, ph = 2 * np.pi * u, np.arccos(2 * v - 1)
    r = 1.0 + 0.18 * np.sin(5 * th) * np.sin(3 * ph)
    pos = np.stack([r * np.sin(ph) * np.cos(th), r * np.sin(ph) * np.sin(th), r * np.cos(ph)], 1)
    scale = 0.03 + 0.02 * (0.5 + 0.5 * np.sin(4 * th))[:, None] * np.ones((n, 3))
    rot = np.stack([np.cos(th / 2), np.sin(th / 2) * 0, np.sin(th / 2) * 0, np.sin(th / 2)], 1)
    rot /= np.linalg.norm(rot, axis=1, keepdims=True)
    col = 0.5 + 0.5 * np.stack([np.sin(th), np.cos(ph), np.sin(th + ph)], 1)
    opac = (0.6 + 0.4 * np.cos(2 * ph))[:, None]
    return np.concatenate([pos, scale, rot, col, opac], 1).astype(np.float32)   # [N,14]


def to_grid(cloud):
    """Sort gaussians by a Morton key over (x,y) → an [ATTR, S, S] image (spatial
    locality), and normalize each attribute channel to [0,1] (carry the ranges)."""
    n = cloud.shape[0]
    p = cloud[:, :2]
    q = ((p - p.min(0)) / (np.ptp(p, 0) + 1e-9) * 255).astype(np.int32)
    order = np.argsort([morton((q[i, 0], q[i, 1])) for i in range(n)])
    g = cloud[order]                                  # [N,14]
    lo, hi = g.min(0), g.max(0)
    gn = (g - lo) / (hi - lo + 1e-9)                  # [N,14] in [0,1]
    img = gn.T.reshape(ATTR, S, S)                    # [14,S,S]
    return img.astype(np.float32), lo, hi, order


class Codec(nn.Module):
    def __init__(self):
        super().__init__()
        self.g_a = nn.Sequential(nn.Conv2d(ATTR, C, 5, 2, 2), GDN2d(C), nn.Conv2d(C, C, 5, 2, 2))
        self.g_s = nn.ModuleList([nn.ConvTranspose2d(C, C, 5, 2, 2, output_padding=1), GDN2d(C),
                                  nn.ConvTranspose2d(C, ATTR, 5, 2, 2, output_padding=1)])

    def synth(self, y):
        x = y
        for m in self.g_s:
            x = m(x)
        return x

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


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


def fit_cdf(sym):
    sym = np.asarray(sym, np.int64).reshape(-1)
    lo, hi = int(sym.min()), int(sym.max())
    n = hi - lo + 1
    hist = np.bincount(sym - 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, lo


def gs_decode_attr(model, xi):
    """Integer g_s on the latent xi ([C,sh,sw]@2^afrac) → [ATTR,S,S] u8 attribute
    image (round-half-up), the byte-exact cloud the Rust port reproduces."""
    x = run_int_net(model, xi)
    half = 1 << (AFRAC - 1)
    x = np.clip(x, 0, ONE)
    return ((x * 255 + half) >> AFRAC).clip(0, 255).astype(np.uint8)             # [ATTR,S,S]


def main(outdir):
    os.makedirs(outdir, exist_ok=True)
    torch.manual_seed(0)
    cloud = make_cloud(S * S)
    img, lo, hi, order = to_grid(cloud)
    x = torch.tensor(img[None])                        # [1,14,S,S]

    m = Codec()
    opt = torch.optim.Adam(m.parameters(), 1.5e-3)
    for step in range(3000):
        xr, y = m(x)
        loss = F.mse_loss(xr, x) + 1e-5 * y.abs().mean()
        opt.zero_grad(); loss.backward(); opt.step()
    m.eval()

    # factorized latent: per-channel median, integer symbols
    with torch.no_grad():
        y = m.g_a(x)[0].numpy()                         # [C,sh,sw]
    Cc, sh, sw = y.shape
    median = np.round(np.median(y.reshape(Cc, -1), 1))
    sym = np.round(y - median[:, None, None]).astype(np.int64)
    median_q = np.round(median * ONE).astype(np.int64)
    xi = (sym * ONE + median_q[:, None, None]).astype(np.int64)

    # PTQ g_s → WHM1
    gs_layers = [ptq_convT(m.g_s[0]), ptq_igdn(m.g_s[1]), ptq_convT(m.g_s[2])]

    attr_int = gs_decode_attr({'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': gs_layers}, xi)
    gt = (img * 255).astype(np.float64)
    snr_i = 10 * np.log10((gt ** 2).sum() / ((attr_int.astype(np.float64) - gt) ** 2).sum() + 1e-12)
    attr_hash = format(fnv1a(np.ascontiguousarray(attr_int).tobytes()), '016x')

    # factorized rANS over the symbols (per-channel data-fit CDF)
    cdfs, offs, flat, idx = [], [], [], []
    for ch in range(Cc):
        cdf, off = fit_cdf(sym[ch])
        cdfs.append(cdf); offs.append(off)
        flat.extend(sym[ch].reshape(-1).tolist()); idx.extend([ch] * sh * sw)
    bits = RansEncoder().encode_with_indexes(flat, idx, cdfs, [len(c) for c in cdfs], offs)

    open(f'{outdir}/splat_gs.bin', 'wb').write(whm1(gs_layers))
    json.dump({'c': Cc, 'sh': sh, 'sw': sw, 's': S, 'attr': ATTR, 'afrac': AFRAC,
               'cdfs': cdfs, 'offsets': offs, 'median_q': median_q.tolist(),
               'rans_b64': base64.b64encode(bytes(bits)).decode(),
               'attr_hash': attr_hash, 'n_gaussians': S * S,
               'lo': lo.tolist(), 'hi': hi.tolist()},   # per-attr ranges → un-normalize for presentation
              open(f'{outdir}/splat.json', 'w'))
    np.ascontiguousarray(attr_int).tofile(f'{outdir}/splat_attr.u8')
    print(f'int splat codec: {S*S} gaussians × {ATTR} attrs, g_s {len(whm1(gs_layers))}B, rANS {len(bits)}B '
          f'-> integer-decode SNR {snr_i:.1f} dB, attr_hash {attr_hash}')


if __name__ == '__main__':
    main(sys.argv[1] if len(sys.argv) > 1 else '/tmp/int_splat')