"""An EXPERIMENTAL quantization-aware training (QAT) attempt for the
`wai.neural.int_synth` capability (decoder: `wai-rs/src/int_synth.rs`). No
committed model was made with it.

`int_synth` runs an all-integer transposed-conv stack with a per-layer
requantize `>> shift`. The decode is bit-exact on every machine regardless of
weights. This tool trains the weights with straight-through fake-quant (forward =
quantized integers, backward = float gradients) against ONE fixed shift for every
layer; it does not learn the shifts. With a fixed shift, small activations can
collapse to 0 (see the comment in `IntSynth.forward`). The working path is
`tools/wai_int_synth_train.py`: train in
float, then quantise after training with each layer's shift derived from its
measured activation range. The trained corpus vector was made that way
(`tools/wai_int_synth_real.py`).

Inputs: paired (latent, target_image) tensors — e.g. dump `(y_hat, x)` from a
trained float codec on a dataset, or co-train with an analysis transform. This
tool is the decoder-side QAT + export; the latents come from upstream.

Output: `int_synth.json` (== the Rust spec).

Run (needs torch):
  python tools/wai_int_synth_qat.py --data pairs.pt --epochs 20 \
    --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
import torch.nn.functional as F


def fake_quant_weight(w: torch.Tensor, bits: int) -> torch.Tensor:
    """Straight-through symmetric per-tensor weight fake-quant → integer grid."""
    qmax = (1 << (bits - 1)) - 1
    scale = w.detach().abs().max().clamp_min(1e-8) / qmax
    wq = torch.clamp(torch.round(w / scale), -qmax - 1, qmax) * scale
    return w + (wq - w).detach(), scale  # STE: forward quantized, backward identity


class IntSynth(nn.Module):
    """A transposed-conv synthesis whose forward emulates the integer decode:
    fake-quant weights, integer-style requantize via the per-layer `shift`."""
    def __init__(self, chans, kernel=5, stride=2, w_bits=8, shift=8):
        super().__init__()
        self.w_bits, self.shift = w_bits, shift
        self.convs = nn.ModuleList()
        for i in range(len(chans) - 1):
            pad = (kernel - stride + 1) // 2
            self.convs.append(nn.ConvTranspose2d(chans[i], chans[i + 1], kernel,
                                                 stride=stride, padding=pad, output_padding=stride - 1))

    def forward(self, y):
        x = y
        for i, conv in enumerate(self.convs):
            wq, _ = fake_quant_weight(conv.weight, self.w_bits)
            x = F.conv_transpose2d(x, wq, conv.bias, conv.stride, conv.padding, conv.output_padding)
            # Emulate the integer >> shift requantize. NB: plain torch.floor has
            # ZERO gradient — it silently freezes training. Use a straight-through
            # estimator so gradients flow as for the float division. Even so, a
            # FIXED shift can collapse small activations to 0; for a model that
            # actually reconstructs use tools/wai_int_synth_train.py (float train
            # + per-layer activation-calibrated shifts) instead of pure QAT.
            d = x / (1 << self.shift)
            x = d + (torch.floor(d) - d).detach()
            if i < len(self.convs) - 1:
                x = torch.relu(x)
        return x

    def export_spec(self) -> dict:
        layers = []
        for i, conv in enumerate(self.convs):
            qmax = (1 << (self.w_bits - 1)) - 1
            w = conv.weight.detach()
            scale = float(w.abs().max()) / qmax if float(w.abs().max()) > 0 else 1.0
            w_q = torch.clamp(torch.round(w / scale), -qmax - 1, qmax).to(torch.int64)
            b = conv.bias.detach() if conv.bias is not None else torch.zeros(conv.out_channels)
            b_q = torch.round(b / scale * (1 << self.shift)).to(torch.int64)
            pad = conv.padding[0]
            layers.append({
                "in_ch": conv.in_channels, "out_ch": conv.out_channels,
                "k": conv.kernel_size[0], "stride": conv.stride[0],
                "pad": pad, "out_pad": conv.output_padding[0],
                "weight": w_q.reshape(-1).tolist(), "bias": b_q.reshape(-1).tolist(),
                "shift": self.shift, "relu": i < len(self.convs) - 1,
            })
        return {"layers": layers}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="torch .pt with tensors 'y' (latents) and 'x' (target images, [0,1]).")
    ap.add_argument("--chans", type=int, nargs="+", default=[192, 128, 64, 32, 3],
                    help="channel widths: latent → … → 3 (RGB).")
    ap.add_argument("--epochs", type=int, default=20)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--shift", type=int, default=8)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    blob = torch.load(args.data, map_location="cpu", weights_only=False)
    y, x = blob["y"].float(), blob["x"].float()
    net = IntSynth(args.chans, shift=args.shift).train()
    opt = torch.optim.Adam(net.parameters(), lr=args.lr)

    for ep in range(args.epochs):
        opt.zero_grad()
        recon = net(y) / 255.0           # integer domain → [0,1]-ish
        loss = F.mse_loss(recon.clamp(0, 1), x)
        loss.backward()
        opt.step()
        if ep % 5 == 0 or ep == args.epochs - 1:
            print(f"epoch {ep:3d}  mse {loss.item():.5f}")

    out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(net.eval().export_spec(), separators=(",", ":")))
    print(f"wrote an experimental QAT integer synthesis → {out}  (shift={args.shift}); "
          f"decode is bit-exact on any machine by construction, quality is not checked.")
    return 0


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