"""
Produce a compact fp16 variant of the exported ONNX (and best-effort onnxslim).

The 804 MB fp32 export carries ~425 MB of constant-folded fp32 tensors (Swin attention masks /
positional-embedding tables). fp16 (keep_io_types=True so I/O stay float32) ~halves the file and is
native/fast on WebGPU & DirectML. CPU runs fp16 too (via casts), just slower than fp32.

Note: onnxslim's shape-inference pass trips protobuf's 2 GB serialization limit on the fp32 model,
so we convert to fp16 first (smaller) and only then try to slim.
"""
import time
from pathlib import Path

import onnx

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "models" / "mvanet_box_segmenter.onnx"
FP16 = ROOT / "models" / "mvanet_box_segmenter_fp16.onnx"
FP16_SLIM = ROOT / "models" / "mvanet_box_segmenter_fp16_slim.onnx"


def mb(p: Path) -> float:
    return p.stat().st_size / 1e6


def main() -> None:
    print(f"[src ] {SRC.name}  {mb(SRC):.1f} MB")

    # fp16 (input/output kept float32 so the JS harness is unchanged)
    import numpy as np
    from onnx import TensorProto, numpy_helper
    from onnxconverter_common import float16

    def fix_resize_scale_inputs(model) -> int:
        """Resize/Upsample roi & scales inputs must be float32 even when X is fp16 (ONNX spec).
        The fp16 converter wrongly casts the Constants feeding them; cast those back to fp32."""
        inits = {t.name: t for t in model.graph.initializer}
        consts = {n.output[0]: n for n in model.graph.node if n.op_type == "Constant"}
        targets = set()
        for n in model.graph.node:
            if n.op_type in ("Resize", "Upsample"):
                for idx in (1, 2):  # roi, scales
                    if idx < len(n.input) and n.input[idx]:
                        targets.add(n.input[idx])
        fixed = 0
        for name in targets:
            if name in inits and inits[name].data_type == TensorProto.FLOAT16:
                arr = numpy_helper.to_array(inits[name]).astype(np.float32)
                inits[name].CopyFrom(numpy_helper.from_array(arr, name)); fixed += 1
            elif name in consts:
                for a in consts[name].attribute:
                    if a.name == "value" and a.t.data_type == TensorProto.FLOAT16:
                        arr = numpy_helper.to_array(a.t).astype(np.float32)
                        a.t.CopyFrom(numpy_helper.from_array(arr, a.t.name)); fixed += 1
        return fixed

    from onnx import helper

    def wrap_fp32_io(model) -> None:
        """Make external input/output float32 around an all-fp16 interior, via explicit Cast nodes.
        (keep_io_types=True left the fp32 input feeding dtype-agnostic reshape branches, which then
        clashed with cast fp16 branches at a Concat — so we go uniform fp16 + manual I/O casts.)"""
        g = model.graph
        # INPUT: rename the fp16 graph input value, add a fp32 input that Casts into it.
        gi = g.input[0]
        old = gi.name
        internal = old + "_f16"
        for n in g.node:
            n.input[:] = [internal if x == old else x for x in n.input]
        shape = [d.dim_value for d in gi.type.tensor_type.shape.dim]
        g.input.remove(gi)
        g.input.insert(0, helper.make_tensor_value_info(old, TensorProto.FLOAT, shape))
        g.node.insert(0, helper.make_node("Cast", [old], [internal], to=TensorProto.FLOAT16,
                                           name="cast_input_to_fp16"))
        # OUTPUT: rename the fp16 graph output, Cast it back to fp32.
        go = g.output[0]
        oold = go.name
        ointernal = oold + "_f16"
        for n in g.node:
            n.output[:] = [ointernal if x == oold else x for x in n.output]
            n.input[:] = [ointernal if x == oold else x for x in n.input]
        oshape = [d.dim_value for d in go.type.tensor_type.shape.dim]
        g.output.remove(go)
        g.output.insert(0, helper.make_tensor_value_info(oold, TensorProto.FLOAT, oshape))
        g.node.append(helper.make_node("Cast", [ointernal], [oold], to=TensorProto.FLOAT,
                                        name="cast_output_to_fp32"))

    # Shape inference must be ON: otherwise the converter can't see branch types (e.g. the global-view
    # Resize branch) and leaves a mixed fp16/fp32 Concat in SplitMultiView. keep_io_types=True keeps
    # the model's input/output float32 so the Node harness is unchanged.
    # onnxconverter_common leaves several mixed fp16/fp32 boundaries (Concat in SplitMultiView, the
    # WindowSDPA scale Div, ...) that onnx.checker accepts but ORT rejects. General fix: infer types,
    # then Cast any fp32 input of a float-compute op to fp16 so every such op is uniformly fp16.
    FLOAT_OPS = {"Add", "Sub", "Mul", "Div", "Concat", "Where", "Sum", "Min", "Max", "Pow",
                 "MatMul", "Gemm", "PRelu", "Sqrt", "Reciprocal", "Neg", "Softmax", "Erf", "Sigmoid"}

    def reconcile_float_types(model) -> int:
        """Forward-propagate per-tensor dtypes ourselves (onnx shape-inference under-types this graph),
        then Cast any fp32 input of a float-compute op to fp16 so ORT's strict type check passes."""
        from onnx import TensorProto, helper
        FP16, FP32, INT64, BOOL = (TensorProto.FLOAT16, TensorProto.FLOAT,
                                   TensorProto.INT64, TensorProto.BOOL)
        INT_OUT = {"Shape", "Size", "NonZero", "ArgMax", "ArgMin"}
        BOOL_OUT = {"Equal", "Greater", "Less", "GreaterOrEqual", "LessOrEqual",
                    "And", "Or", "Not", "Xor", "IsNaN", "IsInf"}
        dt = {t.name: t.data_type for t in model.graph.initializer}
        for v in model.graph.input:
            dt[v.name] = v.type.tensor_type.elem_type

        def out_dtype(n):
            op = n.op_type
            if op == "Constant":
                for a in n.attribute:
                    if a.name == "value":
                        return a.t.data_type
                    if a.name in ("value_float", "value_floats"):
                        return FP32
                    if a.name in ("value_int", "value_ints"):
                        return INT64
                return None
            if op == "ConstantOfShape":
                for a in n.attribute:
                    if a.name == "value":
                        return a.t.data_type
                return FP32
            if op in ("Cast", "CastLike"):
                return next((a.i for a in n.attribute if a.name == "to"), None)
            if op in INT_OUT:
                return INT64
            if op in BOOL_OUT:
                return BOOL
            return dt.get(n.input[0]) if n.input else None

        for _ in range(3):  # a few passes in case nodes aren't perfectly topo-sorted
            for n in model.graph.node:
                d = out_dtype(n)
                if d is not None:
                    for o in n.output:
                        if o:
                            dt[o] = d

        out_nodes, c = [], 0
        for n in model.graph.node:
            if n.op_type in FLOAT_OPS:
                for idx, i in enumerate(list(n.input)):
                    if i and dt.get(i) == FP32:
                        co = f"{i}__f16_{c}"
                        out_nodes.append(helper.make_node("Cast", [i], [co], to=FP16, name=f"recon_{c}"))
                        n.input[idx] = co
                        dt[co] = FP16
                        c += 1
            out_nodes.append(n)
        model.graph.ClearField("node")
        model.graph.node.extend(out_nodes)
        return c

    # Pipeline order matters:
    #   convert fp16  ->  fix Resize scales  ->  onnxslim (folds 3377 Constants, dedupes)
    #   ->  reconcile float types  ->  clear stale value_info  ->  save.
    # reconcile + value_info-clear MUST be the LAST steps: onnxslim eliminates "redundant" casts
    # (undoing reconcile) and the original fp32 value_info conflicts with the fp16 tensors (ORT rejects).
    tmp1 = FP16.with_name("_fp16_pre.onnx")
    m = onnx.load(str(SRC))
    t = time.time()
    m16 = float16.convert_float_to_float16(m, keep_io_types=True, disable_shape_infer=False)
    fix_resize_scale_inputs(m16)
    onnx.save(m16, str(tmp1))

    from onnxslim import slim
    slim(str(tmp1), str(FP16))                       # slim in place into the final path

    ms = onnx.load(str(FP16))
    n_fixed = fix_resize_scale_inputs(ms)            # re-assert in case slim refolded
    n_recon = reconcile_float_types(ms)
    ms.graph.ClearField("value_info")               # drop stale fp32 annotations; ORT re-infers
    onnx.checker.check_model(ms)
    onnx.save(ms, str(FP16))
    tmp1.unlink(missing_ok=True)
    print(f"[fp16] slim+reconcile: resize-fixes={n_fixed} reconcile-casts={n_recon}")
    print(f"[fp16] {FP16.name}  {mb(FP16):.1f} MB  ({time.time()-t:.1f}s)")


if __name__ == "__main__":
    main()