"""PyTorch checkpoint -> ONNX (+ parity check, + optional INT8 dynamic quantization).

The exported graph takes (state, question_ids, choice_ids) with dynamic batch,
number of choices and state size, and returns (probs, logits, embedding); the
calibration temperature is baked in. Everything the runtime needs to
pre-process inputs is stored in the model's metadata (key "minijev").
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch

from minijev.inference.runtime import META_KEY, OnnxDecider, TorchDecider
from minijev.model import STATE_INPUTS, ExportWrapper


def export_onnx(checkpoint: str | Path, out: str | Path | None = None, opset: int = 18,
                verify: bool = True, quantize: bool = True) -> dict:
    import onnx

    checkpoint = Path(checkpoint)
    out = Path(out) if out else checkpoint.with_suffix(".onnx")
    torch_dec = TorchDecider(checkpoint)
    model, cfg, task = torch_dec.model, torch_dec.cfg, torch_dec.task
    state_key = STATE_INPUTS[task.state_encoder]
    wrapper = ExportWrapper(model, state_key).eval()

    states = task.example_states(4)
    question = task.default_question
    choices = list(task.default_choices)
    x = torch_dec.encode(states)
    q, c = torch_dec._tokens(question, tuple(choices))
    B = len(states)
    args = (torch.as_tensor(x), torch.as_tensor(np.repeat(q, B, 0)), torch.as_tensor(np.repeat(c, B, 0)))

    state_axes = {0: "batch", 2: "height", 3: "width"} if x.ndim == 4 else {0: "batch", 1: "length"}
    dynamic_axes = {state_key: state_axes, "question_ids": {0: "batch"},
                    "choice_ids": {0: "batch", 1: "n_choices"},
                    "probs": {0: "batch", 1: "n_choices"}, "logits": {0: "batch", 1: "n_choices"},
                    "embedding": {0: "batch"}}
    # TorchScript-based exporter: mature, handles SDPA with masks and dynamic axes well
    torch.onnx.export(wrapper, args, str(out), input_names=[state_key, "question_ids", "choice_ids"],
                      output_names=["probs", "logits", "embedding"], dynamic_axes=dynamic_axes,
                      opset_version=opset, dynamo=False)
    meta = {"task": task.name, "task_cfg": cfg["task"], "state_input": state_key,
            "max_question_len": torch_dec.max_question_len, "max_choice_len": torch_dec.max_choice_len,
            "temperature": model.temperature, "params": model.num_parameters(), "config_name": cfg.get("name"),
            "default_question": question, "default_choices": choices}
    m = onnx.load(str(out))
    entry = m.metadata_props.add()
    entry.key, entry.value = META_KEY, json.dumps(meta)
    onnx.save(m, str(out))
    report = {"onnx": str(out), "size_mb": round(out.stat().st_size / 2**20, 3), "meta": meta}
    if verify:
        report["parity_fp32"] = parity(torch_dec, OnnxDecider(out), task)
    if quantize:
        q_out = out.with_name(out.stem + ".int8.onnx")
        quantize_int8(out, q_out)
        report["onnx_int8"] = str(q_out)
        report["size_int8_mb"] = round(q_out.stat().st_size / 2**20, 3)
        if verify:
            report["parity_int8"] = parity(torch_dec, OnnxDecider(q_out), task)
    return report


def parity(ref, other, task, n: int = 256, seed: int = 1) -> dict:
    """Compare probabilities on real states, including 2- and 3-choice subsets."""
    states = task.example_states(n, seed=seed)
    choices = list(task.default_choices)
    res = {}
    for name, ch in (("all", choices), ("subset", choices[: max(2, len(choices) - 1)])):
        p_ref = ref.probs(states, task.default_question, ch)
        p_oth = other.probs(states, task.default_question, ch)
        res[name] = {"max_abs_diff": float(np.abs(p_ref - p_oth).max()),
                     "argmax_agreement": float((p_ref.argmax(1) == p_oth.argmax(1)).mean())}
    return res


def quantize_int8(src: str | Path, dst: str | Path) -> None:
    """Dynamic INT8 quantization of MatMul/Gemm weights (activations stay float).

    Conv layers are left in fp32: dynamic ConvInteger kernels are slow on most
    CPUs and the stem is a small share of the compute.
    """
    from onnxruntime.quantization import QuantType, quantize_dynamic
    from onnxruntime.quantization.shape_inference import quant_pre_process

    pre = Path(str(dst) + ".pre.onnx")
    try:
        quant_pre_process(str(src), str(pre), skip_symbolic_shape=True)
        source = pre
    except Exception:
        source = Path(src)
    quantize_dynamic(str(source), str(dst), weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gemm"])
    if pre.exists():
        pre.unlink()
    import onnx
    m = onnx.load(str(dst))
    src_meta = {p.key: p.value for p in onnx.load(str(src)).metadata_props}
    del m.metadata_props[:]
    for k, v in src_meta.items():
        e = m.metadata_props.add()
        e.key, e.value = k, v
    onnx.save(m, str(dst))