#!/usr/bin/env python3
"""Train, quantise and export the `wai.video.int_cond` reference vector.

`wai.video.int_cond` is conditional inter-frame coding at the integer tier
(extensions/integer-payloads.md §17). Each frame's temporal context is the previous
frame's decoded *feature*, not its pixels; an integer context network over that
context gives every latent element's scale sigma and mean mu; the latent is decoded
against the table its sigma selects; and an integer conditional decoder turns the
latent and the context into the frame's feature, which an integer synthesis turns
into RGB and which the next frame reads as its context. An intra frame reads the
zero feature. Every network is a WHM1 network lowered by `whm1/1`, so the decode is
byte-identical on every machine.

This tool trains a small model of that topology on seeded procedural clips (moving
shapes over panning gratings; no external content), quantises it after training
(per-output-channel int12 weights, activations at 2^12, the int_hyper family's
quantisation), and encodes a held-out procedural clip with the integer networks
themselves, so the encoder computes exactly the sigma, mu, features and frames every
decoder computes. The vector is a real decode of a trained model; its quality is
what this small model gives and is reported, not claimed. What lowering a pretrained
model's own weights onto this topology takes is integer-payloads §17.7.

  python3 tools/wai_int_cond_export.py <out-dir> [--steps N] [--threads T] [--seed S]

Requires numpy and torch (CPU is enough). Writes into <out-dir>:

  model.wcb      the one-file parameter set: magic WCB1, then five sections, each a u32
                 length and its bytes: ctx, prior, dec, recon (WHM1 networks) and the
                 entropy tables (JSON: afrac, scale_edges, cdfs, cdf_lengths, offsets).
  clip.wic1      the WIC1 clip payload: its frames, each a WCF1 frame.
  sha256.json    each decoded frame's SHA-256 and the clip's (SHA-256 over the frame
                 digests), the sidecar the conformance verifiers decode to.
  expected.json  what the export measured: bytes per frame, PSNR per frame against the
                 source, the source's SHA-256, the training record.

Training is floating point and not reproducible bit for bit across machines or thread
counts; the committed files are the vector, and every decode of them is integer.
"""
import argparse
import hashlib
import json
import math
import os
import struct
import sys

import numpy as np

H = W = 64          # frame
M = 8               # latent channels
CF = 16             # reference feature channels
CC = 16             # temporal context channels
NA = 16             # the encoder's analysis channels (encoder only)
AFRAC = 12          # activations at 2^12
GF = 16             # WHM1 header field (no inverse GDN in this model)
SHIFT = 24          # requantisation shift
WBITS = 12          # weights: per-output-channel int12
LEAKY_Q, LEAKY_F = 655, 16   # leaky slope 655 / 2^16
SCALE_MIN, SCALE_MAX, SCALE_LEVELS = 0.11, 32.0, 24
TAIL = 6.0          # a table spans mu ± ceil(TAIL·scale); the rest is the escape
LAMBDA = 0.013 * 255 * 255   # rate-distortion weight on MSE of [0, 1] pixels
CLIP_FRAMES = 6
TRAIN_T = 4
I64_MAX = (1 << 63) - 1
CAPABILITY = "wai.video.int_cond"


# ---------------------------------------------------------------- procedural clips

def procedural_clip(rng, n, h=H, w=W):
    """`n` RGB8 frames [n, h, w, 3]: two coloured gratings panning at sub-pixel speeds,
    and two to four shapes (discs, boxes, a textured square) moving across them."""
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float64)
    g = []
    for _ in range(2):
        f = rng.uniform(0.04, 0.25, size=2) * rng.choice([-1, 1], size=2)
        g.append((f, rng.uniform(0, 2 * math.pi), rng.uniform(-1.5, 1.5, size=2), rng.uniform(0.2, 0.9, size=3)))
    base = rng.uniform(0.1, 0.6, size=3)
    shapes = []
    for _ in range(int(rng.integers(2, 5))):
        shapes.append(dict(kind=int(rng.integers(0, 3)), c=rng.uniform(0, 1, size=3), r=rng.uniform(4, 14),
                           p=rng.uniform(0, [h, w]), v=rng.uniform(-3, 3, size=2), t=rng.uniform(0.2, 0.8)))
    out = np.zeros((n, h, w, 3))
    for t in range(n):
        img = np.broadcast_to(base, (h, w, 3)).copy()
        for f, ph, v, amp in g:
            s = np.sin((yy - v[0] * t) * f[0] * 2 * math.pi / 8 + (xx - v[1] * t) * f[1] * 2 * math.pi / 8 + ph)
            img += 0.25 * s[..., None] * amp
        for s in shapes:
            cy, cx = s["p"] + s["v"] * t
            if s["kind"] == 0:
                m = np.clip(s["r"] - np.hypot(yy - cy, xx - cx) + 0.5, 0, 1)
            elif s["kind"] == 1:
                m = np.clip(s["r"] - np.maximum(abs(yy - cy), abs(xx - cx)) + 0.5, 0, 1)
            else:
                m = np.clip(s["r"] - np.maximum(abs(yy - cy), abs(xx - cx)) + 0.5, 0, 1)
                tex = 0.5 + 0.5 * np.sign(np.sin((yy - cy) * 1.3) * np.sin((xx - cx) * 1.3))
                img = img * (1 - m[..., None]) + (s["c"] * (0.6 + 0.4 * tex[..., None])) * m[..., None]
                continue
            img = img * (1 - m[..., None]) + s["c"] * m[..., None]
        out[t] = img
    return np.clip(np.round(out * 255), 0, 255).astype(np.uint8)


# ---------------------------------------------------------------- the float model

def build_torch(torch):
    nn = torch.nn
    leaky = lambda: nn.LeakyReLU(LEAKY_Q / (1 << LEAKY_F))

    class CondCodec(nn.Module):
        def __init__(self):
            super().__init__()
            # Encoder only: analysis of the frame, then conditioned on the context.
            self.g_a = nn.Sequential(nn.Conv2d(3, NA, 5, 2, 2), nn.ReLU(), nn.Conv2d(NA, NA, 5, 2, 2))
            self.g_c = nn.Sequential(nn.Conv2d(NA + CC, NA, 3, 1, 1), nn.ReLU(), nn.Conv2d(NA, M, 3, 1, 1))
            # Decode side: the four networks of the set.
            self.ctx = nn.Sequential(nn.Conv2d(CF, CC, 3, 1, 1), nn.ReLU(), nn.Conv2d(CC, CC, 3, 1, 1))
            self.prior = nn.Sequential(nn.Conv2d(CC, CC, 3, 1, 1), leaky(), nn.Conv2d(CC, 2 * M, 1, 1, 0))
            self.dec = nn.Sequential(nn.Conv2d(M + CC, CF, 3, 1, 1), nn.ReLU(), nn.Conv2d(CF, CF, 3, 1, 1))
            self.recon = nn.Sequential(nn.ConvTranspose2d(CF, CF, 5, 2, 2, 1), leaky(),
                                       nn.ConvTranspose2d(CF, 3, 5, 2, 2, 1))
            with torch.no_grad():
                self.prior[2].bias[:M].fill_(1.0)
                self.recon[2].bias.fill_(0.5)

        def step(self, x, feat, train=True):
            c = self.ctx(feat)
            p = self.prior(c)
            sigma, mu = p[:, :M], p[:, M:]
            y = self.g_c(torch.cat([self.g_a(x), c], 1))
            s = sigma.clamp(SCALE_MIN, SCALE_MAX)
            r = y - mu
            noisy = r + (torch.rand_like(r) - 0.5) if train else torch.round(r)
            normal = torch.distributions.Normal(0.0, 1.0)
            lik = normal.cdf((noisy + 0.5) / s) - normal.cdf((noisy - 0.5) / s)
            bits = -torch.log2(lik.clamp_min(1e-9)).sum()
            y_hat = mu + (r + (torch.round(r) - r).detach())      # straight-through rounding
            f = self.dec(torch.cat([y_hat, c], 1))
            return self.recon(f), f, bits

    return CondCodec()


def train(steps, seed, threads):
    import torch
    torch.manual_seed(seed)
    torch.set_num_threads(threads)
    rng = np.random.default_rng(seed)
    net = build_torch(torch)
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    batch = 8
    log = []
    for it in range(steps):
        if it == int(steps * 0.75):
            for g in opt.param_groups:
                g["lr"] = 1e-4
        clips = np.stack([procedural_clip(rng, TRAIN_T) for _ in range(batch)])           # [B, T, H, W, 3]
        x = torch.from_numpy(clips.astype(np.float32) / 255).permute(0, 1, 4, 2, 3)      # [B, T, 3, H, W]
        feat = torch.zeros(batch, CF, H // 4, W // 4)
        loss, bits_all, mse_all = 0.0, 0.0, 0.0
        for t in range(TRAIN_T):
            rec, feat, bits = net.step(x[:, t], feat)
            mse = torch.mean((rec - x[:, t]) ** 2)
            bpp = bits / (batch * H * W)
            loss = loss + LAMBDA * mse + bpp
            bits_all += float(bpp)
            mse_all += float(mse)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0)
        opt.step()
        if it % 250 == 0 or it == steps - 1:
            psnr = 10 * math.log10(1 / max(mse_all / TRAIN_T, 1e-12))
            log.append({"step": it, "bpp": round(bits_all / TRAIN_T, 4), "psnr_db": round(psnr, 2)})
            print(f"step {it}: bpp {bits_all / TRAIN_T:.4f}, psnr {psnr:.2f} dB", flush=True)
    return net, log, torch.__version__


# ---------------------------------------------------------------- quantisation

def q_conv(layer):
    w = layer.weight.detach().double().numpy()                         # [out, in, k, k]
    s = np.abs(w).max(axis=(1, 2, 3)) / (2 ** (WBITS - 1) - 1) + 1e-12
    wi = np.round(w / s[:, None, None, None]).clip(-(2 ** (WBITS - 1)), 2 ** (WBITS - 1) - 1).astype(np.int64)
    return {"type": "conv", "in_ch": int(w.shape[1]), "out_ch": int(w.shape[0]), "k": int(w.shape[2]),
            "stride": int(layer.stride[0]), "pad": int(layer.padding[0]), "weight": wi,
            "mult": np.round(s * 2 ** SHIFT).astype(np.int64),
            "bias": np.round(layer.bias.detach().double().numpy() * 2 ** AFRAC).astype(np.int64)}


def q_convt(layer):
    w = layer.weight.detach().double().numpy()                         # [in, out, k, k]
    s = np.abs(w).max(axis=(0, 2, 3)) / (2 ** (WBITS - 1) - 1) + 1e-12
    wi = np.round(w / s[None, :, None, None]).clip(-(2 ** (WBITS - 1)), 2 ** (WBITS - 1) - 1).astype(np.int64)
    return {"type": "convT", "in_ch": int(w.shape[0]), "out_ch": int(w.shape[1]), "k": int(w.shape[2]),
            "stride": int(layer.stride[0]), "pad": int(layer.padding[0]), "opad": int(layer.output_padding[0]),
            "weight": wi, "mult": np.round(s * 2 ** SHIFT).astype(np.int64),
            "bias": np.round(layer.bias.detach().double().numpy() * 2 ** AFRAC).astype(np.int64)}


def quantise(seq):
    import torch.nn as nn
    layers = []
    for m in seq:
        if isinstance(m, nn.Conv2d):
            layers.append(q_conv(m))
        elif isinstance(m, nn.ConvTranspose2d):
            layers.append(q_convt(m))
        elif isinstance(m, nn.ReLU):
            layers.append({"type": "relu"})
        elif isinstance(m, nn.LeakyReLU):
            layers.append({"type": "leaky"})
        else:
            raise ValueError(type(m))
    return {"afrac": AFRAC, "gf": GF, "shift": SHIFT, "layers": layers}


def whm1(model):
    """The WHM1 bytes of an integer network (extensions/integer-payloads.md §9)."""
    out = bytearray(b"WHM1") + struct.pack("<IIII", model["afrac"], model["gf"], model["shift"], len(model["layers"]))
    for L in model["layers"]:
        t = L["type"]
        if t in ("conv", "convT"):
            assert np.abs(L["weight"]).max() < 2 ** 15
            if t == "convT":
                out += struct.pack("<BIIIIII", 0, L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"], L["opad"])
            else:
                out += struct.pack("<BIIIII", 1, L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"])
            out += L["weight"].astype("<i2").tobytes() + L["mult"].astype("<i8").tobytes() + L["bias"].astype("<i8").tobytes()
        elif t == "relu":
            out += struct.pack("<B", 3)
        elif t == "leaky":
            out += struct.pack("<BqI", 4, LEAKY_Q, LEAKY_F)
    return bytes(out)


# ---------------------------------------------------------------- the integer networks

def mag(x):
    return int(np.max(np.abs(x.astype(object)))) if x.size else 0


def run_int(model, x):
    """Run an integer network on x ([c, h, w] int64 at 2^AFRAC) with the arithmetic of
    integer-operators §13.2: a convolution's sum, then floor((sum · mult) / 2^shift) +
    bias; a leaky ReLU's (v · 655) >> 16 below zero. Before every layer the bound that
    keeps the sum and the requantised value inside int64 is checked, since numpy wraps."""
    for L in model["layers"]:
        t = L["type"]
        if t in ("conv", "convT"):
            k, cin, co = L["k"], L["in_ch"], L["out_ch"]
            assert x.shape[0] == cin
            acc = cin * k * k * int(np.abs(L["weight"]).max()) * mag(x)
            assert acc <= I64_MAX and ((acc * int(L["mult"].max())) >> SHIFT) + 1 + int(np.abs(L["bias"]).max()) <= I64_MAX
            s, p = L["stride"], L["pad"]
            h, w = x.shape[1:]
            if t == "conv":
                ho, wo = (h + 2 * p - k) // s + 1, (w + 2 * p - k) // s + 1
                xp = np.zeros((cin, h + 2 * p, w + 2 * p), dtype=np.int64)
                xp[:, p:p + h, p:p + w] = x
                raw = np.zeros((co, ho, wo), dtype=np.int64)
                for ky in range(k):
                    for kx in range(k):
                        sub = xp[:, ky:ky + s * ho:s, kx:kx + s * wo:s]
                        raw += np.tensordot(L["weight"][:, :, ky, kx], sub, axes=([1], [0]))
            else:
                op = L["opad"]
                ho, wo = (h - 1) * s + k + op - 2 * p, (w - 1) * s + k + op - 2 * p
                full = np.zeros((co, ho + 2 * p, wo + 2 * p), dtype=np.int64)
                for ky in range(k):
                    for kx in range(k):
                        contrib = np.tensordot(L["weight"][:, :, ky, kx], x, axes=([0], [0]))   # [co, h, w]
                        full[:, ky:ky + s * h:s, kx:kx + s * w:s] += contrib
                raw = full[:, p:p + ho, p:p + wo]
            x = ((raw.astype(object) * L["mult"][:, None, None]) >> SHIFT).astype(np.int64) + L["bias"][:, None, None]
        elif t == "relu":
            x = np.maximum(x, 0)
        elif t == "leaky":
            assert mag(x) * LEAKY_Q <= I64_MAX
            x = np.where(x >= 0, x, (x * LEAKY_Q) >> LEAKY_F)
    return x


def to_rgb8(x):
    one, half = 1 << AFRAC, 1 << (AFRAC - 1)
    return ((np.clip(x, 0, one) * 255 + half) >> AFRAC).astype(np.uint8).transpose(1, 2, 0).copy()


def bucket(sigma, edges):
    """integer-operators E3 `ceil`: b = N − 1 − #{i < N − 1 : max(σ, e_0) ≤ e_i}."""
    s = np.maximum(sigma, edges[0])
    b = np.full(s.shape, len(edges) - 1, dtype=np.int64)
    for e in edges[:-1]:
        b -= (s <= e).astype(np.int64)
    return b


# ---------------------------------------------------------------- tables and the coder

def gaussian_tables():
    scales = np.exp(np.linspace(math.log(SCALE_MIN), math.log(SCALE_MAX), SCALE_LEVELS))
    edges = [int(math.floor(s * 2 ** AFRAC)) for s in scales]
    cdfs, offsets = [], []
    phi = lambda v: 0.5 * (1 + math.erf(v / math.sqrt(2)))
    for s in scales:
        half = max(1, int(math.ceil(TAIL * s)))
        p = [phi((v + 0.5) / s) - phi((v - 0.5) / s) for v in range(-half, half + 1)]
        esc = max(1.0 - sum(p), 0.0)
        f = [max(1, int(round(q * 65536))) for q in p] + [max(1, int(round(esc * 65536)))]
        f[half] += 65536 - sum(f)
        assert min(f) >= 1
        cdf = [0]
        for q in f:
            cdf.append(cdf[-1] + q)
        assert cdf[-1] == 65536 and 3 <= len(cdf) <= 4097
        cdfs.append(cdf)
        offsets.append(-half)
    return edges, cdfs, offsets


def coded(value, cdf, offset):
    """integer-payloads §2: a value as its table symbol and, outside the table, the
    escape's 4-bit groups."""
    n = len(cdf)
    v = value - offset
    if 0 <= v < n - 2:
        return [("t", cdf[v], cdf[v + 1] - cdf[v])]
    raw = -2 * v - 1 if v < 0 else 2 * (v - (n - 2))
    assert raw < 1 << 32
    nib = 0
    while raw >> (4 * nib):
        nib += 1
    return [("t", cdf[n - 2], cdf[n - 1] - cdf[n - 2]), ("b", nib)] + [("b", (raw >> (4 * j)) & 15) for j in range(nib)]


def rans_encode(values, tables):
    """The rANS encoder of integer-payloads §2 (precision 16, 4-bit bypass), its state
    starting at 2^31, so the stream ends canonically."""
    syms = [s for v, (cdf, off) in zip(values, tables) for s in coded(v, cdf, off)]
    state, words = 1 << 31, []
    lo = (1 << 31) >> 16
    for s in reversed(syms):
        if s[0] == "t":
            start, freq = s[1], s[2]
            while state >= (lo << 32) * freq:
                words.append(state & 0xFFFFFFFF)
                state >>= 32
            state = ((state // freq) << 16) + (state % freq) + start
        else:
            if state >= (lo << 32) << 12:
                words.append(state & 0xFFFFFFFF)
                state >>= 32
            state = (state << 4) | s[1]
    return struct.pack("<II", state & 0xFFFFFFFF, state >> 32) + b"".join(struct.pack("<I", x) for x in reversed(words))


# ---------------------------------------------------------------- encode

def encode(nets, edges, cdfs, offsets, frames):
    one, half = 1 << AFRAC, 1 << (AFRAC - 1)
    sh, sw = H // 4, W // 4
    feat = np.zeros((CF, sh, sw), dtype=np.int64)
    payloads, decoded, stats = [], [], []
    for t, img in enumerate(frames):
        c = run_int(nets["ctx"], feat)
        p = run_int(nets["prior"], c)
        assert p.shape == (2 * M, sh, sw), p.shape
        sigma, mu = p[:M], p[M:]
        x = (img.transpose(2, 0, 1).astype(np.int64) << AFRAC) // 255
        y = run_int(nets["g_c"], np.concatenate([run_int(nets["g_a"], x), c]))
        r = (y - mu + half) >> AFRAC                                        # round half up
        y_hat = r * one + mu
        feat = run_int(nets["dec"], np.concatenate([y_hat, c]))
        assert feat.shape == (CF, sh, sw)
        rgb = to_rgb8(run_int(nets["recon"], feat))
        assert rgb.shape == (H, W, 3)
        b = bucket(sigma, edges).reshape(-1)
        stream = rans_encode([int(v) for v in r.reshape(-1)], [(cdfs[i], offsets[i]) for i in b])
        kind = 0 if t == 0 else 1
        payloads.append(struct.pack("<4sBBHIIIIII", b"WCF1", 1, kind, 0, H, W, M, sh, sw, CF)
                        + struct.pack("<I", len(stream)) + stream)
        decoded.append(rgb)
        mse = float(np.mean((rgb.astype(np.float64) - img.astype(np.float64)) ** 2))
        esc = sum(1 for v, i in zip(r.reshape(-1), b) if not (0 <= v - offsets[i] < len(cdfs[i]) - 2))
        stats.append({"frame": t, "kind": "intra" if kind == 0 else "inter", "bytes": len(payloads[-1]),
                      "psnr_db": round(10 * math.log10(255 * 255 / max(mse, 1e-9)), 2), "escapes": esc,
                      "buckets_used": len(set(b.tolist()))})
    return payloads, decoded, stats


def bundle(nets, edges, cdfs, offsets):
    tables = {"afrac": AFRAC, "scale_edges": edges, "cdfs": cdfs, "cdf_lengths": [len(c) for c in cdfs], "offsets": offsets}
    out = bytearray(b"WCB1")
    for part in (whm1(nets["ctx"]), whm1(nets["prior"]), whm1(nets["dec"]), whm1(nets["recon"]),
                 json.dumps(tables, separators=(",", ":")).encode()):
        out += struct.pack("<I", len(part)) + part
    return bytes(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("outdir")
    ap.add_argument("--steps", type=int, default=4000)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--seed", type=int, default=20261109)
    a = ap.parse_args()
    net, log, torch_version = train(a.steps, a.seed, a.threads)
    nets = {name: quantise(getattr(net, name)) for name in ("g_a", "g_c", "ctx", "prior", "dec", "recon")}
    edges, cdfs, offsets = gaussian_tables()
    source = procedural_clip(np.random.default_rng(a.seed + 1), CLIP_FRAMES)
    payloads, decoded, stats = encode(nets, edges, cdfs, offsets, source)
    clip = struct.pack("<4sB3xI", b"WIC1", 1, len(payloads)) + b"".join(struct.pack("<I", len(p)) + p for p in payloads)
    wcb = bundle(nets, edges, cdfs, offsets)
    digests = [hashlib.sha256(f.tobytes()).digest() for f in decoded]
    side = {"clip": {"buffer": "RGB8: rows top to bottom, pixels left to right, R, G, B bytes",
                     "frames_sha256": [d.hex() for d in digests], "clip_sha256": hashlib.sha256(b"".join(digests)).hexdigest()}}
    expected = {
        "capability": CAPABILITY, "h": H, "w": W, "n_frames": len(payloads),
        "latent": [M, H // 4, W // 4], "feature_channels": CF, "context_channels": CC,
        "model_sha256": hashlib.sha256(wcb).hexdigest(), "model_bytes": len(wcb),
        "payload_sha256": hashlib.sha256(clip).hexdigest(), "payload_bytes": len(clip),
        "source_sha256": hashlib.sha256(source.tobytes()).hexdigest(),
        "frames": stats,
        "training": {"seed": a.seed, "steps": a.steps, "batch": 8, "unrolled_frames": TRAIN_T, "lambda_mse": LAMBDA,
                     "data": "seeded procedural clips (panning gratings, moving shapes), generated by this tool",
                     "torch": torch_version, "log": log},
        "weights": "trained by this tool, then quantised after training: per-output-channel int12 weights, activations at 2^12",
    }
    os.makedirs(a.outdir, exist_ok=True)
    for name, data in (("model.wcb", wcb), ("clip.wic1", clip)):
        with open(os.path.join(a.outdir, name), "wb") as f:
            f.write(data)
    for name, obj in (("sha256.json", side), ("expected.json", expected)):
        with open(os.path.join(a.outdir, name), "w") as f:
            json.dump(obj, f, indent=1)
            f.write("\n")
    print(f"int_cond vector: {len(payloads)} frames of {W}x{H}, {len(clip)} payload bytes, model {len(wcb)} bytes")
    for s in stats:
        print(f"  frame {s['frame']} ({s['kind']}): {s['bytes']} bytes, {s['psnr_db']} dB, escapes {s['escapes']}")


if __name__ == "__main__":
    sys.exit(main())