"""Export a package folder (config.json, backbone/, heads.safetensors) to ONNX for the in-browser demo.

  python onnx_export.py DIR            # writes DIR/onnx/model_fp32.onnx and DIR/onnx/model_int8.onnx

The graph takes input_ids and attention_mask (int64, dynamic batch and sequence length) and returns `probs`
[batch, n_heads], the sigmoid of each head's logit, in the order of config.json "heads". The int8 file is dynamic
quantisation of the weights with per-channel scales: the embedding table (19.3 of the 32 million parameters) and the
attention and MLP up-projection matrices. The MLP down-projection stays fp32: quantising it cost up to 0.02 AUROC
on the delegate head, the other layers cost under 0.001 (see PARITY.md).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
from cue import CueModel  # noqa: E402


class ProbWrapper(torch.nn.Module):
    def __init__(self, m: CueModel):
        super().__init__()
        self.backbone = m.backbone
        self.w = torch.nn.Parameter(torch.cat([m.heads[q].weight for q in m.head_names], 0).detach().clone(), requires_grad=False)
        self.b = torch.nn.Parameter(torch.cat([m.heads[q].bias for q in m.head_names], 0).detach().clone(), requires_grad=False)

    def forward(self, input_ids, attention_mask):
        b = self.backbone
        h = b.model(input_ids=input_ids, attention_mask=attention_mask)[0]
        mk = attention_mask.unsqueeze(-1).to(h.dtype)
        h = (h * mk).sum(1) / mk.sum(1)          # mean pooling (config classifier_pooling = "mean")
        v = b.head(h)
        return torch.sigmoid(v @ self.w.T + self.b)


def export(folder: Path, opset: int = 17):
    from transformers import AutoModelForSequenceClassification
    from onnxruntime.quantization import QuantType, quantize_dynamic

    folder = Path(folder)
    cfg = json.loads((folder / "config.json").read_text())
    m = CueModel.from_pretrained(folder)
    if m.backbone.config.classifier_pooling != "mean":
        sys.exit("only mean pooling is exported")
    # eager attention exports cleanly; the weights are identical
    m.backbone = AutoModelForSequenceClassification.from_pretrained(folder / cfg["backbone_dir"], attn_implementation="eager").eval()
    w = ProbWrapper(m).eval()
    out = folder / "onnx"
    out.mkdir(exist_ok=True)
    ids = torch.randint(5, 1000, (2, 24))
    mask = torch.ones_like(ids)
    mask[1, 18:] = 0
    fp32 = out / "model_fp32.onnx"
    torch.onnx.export(w, (ids, mask), str(fp32), input_names=["input_ids", "attention_mask"], output_names=["probs"],
                      dynamic_axes={"input_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"},
                                    "probs": {0: "batch"}}, opset_version=opset, dynamo=False)
    int8 = out / "model_int8.onnx"
    import onnx
    names = [n.name for n in onnx.load(str(fp32)).graph.node if n.op_type == "MatMul"]
    keep = ("/attn/Wqkv/", "/attn/Wo/", "/mlp/Wi/")   # int8; the MLP down-projection, head and rotary stay fp32
    quantize_dynamic(str(fp32), str(int8), weight_type=QuantType.QInt8, per_channel=True,
                     nodes_to_exclude=[n for n in names if not any(k in n for k in keep)])
    for p in (fp32, int8):
        print(f"{p.name}: {p.stat().st_size / 1e6:.1f} MB")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("dir")
    ap.add_argument("--opset", type=int, default=17)
    a = ap.parse_args()
    export(Path(a.dir), a.opset)