"""
Export finegrain/finegrain-box-segmenter (MVANet, Refiners) to ONNX.

The box prompt is NOT part of the network (it is a host-side crop in
refiners.solutions.BoxSegmenter.run). So the exported graph is purely:
    input  float32 [1,3,1024,1024]  (RGB, ImageNet-normalized)
    logits float32 [1,1,1024,1024]  (RAW logits; sigmoid is applied host-side)

See ../.research/*.md for the full sourced analysis.

Usage:
    python export_onnx.py                       # dynamo exporter, opset 18
    python export_onnx.py --patch-pool          # + replace adaptive_avg_pool2d with avg_pool2d
    python export_onnx.py --exporter legacy      # TorchScript exporter, opset 17
"""
import argparse
import json
import os
import time
import traceback
from pathlib import Path

import torch


def patch_pool() -> None:
    """Replace adaptive_avg_pool2d in refiners MVANet `Pool` with an exact avg_pool2d.

    In Pool.forward, output size = (h//ratio, w//ratio) and `assert h % ratio == 0`,
    so each output cell covers exactly `ratio` input pixels -> avg_pool2d(kernel=ratio,
    stride=ratio) is numerically identical, but exports to a plain ONNX AveragePool.
    """
    from torch.nn.functional import avg_pool2d
    from refiners.foundationals.swin.mvanet import utils as U

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, _, h, w = x.shape
        assert h % self.ratio == 0 and w % self.ratio == 0
        r = avg_pool2d(x, kernel_size=self.ratio, stride=self.ratio)
        return torch.unflatten(r, 0, (b, -1))

    U.Pool.forward = forward
    print("[patch] Pool.forward -> F.avg_pool2d(kernel=stride=ratio)")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--exporter", choices=["dynamo", "legacy"], default="dynamo")
    ap.add_argument("--opset", type=int, default=None)
    ap.add_argument("--patch-pool", action="store_true")
    ap.add_argument("--no-fold", action="store_true",
                    help="legacy: do_constant_folding=False (smaller file; ORT folds at load)")
    ap.add_argument("--out", default="models/mvanet_box_segmenter.onnx")
    args = ap.parse_args()
    opset = args.opset if args.opset is not None else (18 if args.exporter == "dynamo" else 17)

    torch.manual_seed(0)
    torch.set_num_threads(max(1, os.cpu_count() or 4))

    if args.patch_pool:
        patch_pool()

    from refiners.solutions import BoxSegmenter

    print("[load] BoxSegmenter(device=cpu) (downloads model.safetensors v0.1 on first run) ...")
    t0 = time.time()
    seg = BoxSegmenter(device="cpu")
    model = seg.model.eval().float()
    n_params = sum(p.numel() for p in model.parameters())
    dtypes = sorted({str(p.dtype) for p in model.parameters()})
    print(f"[load] ok in {time.time()-t0:.1f}s  params={n_params:,}  param_dtypes={dtypes}")

    x = torch.randn(1, 3, 1024, 1024, dtype=torch.float32)
    print("[sanity] running an eager forward ...")
    t0 = time.time()
    with torch.no_grad():
        y = model(x)
    print(f"[sanity] out={tuple(y.shape)} dtype={y.dtype} "
          f"min={float(y.min()):.4f} max={float(y.max()):.4f}  ({time.time()-t0:.1f}s)")
    assert tuple(y.shape) == (1, 1, 1024, 1024), f"unexpected output shape {tuple(y.shape)}"

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    print(f"[export] exporter={args.exporter} opset={opset} -> {out}")
    t0 = time.time()
    try:
        if args.exporter == "dynamo":
            torch.onnx.export(
                model, (x,), str(out),
                dynamo=True, opset_version=opset,
                input_names=["input"], output_names=["logits"],
                external_data=False, optimize=True, verify=False,
            )
        else:
            with torch.no_grad():
                torch.onnx.export(
                    model, (x,), str(out),
                    dynamo=False, opset_version=opset,
                    input_names=["input"], output_names=["logits"],
                    do_constant_folding=not args.no_fold,
                )
    except Exception:
        print("[export] FAILED:\n" + traceback.format_exc())
        raise SystemExit(2)
    print(f"[export] done in {time.time()-t0:.1f}s  size={out.stat().st_size/1e6:.1f} MB")

    # Inspect the produced model
    import onnx
    m = onnx.load(str(out))
    try:
        onnx.checker.check_model(m)
        print("[check] onnx.checker: OK")
    except Exception as e:
        print(f"[check] onnx.checker WARN: {e}")
    ops = sorted({n.op_type for n in m.graph.node})
    opsets = {f"{i.domain or 'ai.onnx'}": i.version for i in m.opset_import}
    ins = [(i.name, [d.dim_value or d.dim_param for d in i.type.tensor_type.shape.dim]) for i in m.graph.input]
    outs = [(o.name, [d.dim_value or d.dim_param for d in o.type.tensor_type.shape.dim]) for o in m.graph.output]
    meta = {
        "exporter": args.exporter, "opset": opset, "patch_pool": args.patch_pool,
        "params": n_params, "size_mb": round(out.stat().st_size / 1e6, 1),
        "opset_import": opsets, "n_nodes": len(m.graph.node),
        "op_types": ops, "inputs": ins, "outputs": outs,
    }
    (out.parent / (out.stem + ".meta.json")).write_text(json.dumps(meta, indent=2))
    print("[inspect] inputs:", ins)
    print("[inspect] outputs:", outs)
    print("[inspect] n_nodes:", len(m.graph.node), "opset_import:", opsets)
    print("[inspect] op_types:", ops)
    print("[done] wrote", out, "and", out.parent / (out.stem + ".meta.json"))


if __name__ == "__main__":
    main()