"""Train an integer synthesis for `wai.neural.int_synth` and export the
`int_synth.json` the Rust decoder runs — the WORKING path: float training +
post-training quantization with PER-LAYER ACTIVATION-SCALE CALIBRATION.

Why not pure QAT-through-the-shift: the integer requantize `>> shift` has zero
gradient (floor), so straight QAT can't learn; even with a straight-through
floor, a fixed shift collapses small activations to 0. The robust recipe is to
train in float, then quantize each layer with a shift derived from the MEASURED
activation range so the integer pipeline tracks the float one:

  s_w   = max|W| / 127                 (per-layer weight scale)
  s_i   = max|activation_i| / 127      (per-layer activation scale, calibrated)
  shift = round(log2( s_i / (s_w · s_{i-1}) ))         (≥ 0)
  W_int = round(W / s_w)               b_int = round(b / (s_w · s_{i-1}))
  latent_int = round(latent / s_0)

The integer forward (`int_synth.rs`) is then bit-exact across machines BY
CONSTRUCTION, and the Rust decode matches this tool's integer reference to
max|Δ|=0 (`int_synth_trained_model_matches_integer_forward`). Quality scales
with model size / data / per-channel scales; this reference run hits ~16 dB on a
tiny synthetic set in seconds.

Run (needs torch; bring images as a [N,3,H,W] float tensor in [0,1], H=W=2^k·S):
  python tools/wai_int_synth_train.py --data imgs.pt --out wai-web/demo/models/int_synth/int_synth.json
"""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", help="torch .pt: float images [N,3,H,W] in [0,1]. Omit for a synthetic demo set.")
    ap.add_argument("--chans", type=int, nargs="+", default=[24, 24, 16, 8, 3])
    ap.add_argument("--latent-side", type=int, default=4)
    ap.add_argument("--epochs", type=int, default=600)
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    torch.manual_seed(0)
    C, LS = args.chans[0], args.latent_side

    if args.data:
        data = torch.load(args.data, map_location="cpu").float()
    else:
        SZ = LS * 16
        imgs = []
        for k in range(24):
            xx, yy = np.meshgrid(np.linspace(0, 1, SZ), np.linspace(0, 1, SZ))
            imgs.append(np.stack([np.sin((xx + .1 * k) * 8) * .5 + .5,
                                  np.cos((yy + .07 * k) * 6) * .5 + .5, (xx * yy + .05 * k) % 1.0], 0))
        data = torch.from_numpy(np.stack(imgs).astype(np.float32))

    dn = lambda a, b: nn.Conv2d(a, b, 5, 2, 2)
    up = lambda a, b: nn.ConvTranspose2d(a, b, 5, 2, 2, output_padding=1)
    n_down = int(round(math.log2(data.shape[-1] // LS)))
    enc_layers = []
    ch = 3
    for _ in range(n_down):
        enc_layers += [dn(ch, C), nn.ReLU()]; ch = C
    enc = nn.Sequential(*enc_layers[:-1])
    dec = nn.ModuleList([up(args.chans[i], args.chans[i + 1]) for i in range(len(args.chans) - 1)])

    def dfwd(y):
        x = y
        for i, c in enumerate(dec):
            x = c(x)
            if i < len(dec) - 1:
                x = torch.relu(x)
        return x

    opt = torch.optim.Adam(list(enc.parameters()) + [p for c in dec for p in c.parameters()], lr=args.lr)
    for ep in range(args.epochs):
        opt.zero_grad()
        recon = dfwd(enc(data) * 64.0)
        loss = F.mse_loss(recon.clamp(0, 1) * 255, data * 255)
        loss.backward(); opt.step()
        if ep % 100 == 0 or ep == args.epochs - 1:
            print(f"  epoch {ep:3d}  px-mse {loss.item():.1f}")

    # calibrate activation scales, then quantize.
    with torch.no_grad():
        x = enc(data) * 64.0
        amax = []
        for i, c in enumerate(dec):
            x = c(x)
            if i < len(dec) - 1:
                x = torch.relu(x)
            amax.append(float(x.abs().max()))
        s0 = float((enc(data) * 64.0).abs().max()) / 127.0
    scales = [s0] + [a / 127.0 for a in amax[:-1]] + [1.0 / 255.0]  # final: int = pixel
    layers, sin = [], s0
    for i, c in enumerate(dec):
        sw = float(c.weight.detach().abs().max()) / 127.0
        shift = max(int(round(math.log2(scales[i + 1] / (sw * sin)))), 0)
        layers.append({
            "in_ch": c.in_channels, "out_ch": c.out_channels, "k": 5, "stride": 2, "pad": 2, "out_pad": 1,
            "weight": torch.clamp(torch.round(c.weight.detach() / sw), -128, 127).to(torch.int64).reshape(-1).tolist(),
            "bias": torch.round(c.bias.detach() / (sw * sin)).to(torch.int64).reshape(-1).tolist(),
            "shift": shift, "relu": i < len(dec) - 1,
        })
        sin = scales[i + 1]
    out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({"layers": layers}, separators=(",", ":")))
    print(f"wrote calibrated integer synthesis → {out}  shifts={[l['shift'] for l in layers]}; "
          f"decode is bit-exact on any machine. Latent scale s0={s0:.4g} (encoder is not exported; "
          f"quantize your latents as round(latent/s0) ∈ int8).")
    return 0


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