"""Export JPEG-AI decoder weights from the ISO reference ONNX assets.

The byte-exact decoder in `wai-rs` (`jpegai_synthesis`, `jpegai_pipeline`,
`jpegai_icci`) is driven by weights lifted straight out of the reference
`jpegai-vm-assets` ONNX files. Two regimes, and using the wrong one silently
produces plausible-but-wrong output:

  float nets  — synthesis (model_y / model_uv), hyper_decoder, MCM stages.
                Exported as 16-bit per-tensor symmetric PTQ:
                  conv weight [O, I/g, kh, kw] -> int16 + per-tensor
                  mult = round(scale * 2^31), shift = 31
                  bias / LayerNorm / temperature -> Q16 (round(v * 2^16))
                This is what `jpegai_synthesis::Weights::from_json` expects.

  integer net — hyper_scale_decoder. Already integer end to end (int32 conv
                weights, integer-valued Pow exponents); exported VERBATIM with
                no quantization, for `jpegai_synthesis::IntNet::from_json`.
                Quantizing it would corrupt an exactly-representable net.

Reading the file is dependency-free: ONNX is protobuf, and we only need the
initializer table, so there is no `onnx` / `protobuf` / `torch` requirement.

  ModelProto.graph = 7
  GraphProto.initializer = 5
  TensorProto{ dims=1, data_type=2, name=8, raw_data=9 }
  data_type: 1 = float32, 6 = int32, 7 = int64

Usage
-----
    # float net -> Q16 PTQ (synthesis / hyper_decoder / MCM stage)
    python wai_jpegai_export_weights.py ptq  synthesis16.onnx  synth16_y.json

    # integer net -> verbatim
    python wai_jpegai_export_weights.py int  hyper_scale_decoder.onnx  hs_y.json

    # just look at what is inside
    python wai_jpegai_export_weights.py list model_15.onnx

Caveat worth knowing: the eICCI post-filter models under `post_filters/` export
the model's `forward`, which the reference decoder never calls (it calls
`process_dwt444_2` + `process_comp`), and their initializers are BatchNorm-folded
and incomplete — no `BNY`, no `scaleU`/`scaleV`, no per-residual-block `scale`.
For eICCI take the reference's PyTorch `state_dict` instead; `list` is still
useful there for inspection.
"""

import json
import struct
import sys
from collections import OrderedDict

SHIFT = 31          # conv requant shift, matches int_transformer::conv_grouped
Q16 = 1 << 16       # bias / LN / temperature fixed-point scale
INT16_MAX = 32767


# ---- dependency-free protobuf reader ----

def _varint(buf, i):
    shift = value = 0
    while True:
        byte = buf[i]
        i += 1
        value |= (byte & 0x7F) << shift
        if not byte & 0x80:
            return value, i
        shift += 7


def _fields(buf):
    """Yield (field_number, wire_type, payload) for one protobuf message."""
    i = 0
    while i < len(buf):
        tag, i = _varint(buf, i)
        field, wire = tag >> 3, tag & 7
        if wire == 0:
            value, i = _varint(buf, i)
            yield field, 0, value
        elif wire == 2:
            length, i = _varint(buf, i)
            yield field, 2, buf[i:i + length]
            i += length
        elif wire == 1:
            yield field, 1, buf[i:i + 8]
            i += 8
        elif wire == 5:
            yield field, 5, buf[i:i + 4]
            i += 4
        else:
            raise ValueError(f"unsupported wire type {wire}")


def _parse_tensor(buf):
    dims, name, dtype, raw = [], "", 0, b""
    for field, wire, payload in _fields(buf):
        if field == 1:
            if wire == 0:
                dims.append(payload)
            elif wire == 2:                      # packed dims
                j = 0
                while j < len(payload):
                    d, j = _varint(payload, j)
                    dims.append(d)
        elif field == 2 and wire == 0:
            dtype = payload
        elif field == 8 and wire == 2:
            name = payload.decode("utf-8", "replace")
        elif field == 9 and wire == 2:
            raw = payload
    return name, dims, dtype, raw


def initializers(path):
    """name -> (dims, data_type, raw_bytes) for every initializer, in file order."""
    out = OrderedDict()
    model = open(path, "rb").read()
    for field, wire, payload in _fields(model):
        if field == 7 and wire == 2:             # ModelProto.graph
            for gfield, gwire, gpayload in _fields(payload):
                if gfield == 5 and gwire == 2:   # GraphProto.initializer
                    name, dims, dtype, raw = _parse_tensor(gpayload)
                    out[name] = (dims, dtype, raw)
    return out


def _prod(xs):
    n = 1
    for x in xs:
        n *= x
    return n


def _decode(dims, dtype, raw):
    """raw bytes -> list of numbers, or None if the tensor is empty/unsupported."""
    n = _prod(dims)
    if not raw:
        return None
    if dtype == 1 and len(raw) == 4 * n:
        return list(struct.unpack(f"<{n}f", raw))
    if dtype == 6 and len(raw) == 4 * n:
        return list(struct.unpack(f"<{n}i", raw))
    if dtype == 7 and len(raw) == 8 * n:
        return list(struct.unpack(f"<{n}q", raw))
    return None


# ---- export modes ----

def export_ptq(src, dst, strip_prefix=True):
    """Float net -> 16-bit symmetric PTQ, the layout `Weights::from_json` reads."""
    import re
    result = OrderedDict()
    n_conv = n_scalar = total = 0
    worst = 0.0
    for name, (dims, dtype, raw) in initializers(src).items():
        values = _decode(dims, dtype, raw)
        if values is None or dtype != 1:
            continue
        total += len(values)
        # the reference wraps modules as `hop_<x>.inst.<name>`; the decoder
        # addresses them by the bare module path
        key = re.sub(r"^hop_\w+\.inst\.", "", name) if strip_prefix else name
        if len(dims) == 4 and dims[0] > 1:       # conv weight [O, I/g, kh, kw]
            amax = max((abs(v) for v in values), default=0.0) or 1e-12
            scale = amax / INT16_MAX
            quant = [max(-32768, min(INT16_MAX, round(v / scale))) for v in values]
            worst = max(worst, max(abs(v - round(v / scale) * scale) for v in values))
            result[key] = {"kind": "conv", "dims": dims, "shift": SHIFT,
                           "mult": round(scale * (1 << SHIFT)), "w": quant}
            n_conv += 1
        else:                                    # bias / LN gamma,beta / temperature
            result[key] = {"kind": "q16", "dims": dims,
                           "v": [round(v * Q16) for v in values]}
            n_scalar += 1
    json.dump(result, open(dst, "w"))
    print(f"{n_conv} conv + {n_scalar} scalar tensors ({total:,} params) -> {dst}")
    print(f"conv round-trip worst |error| = {worst:.3e} "
          f"({'lossless at 16 bits' if worst < 1e-3 else 'CHECK'})")


def export_int(src, dst):
    """Integer net -> verbatim. Pow exponents are integer-valued floats."""
    result = OrderedDict()
    for name, (dims, dtype, raw) in initializers(src).items():
        values = _decode(dims, dtype, raw)
        if values is None:
            continue
        if dtype == 1:                           # integer-valued float (Pow exponent)
            values = [int(round(v)) for v in values]
        result[name] = {"dims": dims, "v": values}
    json.dump(result, open(dst, "w"))
    print(f"{len(result)} integer tensors -> {dst}")
    for name, entry in result.items():
        lo, hi = min(entry["v"]), max(entry["v"])
        print(f"  {name:24s} dims={entry['dims']} range=[{lo}, {hi}]")


def list_tensors(src):
    dtypes = {1: "f32", 6: "i32", 7: "i64"}
    total = 0
    for name, (dims, dtype, raw) in initializers(src).items():
        n = _prod(dims)
        total += n
        print(f"  {name:38s} {dtypes.get(dtype, dtype):4s} dims={dims}")
    print(f"{total:,} params")


def main(argv):
    if len(argv) < 3:
        print(__doc__)
        return 1
    mode, src = argv[1], argv[2]
    if mode == "list":
        list_tensors(src)
    elif mode == "ptq":
        export_ptq(src, argv[3])
    elif mode == "int":
        export_int(src, argv[3])
    else:
        print(f"unknown mode {mode!r}; expected list | ptq | int")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))