#!/usr/bin/env python3
"""Train an int_synth decoder and emit the committable conformance vector `photo` of
`neural-conformance/`: the integer model (`int_synth.json`), a coded picture
(`latent.bin`, the wire payload), the integer-forward reference RGB (`ref_rgb.bin`) and
`meta.json`. The Rust decoder (`int_synth.rs`) reproduces the reference BIT-EXACTLY on
every machine: the determinism is exact by construction, and this vector shows it on a
model trained in float and quantised after training, not only on seeded weights.

**The pictures are generated here; the tool reads no image file.** The pictures the
decoder is trained on, and the picture the vector codes, are written by `picture()`
below from a seed: 128 × 128 RGB8 dead-leaves pictures. A picture starts as a vertical
gradient between two colours; then 400 discs are painted back to front, each with a
centre uniform over the picture, a radius ⌊32·16 / (16 + u)⌋ (at least 2) for u uniform
in 0..239, a colour each channel of which is uniform in 16..239, and a linear shading
⌊(gx·dx + gy·dy) / 8⌋ for gx, gy uniform in −3..3, where (dx, dy) is the pixel's offset
from the centre; then every sample gets noise uniform in −2..2, and every value is
clamped to 0..255. Every draw is a SplitMix64 counter stream keyed by (seed, stream),
and the generator is integer arithmetic only, so a seed gives the same pictures on
every machine and every Python. The coded picture is stream 0 of the seed and the
training pictures are streams 1 to N: the decoder never sees the picture it codes.

Pipeline (the working recipe; see wai_int_synth_train.py for why training through the
`>> shift` fails): float-train an autoencoder on the training pictures, with torch
seeded from the same seed, deterministic algorithms and a fixed thread count on the
CPU: Adam steps on batches drawn by a generator seeded from the same seed, the learning
rate falling to a tenth for the last quarter of the steps, on the squared error of the
decoder's output before the clamp, so that an output below 0 or above 1 still has a
gradient; then quantise the DECODER after training, with each layer's activation scale
measured on the training pictures (its largest magnitude over 127, and 1/255 at the
output), its shift the largest that keeps its weights within ±127, and its weight scale
the one that makes the shift carry that activation scale exactly, so the integer
pipeline tracks the float one. Encode
the coded picture through the float encoder, quantise 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 to 0..255).
That numpy decode is the reference the Rust decoder must match to max|Δ|=0.

  python tools/wai_int_synth_real.py --out <dir> [--seed S] [--pictures N] [--steps K]
  python tools/wai_int_synth_real.py --picture-only <out.ppm> [--seed S] [--stream K]
  python tools/wai_int_synth_real.py --check <vector-dir>

Training needs numpy and torch. `--picture-only` writes one generated picture as a
binary PPM, and `--check` regenerates the picture a committed vector codes from the
seed in its `meta.json`, requires the vector's reference decode (`ref_rgb.bin`) to
reconstruct it at the PSNR `meta.json` records, and requires the vector's record in the
corpus's `licences.json` to state the seed and streams `meta.json` records; both need
only the Python standard library. The committed vector's command is in
`neural-conformance/README.md`.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import struct
import sys
from pathlib import Path

SIDE = 128                       # the pictures' side; the model's output size
DISCS = 400                      # discs per picture
THREADS = 4                      # fixed: CPU reductions depend on it
SEED = 20261104                  # the committed vector's seed
PICTURES = 512                   # training pictures: streams 1 to PICTURES
STEPS, BATCH, LR = 3000, 32, 2e-4

# ---- the generated pictures (integer arithmetic only) --------------------------------
M64 = (1 << 64) - 1
GOLDEN = 0x9E3779B97F4A7C15


def _mix64(x):
    """SplitMix64's finalizer."""
    x &= M64
    x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & M64
    x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & M64
    return x ^ (x >> 31)


class Rng:
    """A SplitMix64 stream, keyed by (seed, stream)."""

    def __init__(self, seed, stream):
        self.s = _mix64(seed * GOLDEN + stream * 0xD1B54A32D192ED03)

    def next(self):
        self.s = (self.s + GOLDEN) & M64
        return _mix64(self.s)

    def int(self, lo, hi):
        """Uniform integer in [lo, hi]."""
        return lo + self.next() % (hi - lo + 1)


def picture(seed, stream, side=SIDE):
    """One dead-leaves picture as RGB8 bytes: rows top to bottom, pixels left to right,
    R, G, B. The module docstring states the formula."""
    r = Rng(seed, stream)
    top = [r.int(16, 239) for _ in range(3)]
    bottom = [r.int(16, 239) for _ in range(3)]
    px = bytearray(side * side * 3)
    for y in range(side):
        row = [(top[c] * (side - 1 - y) + bottom[c] * y) // (side - 1) for c in range(3)]
        for x in range(side):
            p = (y * side + x) * 3
            px[p:p + 3] = bytes(row)
    rmax = side // 4
    for _ in range(DISCS):
        cx, cy = r.int(0, side - 1), r.int(0, side - 1)
        rad = max(2, rmax * 16 // (16 + r.int(0, 239)))
        base = [r.int(16, 239) for _ in range(3)]
        gx, gy = r.int(-3, 3), r.int(-3, 3)
        for y in range(max(0, cy - rad), min(side, cy + rad + 1)):
            dy = y - cy
            for x in range(max(0, cx - rad), min(side, cx + rad + 1)):
                dx = x - cx
                if dx * dx + dy * dy > rad * rad:
                    continue
                shade = (gx * dx + gy * dy) // 8
                p = (y * side + x) * 3
                for c in range(3):
                    px[p + c] = min(255, max(0, base[c] + shade))
    for i in range(len(px)):
        px[i] = min(255, max(0, px[i] + r.int(-2, 2)))
    return bytes(px)


def write_ppm(path, rgb, w, h):
    Path(path).write_bytes(b"P6\n%d %d\n255\n" % (w, h) + bytes(rgb))


def psnr_db(a, b):
    """PSNR of two RGB8 buffers, in dB, from their integer sum of squared differences."""
    assert len(a) == len(b) and len(a) > 0
    sse = sum((x - y) * (x - y) for x, y in zip(a, b))
    return 99.0 if sse == 0 else 10 * math.log10(255 * 255 * len(a) / sse)


# ---- the integer forward, byte-for-byte as int_synth.rs runs it ----
def conv_transpose2d_int(inp, layer):
    import numpy as np

    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, layers, h, w):
    import numpy as np

    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


GENERATOR = "tools/wai_int_synth_real.py"   # this tool, as a licence record names it


def licence_record_fails(vdir, meta):
    """The vector's record in the corpus's `licences.json` (two directories up) states
    the pictures meta.json says the vector was made from: content generated by this
    tool, the coded picture `seed S, stream K`, and the training pictures
    `seed S, streams A to B` from meta.json's `training_streams`, a range that leaves
    out the coded picture's stream."""
    fails = []
    seed, stream = meta["seed"], meta["stream"]
    tr_streams = meta.get("training_streams")
    if not (isinstance(tr_streams, list) and len(tr_streams) == 2
            and all(isinstance(x, int) for x in tr_streams) and tr_streams[0] <= tr_streams[1]):
        return [f"meta.json's training_streams {tr_streams!r} is not [first, last]"]
    a, b = tr_streams
    if a <= stream <= b:
        fails.append(f"meta.json's training streams {a} to {b} include the coded picture's stream {stream}")
    try:
        records = json.loads((vdir.parent.parent / "licences.json").read_text())["vectors"]
        rec = records[vdir.name]
        tr = rec["training"]
    except (OSError, ValueError, KeyError, TypeError) as e:
        return fails + [f"licences.json has no record with a training record for {vdir.name} ({e!r})"]
    want = {
        "": f"the dead-leaves picture the tool generates from seed {seed}, stream {stream}",
        "training ": f"the dead-leaves pictures the tool generates from seed {seed}, streams {a} to {b}",
    }
    for label, r in (("", rec), ("training ", tr)):
        if not isinstance(r, dict):
            fails.append(f"licences.json's {label}record for {vdir.name} is not an object")
            continue
        if r.get("content") != "generated" or r.get("generator") != GENERATOR:
            fails.append(f"licences.json's {label}record for {vdir.name} does not say its content is "
                         f"generated by {GENERATOR}")
        if r.get("input") != want[label]:
            fails.append(f"licences.json's {label}input for {vdir.name} is {r.get('input')!r}; "
                         f"meta.json says {want[label]!r}")
    return fails


def check(vdir):
    """The committed vector codes the picture its seed generates: its reference decode
    (`ref_rgb.bin`, whose SHA-256 is the corpus's digest for the vector in
    `sha256.json` two directories up) reconstructs that picture at the PSNR meta.json
    records, to two decimals; and the vector's licence record states the pictures
    meta.json says it was made from (`licence_record_fails`)."""
    vdir = Path(vdir)
    meta = json.loads((vdir / "meta.json").read_text())
    w, h = meta["width"], meta["height"]
    head = (vdir / "latent.bin").read_bytes()[:12]
    ph, pw, _c, _s, _r = struct.unpack("<IIHBB", head)
    ref = (vdir / "ref_rgb.bin").read_bytes()
    digests = json.loads((vdir.parent.parent / "sha256.json").read_text())
    fails = []
    if (w, h) != (SIDE, SIDE):
        fails.append(f"meta.json says {w}x{h}, not the generator's {SIDE}x{SIDE}")
    if (ph, pw) != (h, w):
        fails.append(f"latent.bin's header says {pw}x{ph}, meta.json {w}x{h}")
    if len(ref) != w * h * 3:
        fails.append(f"ref_rgb.bin holds {len(ref)} bytes, not {w * h * 3}")
    want = digests.get(vdir.name, {}).get("sha256")
    if hashlib.sha256(ref).hexdigest() != want:
        fails.append(f"ref_rgb.bin's SHA-256 is not sha256.json's {vdir.name} digest ({want})")
    if not fails:
        src = picture(meta["seed"], meta["stream"], w)
        got = round(psnr_db(src, ref), 2)
        print(f"{vdir}: seed {meta['seed']} stream {meta['stream']}: the reference decode reconstructs it at {got:.2f} dB")
        if got != meta["psnr_db"]:
            fails.append(f"PSNR {got:.2f} dB against the generated picture, meta.json records {meta['psnr_db']}")
    fails += licence_record_fails(vdir, meta)
    for f in fails:
        print(f"FAIL {vdir}: {f}", file=sys.stderr)
    return 1 if fails else 0


def train(args):
    import numpy as np
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    torch.manual_seed(args.seed)
    torch.use_deterministic_algorithms(True)
    torch.set_num_threads(THREADS)

    C, LS = args.chans[0], args.latent_side
    n_layers = len(args.chans) - 1
    size = LS * (2 ** n_layers)
    assert size == SIDE, f"--latent-side and --chans give {size}x{size}; the pictures are {SIDE}x{SIDE}"
    print(f"image {size}x{size} · latent {LS}x{LS}x{C} · {n_layers} decoder layers · seed {args.seed} · "
          f"{args.pictures} training pictures · {args.steps} steps of {args.batch} · lr {args.lr} · {THREADS} threads")

    def tensor(rgb):
        a = np.frombuffer(rgb, np.uint8).reshape(size, size, 3).astype(np.float32) / 255.0
        return a.transpose(2, 0, 1)

    src = picture(args.seed, 0, size)
    data = torch.from_numpy(np.stack([tensor(picture(args.seed, k, size)) for k in range(1, args.pictures + 1)]))
    test = torch.from_numpy(tensor(src)[None])

    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])
    dec = nn.ModuleList([up(args.chans[i], args.chans[i + 1]) for i in range(n_layers)])

    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)
    # The learning rate falls to a tenth for the last quarter of the steps.
    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=[args.steps * 3 // 4], gamma=0.1)
    order = torch.Generator().manual_seed(args.seed)
    for step in range(args.steps):
        batch = data[torch.randperm(len(data), generator=order)[:args.batch]]
        opt.zero_grad()
        recon = dfwd(enc(batch) * 64.0)
        # Before the clamp: clamped, an output that starts below 0 has no gradient and
        # the decoder can stay black.
        loss = F.mse_loss(recon * 255, batch * 255)
        loss.backward()
        opt.step()
        sched.step()
        if step % 250 == 0 or step == args.steps - 1:
            print(f"  step {step:5d}  batch px-mse {loss.item():7.1f}", flush=True)

    def rgb8(t):
        return (t[0].clamp(0, 1) * 255).round().to(torch.uint8).permute(1, 2, 0).contiguous().numpy().tobytes()

    # ---- calibrate activation scales on the training set, then quantise ----
    with torch.no_grad():
        print(f"  float decode of the coded picture: PSNR {psnr_db(src, rgb8(dfwd(enc(test) * 64.0))):.2f} dB")
        s0max, amax = 0.0, [0.0] * n_layers
        for k in range(0, len(data), args.batch):
            x = enc(data[k:k + args.batch]) * 64.0
            s0max = max(s0max, float(x.abs().max()))
            for i, c in enumerate(dec):
                x = c(x)
                if i < len(dec) - 1:
                    x = torch.relu(x)
                amax[i] = max(amax[i], float(x.abs().max()))
        s0 = s0max / 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):
        # The layer's output is acc · sw · sin / 2^shift, so its scale is the calibrated
        # one exactly when sw = sout / (sin · 2^shift). The shift is the largest that
        # keeps every weight within ±127 under that sw. Rounding the shift alone would
        # leave each layer's gain off by up to √2, and the errors would multiply down
        # the stack into the output's brightness and contrast.
        sout = scales[i + 1]
        wmax = float(c.weight.detach().abs().max())
        shift = math.floor(math.log2(sout * 127 / (wmax * sin)))
        if shift < 0:
            print(f"  layer {i}: no shift >= 0 keeps the weights within ±127; clamping them")
            shift = 0
        sw = sout / (sin * 2 ** shift)
        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 = sout

    # ---- encode the coded picture → int8 latent → wire payload ----
    with torch.no_grad():
        lat = (enc(test) * 64.0)[0].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)
    torch.save({"enc": enc.state_dict(), "dec": dec.state_dict()}, out / "float_model.pt")
    (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).tobytes()
    (out / "ref_rgb.bin").write_bytes(recon)
    write_ppm(out / "original.ppm", src, size, size)
    write_ppm(out / "recon.ppm", recon, size, size)

    psnr = psnr_db(src, recon)
    bpp = (len(payload) * 8) / (size * size)
    (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),
        "seed": args.seed, "stream": 0, "training_streams": [1, args.pictures],
        "steps": args.steps, "batch": args.batch, "lr": args.lr, "threads": THREADS,
        "torch": torch.__version__, "numpy": np.__version__,
    }, indent=2) + "\n")
    print(f"\nwrote the vector → {out}")
    print(f"  model shifts {[l['shift'] for l in layers]} · latent {C}x{LS}x{LS} int8 ({len(payload)} B payload)")
    print(f"  reconstruction of the coded picture: PSNR {psnr:.2f} dB · {bpp:.2f} bpp")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="Train an int_synth decoder on generated pictures and write the corpus vector.")
    ap.add_argument("--out", help="output dir: int_synth.json, latent.bin, ref_rgb.bin and meta.json (the vector), "
                    "two previews and the float model (float_model.pt)")
    ap.add_argument("--seed", type=int, default=SEED)
    ap.add_argument("--pictures", type=int, default=PICTURES, help="training pictures: streams 1 to N of the seed")
    ap.add_argument("--chans", type=int, nargs="+", default=[96, 72, 48, 24, 3])
    ap.add_argument("--latent-side", type=int, default=8)
    ap.add_argument("--steps", type=int, default=STEPS, help="Adam steps, each on a batch drawn from the training pictures")
    ap.add_argument("--batch", type=int, default=BATCH)
    ap.add_argument("--lr", type=float, default=LR)
    ap.add_argument("--picture-only", metavar="OUT.ppm", help="write one generated picture and exit")
    ap.add_argument("--stream", type=int, default=0, help="with --picture-only: 0 the coded picture, 1 to N the training pictures")
    ap.add_argument("--check", metavar="VECTOR_DIR", help="check that a committed vector codes the picture its seed generates")
    args = ap.parse_args()
    if args.check:
        return check(args.check)
    if args.picture_only:
        write_ppm(args.picture_only, picture(args.seed, args.stream), SIDE, SIDE)
        return 0
    if not args.out:
        ap.error("--out is required unless --picture-only or --check is given")
    return train(args)


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