"""Train a REAL int_synth decoder on REAL images and emit a committable
conformance fixture: the integer model (`int_synth.json`), an encoded test
image (`latent.bin`, the wire payload), and the integer-forward reference RGB
+ PSNR. The Rust decoder (`int_synth.rs`) reproduces the reference BIT-EXACTLY
on every machine — that is the whole point of int_synth: the determinism is
exact by construction, and this proves it on a real photograph, not a
constructed toy or a synthetic gradient.

Pipeline (the WORKING recipe — see wai_int_synth_train.py for why pure QAT
through the `>> shift` fails): float-train an autoencoder on random crops of
real photos, then post-training-quantize the DECODER with per-layer
activation-scale calibration so the integer pipeline tracks the float one.
Encode a held-out crop through the float encoder, quantize the latent to int8,
and decode it through the EXPORTED INTEGER model in numpy — the same i64
arithmetic int_synth.rs runs (scatter-accumulate, +bias, arithmetic `>> shift`,
integer ReLU, clamp [0,255]). That numpy decode is the reference the Rust
decoder must match to max|Δ|=0.

  python tools/wai_int_synth_real.py --src <dir-of-pngs> --out <fixture-dir>
"""
from __future__ import annotations

import argparse
import glob
import json
import math
import struct
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image


def load_crops(src: str, size: int, n: int, seed: int = 0, downscale: bool = False) -> torch.Tensor:
    rng = np.random.default_rng(seed)
    paths = sorted(glob.glob(str(Path(src) / "*.png")) + glob.glob(str(Path(src) / "*.jpg")))
    if not paths:
        raise SystemExit(f"no images in {src}")
    if downscale:
        # each image becomes a full recognizable scene at `size`×`size`.
        out = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.LANCZOS), np.float32) / 255.0
               for p in paths]
        return torch.from_numpy(np.stack([o.transpose(2, 0, 1) for o in out]))
    crops = []
    per = max(1, n // len(paths))
    for p in paths:
        im = np.asarray(Image.open(p).convert("RGB"), np.float32) / 255.0
        H, W, _ = im.shape
        for _ in range(per):
            y = int(rng.integers(0, max(1, H - size)))
            x = int(rng.integers(0, max(1, W - size)))
            crops.append(im[y:y + size, x:x + size].transpose(2, 0, 1))
    return torch.from_numpy(np.stack(crops[:n]))


# ---- the integer forward, byte-for-byte as int_synth.rs runs it ----
def conv_transpose2d_int(inp: np.ndarray, layer: dict) -> np.ndarray:
    cin, hin, win = inp.shape
    oc, k, s, p, op = layer["out_ch"], layer["k"], layer["stride"], layer["pad"], layer["out_pad"]
    hout = (hin - 1) * s + k + op - 2 * p
    wout = (win - 1) * s + k + op - 2 * p
    W = np.array(layer["weight"], np.int64).reshape(cin, oc, k, k)
    out = np.zeros((oc, hout, wout), np.int64)
    for ky in range(k):
        for kx in range(k):
            # contribution of this kernel tap: (oc, hin, win)
            contrib = np.einsum("iohw,ihw->ohw", W[:, :, ky, kx][:, :, None, None] + np.zeros((1, 1, hin, win), np.int64), inp)
            ys = np.arange(hin) * s + ky - p
            xs = np.arange(win) * s + kx - p
            ym = (ys >= 0) & (ys < hout)
            xm = (xs >= 0) & (xs < wout)
            if not ym.any() or not xm.any():
                continue
            out[np.ix_(np.arange(oc), ys[ym], xs[xm])] += contrib[:, ym][:, :, xm]
    b = np.array(layer["bias"], np.int64).reshape(oc, 1, 1)
    out = (out + b) >> np.int64(layer["shift"])   # arithmetic shift, matches Rust i64 >>
    if layer["relu"]:
        out = np.maximum(out, 0)
    return out


def decode_int(latent_i8: np.ndarray, layers: list[dict], h: int, w: int) -> np.ndarray:
    cur = latent_i8.astype(np.int64)
    for l in layers:
        cur = conv_transpose2d_int(cur, l)
    assert cur.shape[0] == 3 and cur.shape[1] == h and cur.shape[2] == w, cur.shape
    return np.clip(cur, 0, 255).astype(np.uint8).transpose(1, 2, 0)  # H,W,3


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True, help="dir of real .png/.jpg photos")
    ap.add_argument("--out", required=True, help="fixture output dir")
    ap.add_argument("--chans", type=int, nargs="+", default=[64, 48, 32, 16, 3])
    ap.add_argument("--latent-side", type=int, default=8)
    ap.add_argument("--crops", type=int, default=160)
    ap.add_argument("--epochs", type=int, default=2500)
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--downscale", action="store_true", help="train on full images resized to the model size (recognizable whole scenes) instead of random crops")
    ap.add_argument("--test-idx", type=int, default=0, help="which image (in --downscale mode) is the reconstructed fixture")
    args = ap.parse_args()
    torch.manual_seed(0)
    dev = "mps" if torch.backends.mps.is_available() else "cpu"

    C, LS = args.chans[0], args.latent_side
    n_layers = len(args.chans) - 1
    size = LS * (2 ** n_layers)
    print(f"device {dev} · image {size}×{size} · latent {LS}×{LS}×{C} · {n_layers} decoder layers")

    if args.downscale:
        data = load_crops(args.src, size, args.crops, downscale=True).to(dev)
        test = data[args.test_idx:args.test_idx + 1]   # one full scene the model represents
        print(f"downscale mode: {data.shape[0]} full scenes; test image index {args.test_idx}")
    else:
        data = load_crops(args.src, size, args.crops).to(dev)
        test = load_crops(args.src, size, 1, seed=999).to(dev)  # held-out crop

    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)
    enc = nn.Sequential(*sum([[dn(3 if i == 0 else C, C), nn.ReLU()] for i in range(n_layers)], [])[:-1]).to(dev)
    dec = nn.ModuleList([up(args.chans[i], args.chans[i + 1]) for i in range(n_layers)]).to(dev)

    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 % 250 == 0 or ep == args.epochs - 1:
            print(f"  epoch {ep:4d}  train px-mse {loss.item():7.1f}")

    # ---- calibrate activation scales on the training set, 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]
    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]

    # ---- encode the held-out test image → int8 latent → wire payload ----
    with torch.no_grad():
        lat = (enc(test) * 64.0)[0].cpu().numpy()            # (C, LS, LS)
    lat_i8 = np.clip(np.round(lat / s0), -128, 127).astype(np.int8)

    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    (out / "int_synth.json").write_text(json.dumps({"layers": layers}, separators=(",", ":")))

    # wire payload: u32 H, u32 W, u16 C, u8 S_log2, u8 reserved, then C*S*S int8
    s_log2 = int(round(math.log2(LS)))
    payload = struct.pack("<IIHBB", size, size, C, s_log2, 0) + lat_i8.tobytes()
    (out / "latent.bin").write_bytes(payload)

    # reference = the EXPORTED integer model decoded in numpy (== int_synth.rs)
    recon = decode_int(lat_i8, layers, size, size)
    orig = (test[0].cpu().numpy().transpose(1, 2, 0) * 255).round().astype(np.uint8)
    Image.fromarray(recon).save(out / "recon.png")
    Image.fromarray(orig).save(out / "original.png")
    recon.tofile(out / "ref_rgb.bin")

    mse = float(np.mean((recon.astype(np.float64) - orig.astype(np.float64)) ** 2))
    psnr = 99.0 if mse < 1e-9 else 10 * math.log10(255 * 255 / mse)
    bpp = (len(payload) * 8) / (size * size)

    # side-by-side (original | integer-decoded) for a demo, + metadata
    Image.fromarray(np.concatenate([orig, recon], axis=1)).save(out / "compare.png")
    (out / "meta.json").write_text(json.dumps({
        "width": size, "height": size, "in_channels": C, "s_log2": s_log2,
        "payload_bytes": len(payload), "psnr_db": round(psnr, 2), "bpp": round(bpp, 3),
    }, indent=2))
    print(f"\nwrote fixture → {out}")
    print(f"  model shifts {[l['shift'] for l in layers]} · latent {C}×{LS}×{LS} int8 ({len(payload)} B payload)")
    print(f"  real-image reconstruction: PSNR {psnr:.2f} dB · {bpp:.2f} bpp")
    print(f"  the integer decode is bit-identical on every machine; Rust int_synth must match ref_rgb.bin exactly")
    return 0


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