"""Write a `wai.neural.bmshj2018` parameter set that a sink can decode on its CPU, on
the platform's neural processor through the platform's model runtime, and in a browser
through WebNN (extensions/neural-dispatch.md §6), and one payload to decode.

The set, in one directory (`--out`):

  decoder.onnx          the synthesis transform g_s, float32, for the CPU path
                        (the same graph `wai_bmshj2018_export_onnx.py` writes)
  cdfs_q<N>.json        the entropy bottleneck's tables (`wai_bmshj2018_extract_cdfs.py`)
  decoder.mlmodel       g_s as an ML Program for the platform's model runtime, computed
  decoder.weight.bin    in float16, its input fixed at [1, C, S, S]: the program's
                        specification and its weight file, which a sink places in a
                        package of its own (neural-dispatch §6.1)
  decoder.webnn.json    g_s as a list of layers over one float32 file, for a browser's
  decoder.f32           WebNN graph builder (neural-dispatch §6.2)
  wai-usage.json        the set's usage statement (prior-carriage §18.1): the licence
                        the weights are published under

and beside it the payload `vector.bmshj2018` (a 256 x 256 picture this tool draws from a
seed, coded with the model's own encoder) and `vector.json`, what it was made from.

The weights are CompressAI's pretrained bmshj2018-factorized model, published with the
library under its licence, BSD-3-Clause-Clear; the statement records it. GDN's
reparameterised beta and gamma are taken as values once, so the converters see plain
convolutions, multiplications and square roots.

Usage (needs torch, compressai, coremltools, onnx; the model runtime's conversion runs on
any host, its compile and the compute plan only on the platform):

  python3 tools/wai_bmshj2018_npu_set.py --out <dir> [--quality 3] [--size 256]
"""
from __future__ import annotations

import argparse
import json
import shutil
import struct
import tempfile
from pathlib import Path

import numpy as np
import torch
from compressai.layers import GDN
from compressai.zoo import models

LICENCE = "BSD-3-Clause-Clear"


class PlainGDN(torch.nn.Module):
    """GDN / IGDN with its reparameterised parameters as constants."""

    def __init__(self, g: GDN):
        super().__init__()
        c = g.beta.shape[0]
        self.register_buffer("beta", g.beta_reparam(g.beta).detach().clone())
        self.register_buffer("gamma", g.gamma_reparam(g.gamma).detach().clone().reshape(c, c, 1, 1))
        self.inverse = g.inverse

    def forward(self, x):
        n = torch.sqrt(torch.nn.functional.conv2d(x * x, self.gamma, self.beta))
        return x * n if self.inverse else x / n


class Synthesis(torch.nn.Module):
    """g_s, clamped to [0, 1]: the contract of `wai.neural.bmshj2018`'s decoder."""

    def __init__(self, g_s):
        super().__init__()
        self.g_s = g_s

    def forward(self, latents):
        return self.g_s(latents).clamp(0.0, 1.0)


def plain(g_s) -> torch.nn.Sequential:
    return torch.nn.Sequential(*[PlainGDN(m) if isinstance(m, GDN) else m for m in g_s])


def picture(size: int, seed: int) -> np.ndarray:
    """A picture drawn from `seed`: two gradients and twelve discs, RGB in [0, 1]."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[0:size, 0:size].astype(np.float32) / (size - 1)
    img = np.stack([xx, yy, 0.5 + 0.5 * np.sin(6.0 * (xx + yy))], axis=-1) * 0.6 + 0.2
    for _ in range(12):
        cx, cy, r = rng.uniform(0, 1), rng.uniform(0, 1), rng.uniform(0.04, 0.2)
        col = rng.uniform(0, 1, 3).astype(np.float32)
        img[(xx - cx) ** 2 + (yy - cy) ** 2 < r * r] = col
    return np.clip(img, 0.0, 1.0)


def webnn_layers(seq: torch.nn.Sequential):
    """The layer list and the float32 values decoder.webnn.json and decoder.f32 hold."""
    values: list[np.ndarray] = []
    offset = 0

    def put(t: torch.Tensor):
        nonlocal offset
        a = t.detach().cpu().float().numpy().astype("<f4")
        values.append(a.reshape(-1))
        ref = {"offset": offset, "shape": list(a.shape)}
        offset += a.size
        return ref

    layers = []
    for m in seq:
        if isinstance(m, torch.nn.ConvTranspose2d):
            p, op = m.padding, m.output_padding
            layers.append({
                "op": "conv_transpose2d", "weight": put(m.weight), "bias": put(m.bias),
                "strides": list(m.stride), "padding": [p[0], p[0], p[1], p[1]],
                "output_padding": list(op),
            })
        elif isinstance(m, PlainGDN):
            layers.append({"op": "igdn" if m.inverse else "gdn", "gamma": put(m.gamma), "beta": put(m.beta)})
        else:
            raise SystemExit(f"g_s holds a layer this form does not write: {m!r}")
    return layers, np.concatenate(values)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--quality", type=int, default=3)
    ap.add_argument("--size", type=int, default=256, help="the picture's side, a multiple of 16")
    ap.add_argument("--seed", type=int, default=20261109)
    args = ap.parse_args()
    if args.size % 16:
        raise SystemExit("--size must be a multiple of 16")
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    net = models["bmshj2018-factorized"](quality=args.quality, pretrained=True).eval()
    net.update(force=True)
    seq = plain(net.g_s)
    syn_plain = Synthesis(seq).eval()
    syn_onnx = Synthesis(net.g_s).eval()

    # The payload: the picture coded by the model's own encoder.
    img = picture(args.size, args.seed)
    x = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
    with torch.no_grad():
        coded = net.compress(x)
        y_hat = net.entropy_bottleneck.decompress(coded["strings"][0], coded["shape"])
    rans = coded["strings"][0][0]
    s = coded["shape"][0]
    c = int(net.entropy_bottleneck._quantized_cdf.shape[0])
    s_log2 = int(s).bit_length() - 1
    assert 1 << s_log2 == s and coded["shape"][1] == s
    payload = struct.pack("<IIIIBB", args.size, args.size, len(rans), args.quality, c, s_log2) + rans
    (out / "vector.bmshj2018").write_bytes(payload)

    # The entropy tables (wai_bmshj2018_extract_cdfs.py's form).
    eb = net.entropy_bottleneck
    cdf_len = eb._cdf_length.cpu().int().tolist()
    cdf = eb._quantized_cdf.cpu().int().tolist()
    blob = {
        "model": "bmshj2018-factorized", "quality": args.quality, "precision": 16, "bypass_precision": 4,
        "channels": len(cdf_len), "cdf_length": cdf_len, "offset": eb._offset.cpu().int().tolist(),
        "medians": eb._get_medians().detach().reshape(-1).cpu().float().tolist(),
        "cdf": [row[:cdf_len[i]] for i, row in enumerate(cdf)],
    }
    (out / f"cdfs_q{args.quality}.json").write_text(json.dumps(blob, separators=(",", ":")))

    with torch.no_grad():
        ref = syn_onnx(y_hat)
        diff = float((syn_plain(y_hat) - ref).abs().max())
    if diff != 0.0:
        raise SystemExit(f"the plain g_s differs from CompressAI's by {diff}")

    # The CPU path's ONNX graph, as wai_bmshj2018_export_onnx.py writes it.
    torch.onnx.export(
        syn_onnx, (y_hat,), (out / "decoder.onnx").as_posix(),
        input_names=["latents"], output_names=["image"],
        dynamic_axes={"latents": {2: "S", 3: "S"}, "image": {2: "H", 3: "W"}},
        opset_version=17, do_constant_folding=True, dynamo=False,
    )

    # The ML Program, float16, input fixed at the payload's latent shape.
    import coremltools as ct

    traced = torch.jit.trace(syn_plain, y_hat)
    prog = ct.convert(
        traced, inputs=[ct.TensorType(name="latents", shape=tuple(y_hat.shape))],
        outputs=[ct.TensorType(name="image")], convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16, minimum_deployment_target=ct.target.macOS14,
    )
    with tempfile.TemporaryDirectory() as t:
        pkg = Path(t) / "g_s.mlpackage"
        prog.save(pkg.as_posix())
        data = pkg / "Data" / "com.apple.CoreML"  # allow-name: the package format's fixed directory name
        shutil.copyfile(data / "model.mlmodel", out / "decoder.mlmodel")
        shutil.copyfile(data / "weights" / "weight.bin", out / "decoder.weight.bin")

    # The WebNN form.
    layers, values = webnn_layers(seq)
    (out / "decoder.f32").write_bytes(values.astype("<f4").tobytes())
    (out / "decoder.webnn.json").write_text(json.dumps({
        "format": "wai-bmshj2018-webnn/1",
        "input": {"name": "latents", "shape": list(y_hat.shape)},
        "output": {"name": "image", "shape": [1, 3, args.size, args.size]},
        "clamp": [0.0, 1.0],
        "layers": layers,
    }, indent=1) + "\n")

    (out / "wai-usage.json").write_text(json.dumps({
        "kind": "wai-set-usage", "licence": LICENCE, "terms": None,
        "content_derived": False, "redistribute": True, "ai_usage": None,
    }) + "\n")
    (out / "vector.json").write_text(json.dumps({
        "capability": "wai.neural.bmshj2018", "quality": args.quality, "size": args.size,
        "seed": args.seed, "latent": [1, c, s, s], "payload_bytes": len(payload),
        "picture": "tools/wai_bmshj2018_npu_set.py picture(size, seed)",
        "weights": "CompressAI bmshj2018-factorized, pretrained", "licence": LICENCE,
    }, indent=1) + "\n")
    for f in sorted(out.iterdir()):
        print(f"{f.name:22} {f.stat().st_size:>10,} B")
    return 0


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