"""Exports SAM 3 (facebook/sam3) to ONNX for in-browser inference with ONNX Runtime Web on WebGPU.

Three graphs (see `sam3_graphs.py`): `vision` (once per image), `text` and `detector` (once per
phrase), each at a fixed shape: a 1008 × 1008 image, a 32-token phrase, batch 1.

Files, in `<out>/onnx/`, named by graph and precision:

- `*_fp16.onnx`: fp16 weights and activations, inputs and outputs float32. Standard is the three
  of these.
- `vision_q8f16.onnx`, `text_q8f16.onnx`: as fp16, but every MatMul weight stored at 8 bits
  (`MatMulNBits`, symmetric, block size 128) and expanded on the GPU. Compact is these two and
  `detector_fp16.onnx`, which is small and left at fp16.

Weight-only quantisation keeps every node on WebGPU: the usual dynamic "int8" export runs
`MatMulInteger` and `DynamicQuantizeLinear`, which ONNX Runtime Web's WebGPU backend lacks.

Usage:
    python export_sam3.py --weights <facebook/sam3 checkout> --out <dir>

Requires torch, transformers (with `Sam3Model`), onnx, onnxscript, onnxconverter-common and
onnxruntime. Runs on the CPU. The fp32 graphs go to `<out>/fp32/`, with external data, and
`--skip-fp32` reuses them.
"""

import argparse
from pathlib import Path

import onnx
import torch
from onnxconverter_common import float16
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
from transformers import Sam3Model

from sam3_graphs import IMAGE_SIZE, TEXT_TOKENS, Sam3Detector, Sam3Text, Sam3Vision, position_encoding_72

OPSET = 18
GRAPHS = ("vision", "text", "detector")
QUANTISED = ("vision", "text")


def export_fp32(model, out: Path) -> None:
    out.mkdir(parents=True, exist_ok=True)
    pixels = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE)
    ids = torch.zeros(1, TEXT_TOKENS, dtype=torch.long)
    mask = torch.ones(1, TEXT_TOKENS, dtype=torch.long)
    vision = Sam3Vision(model).eval()
    with torch.no_grad():
        levels = vision(pixels)
    text = Sam3Text(model).eval()
    with torch.no_grad():
        features = text(ids, mask)
    detector = Sam3Detector(model, position_encoding_72(model)).eval()
    jobs = {
        "vision": (vision, (pixels,), ["pixel_values"], ["features_288", "features_144", "features_72"]),
        "text": (text, (ids, mask), ["input_ids", "attention_mask"], ["text_features"]),
        "detector": (
            detector,
            (*levels, features, mask),
            ["features_288", "features_144", "features_72", "text_features", "text_mask"],
            ["pred_logits", "pred_boxes", "presence_logits", "pred_masks", "semantic_seg"],
        ),
    }
    for name, (module, args, inputs, outputs) in jobs.items():
        print(f"exporting {name}")
        with torch.no_grad():
            torch.onnx.export(
                module,
                args,
                str(out / f"{name}.onnx"),
                input_names=inputs,
                output_names=outputs,
                opset_version=OPSET,
                dynamo=True,
                external_data=True,
                optimize=True,
            )


def fix_casts(model: onnx.ModelProto) -> int:
    """Points each Cast to float32 whose output the converter declared float16 at float16.

    `convert_float_to_float16` retypes a Cast's output but leaves its `to`, so a Cast from a
    non-float source (the text mask, cast from bool for pooling) still says float32 and the
    graph doesn't load. Returns how many it fixed.
    """
    declared = {vi.name: vi.type.tensor_type.elem_type for vi in model.graph.value_info}
    fixed = 0
    for node in model.graph.node:
        if node.op_type != "Cast" or declared.get(node.output[0]) != onnx.TensorProto.FLOAT16:
            continue
        for attribute in node.attribute:
            if attribute.name == "to" and attribute.i == onnx.TensorProto.FLOAT:
                attribute.i = onnx.TensorProto.FLOAT16
                fixed += 1
    return fixed


def to_fp16(src: Path, dst: Path) -> None:
    dst.mkdir(parents=True, exist_ok=True)
    for name in GRAPHS:
        model = onnx.load(str(src / f"{name}.onnx"))
        # Stale annotations from the export would fight the converter's own inference.
        del model.graph.value_info[:]
        half = float16.convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=False)
        print(f"fp16 {name}: {fix_casts(half)} casts fixed")
        onnx.save(half, str(dst / f"{name}_fp16.onnx"))


def to_q8f16(onnx_dir: Path, bits: int, block_size: int) -> None:
    for name in QUANTISED:
        print(f"{bits}-bit {name}")
        quantiser = MatMulNBitsQuantizer(
            onnx.load(str(onnx_dir / f"{name}_fp16.onnx")),
            bits=bits,
            block_size=block_size,
            is_symmetric=True,
        )
        quantiser.process()
        onnx.save(quantiser.model.model, str(onnx_dir / f"{name}_q{bits}f16.onnx"))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--weights", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--bits", type=int, default=8)
    parser.add_argument("--block-size", type=int, default=128)
    parser.add_argument("--skip-fp32", action="store_true", help="reuse an earlier fp32 export")
    args = parser.parse_args()

    if not args.skip_fp32:
        model = Sam3Model.from_pretrained(args.weights, attn_implementation="eager").eval()
        export_fp32(model, args.out / "fp32")
        del model
    to_fp16(args.out / "fp32", args.out / "onnx")
    to_q8f16(args.out / "onnx", args.bits, args.block_size)


if __name__ == "__main__":
    main()