#!/usr/bin/env python3
"""The fixture dcvc-drift decodes: the `wai.neural.dcvc_rt` fixture of
tools/wai_validate_synthetic.py, written to a directory of its own, for as many frames as
asked, with the entropy buckets the encoder chose beside it.

    python make_fixture.py --out DIR [--frames N] [--ref-gain G] [--image-gain G] [--check-harness]

The model, its seed, its export, the encode and every ONNX session are the harness's
(`build(MiniDCVC, 2)`, `export_onnx`, `fblob`, `gblob`, `_TDec`, `_session`, `write_wai`).
Four things are added:

- `--frames N` codes N frames. Frame k is the test picture rolled 8 samples right per frame
  and, every 32 frames, 8 samples down, so the first 32 frames are the harness's motion and
  no two of the first 1024 are one picture.
- `--ref-gain G` multiplies the weights the reference feature enters the next feature
  through (the second half of `t_feat`'s input channels). The harness's weights grow the
  feature by about the square root of 3 a frame, which three frames do not notice and a
  longer run saturates; G is what holds it level. G = 1 leaves the weights as they are.
- `--image-gain G` multiplies the last layer of the image head by G and centres it on 0.5.
  The harness's image is clamped in more than nine samples of ten, and a clamped sample
  shows no difference at all; G keeps the samples in range, so a difference in the float
  output can reach them. G = 1 leaves the head as it is.
- `encoder-buckets.bin` holds, frame after frame, the Gaussian bucket the encoder coded each
  latent element with, one byte an element, and `fixture.json` the digests and the
  encoder's readings per frame.

`--check-harness` also runs the harness's own `roundtrip_dcvc` and requires this script's
envelope and reference to be those bytes. It needs `--frames 3` and both gains at 1, and it
writes the harness's fixed paths under /tmp, as the harness does.

Needs torch, compressai, onnxruntime, numpy and Pillow, as the harness does.
"""
from __future__ import annotations

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

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import wai_validate_synthetic as v  # noqa: E402

SIZE, SEED, CF, FPS_X1000 = 256, 2, 48, 30000


def sha256(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def picture(out: Path) -> torch.Tensor:
    """The harness's test picture, as the harness's `main` reads it."""
    from PIL import Image

    path = out / "picture.png"
    v.test_image(path, SIZE)
    img = Image.open(path).convert("RGB").resize((SIZE, SIZE), Image.LANCZOS)
    path.unlink()
    return torch.from_numpy(np.asarray(img, np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)


def frame_of(x: torch.Tensor, k: int) -> np.ndarray:
    base = np.asarray(x[0].permute(1, 2, 0))
    return np.roll(np.roll(base, shift=8 * k, axis=1), shift=8 * (k // 32), axis=0)


def encode(net, x: torch.Tensor, out: Path, n_frames: int) -> dict:
    set_dir = out / "set"
    set_dir.mkdir(parents=True, exist_ok=True)
    gc = net.gaussian_conditional
    v.export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), set_dir / "hyper_synthesis.onnx")
    torch.onnx.export(v._TDec(net).eval(), (torch.zeros(1, 48, 16, 16), torch.zeros(1, CF, 16, 16)),
                      str(set_dir / "temporal_decoder.onnx"), input_names=["latents", "ref_feature"],
                      output_names=["image", "feature"],
                      dynamic_axes={"latents": {2: "h", 3: "w"}, "ref_feature": {2: "h", 3: "w"},
                                    "image": {2: "h", 3: "w"}, "feature": {2: "h", 3: "w"}},
                      opset_version=17, dynamo=False)
    (set_dir / "hyper_cdfs_q3.json").write_text(json.dumps(v.fblob(net.entropy_bottleneck), separators=(",", ":")))
    (set_dir / "gauss_cdfs_q3.json").write_text(json.dumps(v.gblob(gc), separators=(",", ":")))
    hs = v._session(set_dir / "hyper_synthesis.onnx")
    td = v._session(set_dir / "temporal_decoder.onnx")

    feat = torch.zeros(1, CF, 16, 16)
    codes, ref, buckets, frames = [], [], [], []
    for k in range(n_frames):
        xi = torch.from_numpy(np.ascontiguousarray(frame_of(x, k))).permute(2, 0, 1).unsqueeze(0)
        with torch.no_grad():
            y = net.g_a(xi)
            z = net.h_a(y)
            zs = net.entropy_bottleneck.compress(z)
            z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
            params = v.onx(hs, z_hat)
            sc, mn = params.chunk(2, 1)
            idx = gc.build_indexes(sc)
            ys = gc.compress(y, idx, means=mn)
            y_hat = gc.decompress(ys, idx, means=mn)
            img, feat = [torch.from_numpy(t) for t in td.run(None, {"latents": y_hat.numpy(), "ref_feature": feat.numpy()})]
        rgb = (np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8)
        codes.append((zs[0], ys[0]))
        ref.append(rgb)
        buckets.append(idx[0].numpy().astype(np.uint8).reshape(-1).tobytes())
        clipped = float(((img < 0) | (img > 1)).float().mean())
        frames.append({"frame": k, "z_bytes": len(zs[0]), "y_bytes": len(ys[0]),
                       "feature_rms": float(feat.pow(2).mean().sqrt()), "image_clipped": clipped,
                       "rgb_sha256": sha256(rgb.tobytes()), "buckets_sha256": sha256(buckets[-1])})
    _, cz, _, _ = z.shape
    _, cy, _, _ = y.shape
    hdr = struct.pack("<IIIIHHBBHH", SIZE, SIZE, n_frames, FPS_X1000, cy, cz, 4, 2, 3, CF)
    table = b"".join(struct.pack("<II", len(zb), len(yb)) for zb, yb in codes)
    body = b"".join(zb + yb for zb, yb in codes)
    v.write_wai(out / "test.dcvc.wai", "wai.neural.dcvc_rt", "dcvc_rt", hdr, b"", table + body, SIZE, SIZE)
    (out / "ref.rgb").write_bytes(b"".join(r.tobytes() for r in ref))
    (out / "encoder-buckets.bin").write_bytes(b"".join(buckets))
    return {"frames": frames, "elements": cy * 16 * 16}


def check_harness(x: torch.Tensor, out: Path) -> None:
    v.roundtrip_dcvc(v.build(v.MiniDCVC, SEED), x, out / "harness", SIZE)
    for ours, theirs in (("test.dcvc.wai", "/tmp/test.dcvc.wai"), ("ref.rgb", "/tmp/ref_dcvc.rgb")):
        a, b = (out / ours).read_bytes(), Path(theirs).read_bytes()
        if a != b:
            raise SystemExit(f"{ours}: {len(a)} bytes differ from the harness's {theirs} ({len(b)} bytes)")
        print(f"{ours} is the harness's {theirs}, {len(a)} bytes")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--frames", type=int, default=3)
    ap.add_argument("--ref-gain", type=float, default=1.0)
    ap.add_argument("--image-gain", type=float, default=1.0)
    ap.add_argument("--check-harness", action="store_true")
    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    x = picture(args.out)
    net = v.build(v.MiniDCVC, SEED)
    if args.ref_gain != 1.0:
        with torch.no_grad():
            net.t_feat.weight[:, 48:, :, :] *= args.ref_gain
    if args.image_gain != 1.0:
        with torch.no_grad():
            net.t_img[-1].weight *= args.image_gain
            net.t_img[-1].bias.mul_(args.image_gain).add_(0.5)
    meta = encode(net, x, args.out, args.frames)
    import compressai
    import onnxruntime

    record = {"capability": "wai.neural.dcvc_rt", "harness": "tools/wai_validate_synthetic.py", "seed": SEED,
              "size": SIZE, "frames_coded": args.frames, "ref_gain": args.ref_gain, "image_gain": args.image_gain,
              "stack": {"torch": torch.__version__, "compressai": compressai.__version__,
                        "onnxruntime": onnxruntime.__version__, "numpy": np.__version__},
              "files": {p.relative_to(args.out).as_posix(): sha256(p.read_bytes())
                        for p in sorted(args.out.rglob("*")) if p.is_file() and p.name != "fixture.json"},
              **meta}
    (args.out / "fixture.json").write_text(json.dumps(record, indent=1) + "\n")
    last = meta["frames"][-1]
    print(f"{args.frames} frames, ref gain {args.ref_gain}, image gain {args.image_gain}: feature rms {last['feature_rms']:.4g}, "
          f"clipped {last['image_clipped']:.3f}, payload {sum(f['z_bytes'] + f['y_bytes'] for f in meta['frames'])} bytes")
    if args.check_harness:
        check_harness(x, args.out)
    return 0


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