"""Quantize a synthesis transform to the integer `int_synth.json` spec decoded
bit-exactly by `wai-rs/src/int_synth.rs` — the integer-exact `replicate`
path.

`int_synth` runs an all-integer transposed-conv stack (no float, no ONNX), so
its decode is identical on every machine. This tool turns a trained float
synthesis (ConvTranspose2d + ReLU layers) into the integer weights + per-layer
requantize shifts the Rust decoder needs.

IMPORTANT on quality: integer fixed-point inference is bit-exact regardless, but
this tool rounds every layer with ONE fixed shift, so a float model exported here
will decode deterministically and may look much worse than it does in float — the
determinism is free, the quality is the quantisation step. The working path is
`tools/wai_int_synth_train.py` (and `wai_int_synth_real.py`, which wrote the trained
corpus vector): float training, then post-training quantisation with each layer's
shift derived from its measured activation range. `wai_int_synth_qat.py` is an
experimental attempt to train through a fixed shift; no committed model came from it.

Spec written (matches int_synth.rs):
  { "layers": [ { in_ch, out_ch, k, stride, pad, out_pad,
                  weight:[i32; in*out*k*k]  # layout ic,oc,ky,kx
                  bias:[i32; out_ch], shift:u32, relu:bool } ] }

Run (needs torch):
  python tools/wai_int_synth_export.py --weights synth.pt \
    --out wai-web/demo/models/int_synth/int_synth.json
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn as nn


def quantize_layer(conv: nn.ConvTranspose2d, w_bits: int, shift: int, relu: bool) -> dict:
    """Per-tensor symmetric integer quantization of a ConvTranspose2d.

    weight_q = round(weight / w_scale), bias_q = round(bias / (w_scale * in_scale))
    folded so the integer accumulate matches; the decoder applies `>> shift` to
    requantize. We pick `w_scale` from the weight's max-abs to fill the int range.
    """
    w = conv.weight.detach().cpu()                    # [in_ch, out_ch, kH, kW]
    qmax = (1 << (w_bits - 1)) - 1
    w_scale = float(w.abs().max()) / qmax if float(w.abs().max()) > 0 else 1.0
    w_q = torch.clamp(torch.round(w / w_scale), -qmax - 1, qmax).to(torch.int64)

    bias = conv.bias.detach().cpu() if conv.bias is not None else torch.zeros(conv.out_channels)
    # Bias lives in the accumulator domain (before the >> shift).
    bias_q = torch.round(bias / w_scale * (1 << shift)).to(torch.int64)

    in_ch, out_ch, kh, kw = w.shape
    assert kh == kw, "int_synth assumes square kernels"
    sH, sW = conv.stride
    pH, pW = conv.padding
    opH, opW = conv.output_padding
    assert sH == sW and pH == pW and opH == opW, "int_synth assumes symmetric stride/pad/out_pad"

    return {
        "in_ch": in_ch, "out_ch": out_ch, "k": kh,
        "stride": sH, "pad": pH, "out_pad": opH,
        "weight": w_q.reshape(-1).tolist(),
        "bias": bias_q.reshape(-1).tolist(),
        "shift": shift,
        "relu": relu,
        "_w_scale": w_scale,  # informational
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", required=True,
                    help="torch checkpoint: a Sequential of ConvTranspose2d (+ReLU between).")
    ap.add_argument("--out", required=True)
    ap.add_argument("--w-bits", type=int, default=8)
    ap.add_argument("--shift", type=int, default=8,
                    help="per-layer requantize right-shift (one value for every layer).")
    args = ap.parse_args()

    model = torch.load(args.weights, map_location="cpu", weights_only=False)
    model.eval()

    convs = [m for m in model.modules() if isinstance(m, nn.ConvTranspose2d)]
    if not convs:
        raise SystemExit("no ConvTranspose2d layers found in the checkpoint")
    # ReLU after every layer except the last (final layer emits the image).
    layers = [quantize_layer(c, args.w_bits, args.shift, relu=(i < len(convs) - 1))
              for i, c in enumerate(convs)]
    if layers[-1]["out_ch"] != 3:
        print(f"WARNING: final layer out_ch={layers[-1]['out_ch']} != 3 (RGB)")

    out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({"layers": layers}, separators=(",", ":")))
    print(f"wrote {len(layers)}-layer integer synthesis → {out}  "
          f"(w_bits={args.w_bits}, shift={args.shift})")
    print("determinism: bit-exact on any machine by construction; "
          "quality depends on the shift; see wai_int_synth_train.py for calibrated shifts.")
    return 0


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