# Copyright 2026 Sam McLeod
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Half-precision (FP16 / BF16) converter for the Granite Speech 4.1 ONNX
exports.

Wraps `onnxruntime.transformers.float16.convert_float_to_float16` for FP16
and a small graph walker for BF16 (the upstream library only ships FP16).
The output graphs match the conventions used elsewhere in this project:

  - Single external-data sidecar per graph.
  - `ai.onnx` opset 20 / IR 10 only - no `com.microsoft` ops. Verified by
    a domain audit after the conversion completes.
  - `keep_io_types=True` so the model's external inputs/outputs stay FP32.
    Consumers can keep feeding the same dtype as for the FP32 graph; the
    Cast nodes inserted by the converter handle the boundary.

Numeric-sensitive ops are kept at FP32 by default (LayerNormalization,
Softmax, Pow, Sqrt, Reduce*). The converter inserts Cast nodes around
them automatically so the FP32 islands compose with the FP16 body.

Usage:
    python convert_fp.py --input PATH --output PATH --precision fp16
    python convert_fp.py --input PATH --output PATH --precision bf16

The script is self-contained (no project-internal imports) so it ships
inside each Hugging Face bundle alongside the export script.
"""

from __future__ import annotations

import argparse
import sys
import tempfile
import time
from pathlib import Path

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper


# Ops we always leave at FP32. Mostly numeric-sensitive primitives (norms /
# softmax / RMSNorm decomposition) where FP16 overflow / underflow degrades
# accuracy noticeably without any size win on small ops.
DEFAULT_FP32_OPS = (
    "LayerNormalization",
    "RMSNormalization",
    "Softmax",
    "LogSoftmax",
    "Pow",
    "Sqrt",
    "ReduceMean",
    "ReduceSum",
    "ReduceL2",
)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Convert a Granite Speech 4.1 ONNX graph to FP16 or BF16.",
    )
    p.add_argument(
        "--input",
        required=True,
        type=Path,
        help="Path to the FP32 .onnx graph (external sidecar must sit alongside it).",
    )
    p.add_argument(
        "--output",
        required=True,
        type=Path,
        help="Destination .onnx path. A single sidecar named <output>_data is written next to it.",
    )
    p.add_argument(
        "--precision",
        choices=("fp16", "bf16", "fp16w"),
        default="fp16",
        help=(
            "Target precision (default: fp16). 'fp16w' is weights-FP16 / "
            "FP32-compute hybrid: FP32 initializers are rewritten to FP16 in "
            "storage, but a Cast(FP16->FP32) is inserted before each consumer "
            "so arithmetic and IO stay at FP32. Halves weight storage with no "
            "FP16-overflow risk in the body."
        ),
    )
    p.add_argument(
        "--fp16w-skip-small",
        type=int,
        default=1024,
        help="fp16w only: skip initializers with fewer elements than this "
        "(saves negligible bytes and avoids touching norm scale/bias). "
        "Default: 1024.",
    )
    p.add_argument(
        "--fp16w-skip-norm-max-elems",
        type=int,
        default=4096,
        help="fp16w only: skip 1-D initializers <=this size feeding a Mul/Add "
        "whose sibling traces from Sqrt/Div/Pow (LayerNorm/RMSNorm scale/bias "
        "patterns). Default: 4096.",
    )
    p.add_argument(
        "--keep-io-types",
        action="store_true",
        default=True,
        help="Keep the model's external IO at FP32 even though the body runs at "
        "the target precision. Consumers don't need to change call sites. "
        "Default: on.",
    )
    p.add_argument(
        "--no-keep-io-types",
        dest="keep_io_types",
        action="store_false",
        help="Convert external IO to the target precision as well.",
    )
    p.add_argument(
        "--op-block-list",
        default=",".join(DEFAULT_FP32_OPS),
        help="Comma-separated op types kept at FP32. Defaults to numeric-sensitive "
        "ops (LayerNormalization, Softmax, Pow, Sqrt, ReduceMean, ...).",
    )
    p.add_argument(
        "--node-block-list",
        default="",
        help="Comma-separated explicit node names to leave at FP32 (overrides "
        "op-block-list inclusion). Useful for skipping a specific MatMul that "
        "drifts under FP16.",
    )
    p.add_argument(
        "--node-block-pattern",
        default=r"(?i)(layernorm|rmsnorm|/norm/)",
        help="Regex applied to node names; matching nodes are added to the FP32 "
        "block list. Default catches the Granite LLM's RMSNorm subgraphs "
        "(Pow/ReduceMean/Sqrt/Mul/Cast/...) which torch.onnx decomposes into "
        "primitives. Without this, ORT's SimplifiedLayerNormFusion pass at "
        "load time gets confused by the mixed-precision Cast nodes inside "
        "the RMSNorm pattern and refuses to load the graph.",
    )
    p.add_argument(
        "--ir-version",
        type=int,
        default=10,
        help="ONNX IR version to write (default: 10, matches the FP32 exports).",
    )
    return p.parse_args(argv)


# ---------------------------------------------------------------------------
# FP16 path: defer to onnxruntime.transformers.float16
# ---------------------------------------------------------------------------


def _load_fp32_proto(input_path: Path) -> onnx.ModelProto:
    """Load an FP32 graph + its external sidecar.

    Works around onnx's hardlink-attack check: if the source files have
    additional hardlinks (which they do when staged into a bundle directory),
    onnx.load raises ValidationError. We side-step that by reading the bytes
    of the sidecar manually and storing them on the model proto in-memory.
    """
    proto = onnx.load(str(input_path), load_external_data=False)
    sidecar_paths: dict[str, bytes] = {}
    for tensor in proto.graph.initializer:
        if tensor.data_location != onnx.TensorProto.EXTERNAL:
            continue
        ext = {kv.key: kv.value for kv in tensor.external_data}
        location = ext.get("location")
        if not location:
            continue
        side_path = input_path.parent / location
        if side_path.name not in sidecar_paths:
            sidecar_paths[side_path.name] = side_path.read_bytes()
        offset = int(ext.get("offset", 0))
        length = int(ext.get("length", len(sidecar_paths[side_path.name]) - offset))
        tensor.raw_data = sidecar_paths[side_path.name][offset:offset + length]
        tensor.ClearField("data_location")
        tensor.ClearField("external_data")
    return proto


def _ensure_value_info(input_path: Path) -> Path:
    """Run shape inference if the graph has no value_info entries.

    `convert_float_to_float16` keys boundary Cast insertion off
    `graph.value_info` entries. Our torch.onnx exports do not include
    value_info for intermediate tensors (count is 0). Without value_info,
    the helper can't see the boundary between blocked nodes (FP32) and
    surrounding nodes (FP16), so it skips inserting the bridging Casts and
    the resulting graph fails to load with a Type Error on the first FP16
    consumer of an FP32-blocked output.

    `onnx.shape_inference.infer_shapes_path` reads/writes via disk so it
    bypasses protobuf's 2GB single-message limit, which we'd otherwise hit
    on the ~6 GB LLM body graphs.
    """
    proto = onnx.load(str(input_path), load_external_data=False)
    if len(proto.graph.value_info) > 0:
        return input_path

    inferred = input_path.with_name(input_path.stem + ".inferred.onnx")
    if inferred.exists():
        inferred.unlink()
    print(f"  no value_info present; running shape inference -> {inferred.name}")
    t0 = time.time()
    onnx.shape_inference.infer_shapes_path(
        str(input_path),
        str(inferred),
        check_type=False,
        strict_mode=False,
    )
    print(f"  shape inference done in {time.time() - t0:.1f}s")
    return inferred


def convert_to_fp16(
    input_path: Path,
    op_block_list: list[str],
    node_block_list: list[str],
    keep_io_types: bool,
) -> onnx.ModelProto:
    """Use onnxruntime's float16 helper. Returns the converted ModelProto.

    The helper rewrites every FP32 initializer to FP16 in-place and inserts
    Cast nodes around any node whose op_type is in `op_block_list`. With
    `keep_io_types=True`, the graph's external inputs/outputs stay FP32 and
    Cast nodes bridge to the FP16 body.
    """
    from onnxruntime.transformers.float16 import convert_float_to_float16

    inferred_path = _ensure_value_info(input_path)
    print("  loading FP32 graph (with external data)")
    proto = _load_fp32_proto(inferred_path)
    print(f"  loaded graph: value_info={len(proto.graph.value_info)} entries")

    print(
        f"  running convert_float_to_float16 "
        f"(keep_io_types={keep_io_types}, op_block_list={op_block_list}, "
        f"node_block_list_size={len(node_block_list)})"
    )
    t0 = time.time()
    converted = convert_float_to_float16(
        proto,
        keep_io_types=keep_io_types,
        op_block_list=op_block_list or None,
        node_block_list=node_block_list or None,
        # Shape inference re-serialises the proto, which trips the 2GB
        # protobuf single-message limit on the LLM-body graphs (prompt_encode
        # and decode_step have ~6 GB of weights each). _ensure_value_info
        # already populated value_info via the disk-based path.
        disable_shape_infer=True,
    )
    print(f"  fp16 conversion done in {time.time() - t0:.1f}s")

    # Best-effort cleanup of the temp inferred proto (sidecar belongs to the
    # original FP32 graph and must not be touched).
    if inferred_path != input_path and inferred_path.exists():
        try:
            inferred_path.unlink()
        except OSError:
            pass
    return converted


# ---------------------------------------------------------------------------
# FP16w path: weights-FP16 with FP32 compute (storage-only halving)
# ---------------------------------------------------------------------------


def _consumer_input_traces_norm(
    proto_graph: onnx.GraphProto, consumers: list[onnx.NodeProto]
) -> bool:
    """Heuristic: does any consumer have a *sibling* input traced from a
    Sqrt / Div / Pow / ReduceMean op? This catches RMSNorm scale (`Mul` with
    one input from Sqrt's reciprocal) and LayerNorm bias (`Add` after
    LayerNorm body). Conservative: false negatives are fine (we'll just
    convert a few more tensors), false positives keep too many at FP32."""
    norm_outputs: set[str] = set()
    for n in proto_graph.node:
        if n.op_type in ("Sqrt", "Div", "Pow", "ReduceMean"):
            for o in n.output:
                if o:
                    norm_outputs.add(o)
    for c in consumers:
        for inp in c.input:
            if inp in norm_outputs:
                return True
    return False


def convert_to_weights_fp16(
    input_path: Path,
    skip_small_elems: int,
    skip_norm_max_elems: int,
) -> onnx.ModelProto:
    """Rewrite FP32 initializers to FP16 storage and insert Cast(FP16->FP32)
    before each consumer so arithmetic stays FP32.

    Halves the per-tensor weight footprint with no FP16-overflow risk in the
    network body. The trade-off is one extra Cast per converted weight; ORT
    will sometimes fuse those into the downstream MatMul on x86/arm64.

    Skipped:
      - initializers smaller than `skip_small_elems` (saves negligible bytes
        and avoids norm scale/bias).
      - rank-1 tensors with <= `skip_norm_max_elems` whose consumer pattern
        looks like LayerNorm/RMSNorm scale or bias (heuristic).

    Returns the rewritten ModelProto. Caller saves with single sidecar.
    """
    inferred_path = _ensure_value_info(input_path)
    print("  loading FP32 graph (with external data)")
    proto = _load_fp32_proto(inferred_path)
    if inferred_path != input_path and inferred_path.exists():
        try:
            inferred_path.unlink()
        except OSError:
            pass

    init_consumers: dict[str, list[onnx.NodeProto]] = {}
    init_name_set = {init.name for init in proto.graph.initializer}
    for node in proto.graph.node:
        for inp in node.input:
            if inp in init_name_set:
                init_consumers.setdefault(inp, []).append(node)

    converted_inits: dict[str, str] = {}
    new_cast_nodes: list[onnx.NodeProto] = []
    skipped = {"non_float": 0, "small": 0, "norm_like": 0}
    bytes_saved = 0
    cast_uid = 0

    for init in proto.graph.initializer:
        if init.data_type != TensorProto.FLOAT:
            skipped["non_float"] += 1
            continue

        arr = numpy_helper.to_array(init)
        n_elem = int(arr.size)
        if n_elem < skip_small_elems:
            skipped["small"] += 1
            continue

        consumers = init_consumers.get(init.name, [])
        if (
            arr.ndim == 1
            and n_elem <= skip_norm_max_elems
            and consumers
            and any(c.op_type in ("Mul", "Add") for c in consumers)
            and _consumer_input_traces_norm(proto.graph, consumers)
        ):
            skipped["norm_like"] += 1
            continue

        fp16_arr = arr.astype(np.float16, copy=False)
        bytes_saved += arr.nbytes - fp16_arr.nbytes

        init.ClearField("float_data")
        init.ClearField("raw_data")
        init.raw_data = fp16_arr.tobytes()
        init.data_type = TensorProto.FLOAT16
        init.ClearField("data_location")
        init.ClearField("external_data")

        cast_uid += 1
        cast_out = f"{init.name}_fp16w_cast_to_fp32"
        cast_node = helper.make_node(
            "Cast",
            inputs=[init.name],
            outputs=[cast_out],
            name=f"Cast_fp16w_{cast_uid}_{init.name.replace('/', '_')}",
            to=TensorProto.FLOAT,
        )
        new_cast_nodes.append(cast_node)
        converted_inits[init.name] = cast_out

    for node in proto.graph.node:
        new_inputs = list(node.input)
        rewired = False
        for i, name in enumerate(new_inputs):
            if name in converted_inits:
                new_inputs[i] = converted_inits[name]
                rewired = True
        if rewired:
            del node.input[:]
            node.input.extend(new_inputs)

    proto.graph.node.extend(new_cast_nodes)
    proto = _topo_sort(proto)

    print(
        f"  fp16w conversion done: {len(converted_inits)} initializers "
        f"converted ({bytes_saved / 1e9:.2f} GB saved on weights), "
        f"{len(new_cast_nodes)} Cast nodes inserted, skipped={skipped}"
    )
    return proto


# ---------------------------------------------------------------------------
# BF16 path: graph walker (onnxruntime.transformers.float16 doesn't do BF16)
# ---------------------------------------------------------------------------


def _bf16_pack_from_fp32(arr: np.ndarray) -> bytes:
    """Pack an FP32 numpy array into raw BFloat16 bytes (truncating mantissa).

    BFloat16 is the high 16 bits of an FP32 IEEE-754 representation. Round-to-
    nearest-even is overkill for our purposes; truncation is what every
    standard kernel does for inference paths.
    """
    if arr.dtype != np.float32:
        arr = arr.astype(np.float32, copy=False)
    u32 = arr.view(np.uint32)
    # Round-to-nearest-even via add of 0x7FFF + (lsb of high half). This
    # matches torch's bfloat16 cast behaviour better than naked truncation
    # and avoids a small bias on values with all-ones mantissas.
    rounding_bias = 0x7FFF + ((u32 >> 16) & 1)
    rounded = (u32 + rounding_bias) >> 16
    bf16 = rounded.astype(np.uint16)
    return bf16.tobytes()


def _convert_initializer_to_bf16(tensor: onnx.TensorProto) -> None:
    """Mutate a FP32 TensorProto in-place to BFLOAT16 layout."""
    arr = numpy_helper.to_array(tensor)
    if arr.dtype != np.float32:
        return
    raw = _bf16_pack_from_fp32(arr)
    tensor.ClearField("float_data")
    tensor.ClearField("raw_data")
    tensor.raw_data = raw
    tensor.data_type = TensorProto.BFLOAT16
    tensor.ClearField("data_location")
    tensor.ClearField("external_data")


def _bf16_value_info(name: str, shape: list[int | str | None]) -> onnx.ValueInfoProto:
    return helper.make_tensor_value_info(name, TensorProto.BFLOAT16, shape)


def _make_cast(input_name: str, output_name: str, to_dtype: int, name: str) -> onnx.NodeProto:
    return helper.make_node(
        "Cast",
        inputs=[input_name],
        outputs=[output_name],
        name=name,
        to=to_dtype,
    )


def convert_to_bf16(
    input_path: Path,
    op_block_list: list[str],
    node_block_list: list[str],
    keep_io_types: bool,
) -> onnx.ModelProto:
    """Walk the graph and convert every FP32 initializer to BFLOAT16, casting
    around blocked ops so they receive FP32 inputs and emit FP32 outputs.

    Mirrors the high-level idea of `convert_float_to_float16` but emits
    BFLOAT16 (`TensorProto.BFLOAT16`) instead of FLOAT16. Op coverage is
    smaller than FP16's; if the consumer's runtime lacks a BF16 kernel for
    some op the runtime will error at load time. Validate the resulting
    graph before shipping.
    """
    inferred_path = _ensure_value_info(input_path)
    print("  loading FP32 graph (with external data)")
    proto = _load_fp32_proto(inferred_path)
    if inferred_path != input_path and inferred_path.exists():
        try:
            inferred_path.unlink()
        except OSError:
            pass

    block_op_set = set(op_block_list)
    block_node_set = set(node_block_list)

    # 1) Rewrite every FP32 initializer to BFLOAT16 except those produced
    #    only into blocked ops.
    blocked_init_names: set[str] = set()
    init_consumers: dict[str, list[onnx.NodeProto]] = {}
    init_names = {init.name for init in proto.graph.initializer}
    for node in proto.graph.node:
        for inp in node.input:
            if inp in init_names:
                init_consumers.setdefault(inp, []).append(node)

    for init in proto.graph.initializer:
        consumers = init_consumers.get(init.name, [])
        if consumers and all(
            (n.op_type in block_op_set or n.name in block_node_set) for n in consumers
        ):
            blocked_init_names.add(init.name)

    converted_initializers = 0
    for init in proto.graph.initializer:
        if init.name in blocked_init_names:
            continue
        if init.data_type != TensorProto.FLOAT:
            continue
        _convert_initializer_to_bf16(init)
        converted_initializers += 1

    # 2) Insert Cast(FP32) -> blocked op -> Cast(BF16) bridges. This means
    #    blocked ops see FP32 inputs and produce FP32 outputs, and the
    #    surrounding BF16 body uses Cast to bridge.
    cast_uid = 0
    new_nodes: list[onnx.NodeProto] = []
    # Collect tensor names that are BF16 in the graph after step 1. We
    # approximate by tracking outputs of non-blocked nodes plus converted
    # initializers.
    bf16_tensors: set[str] = set()
    for init in proto.graph.initializer:
        if init.data_type == TensorProto.BFLOAT16:
            bf16_tensors.add(init.name)

    # Track FP32 inputs (graph inputs) whose dtype was FLOAT. They stay FP32
    # if keep_io_types=True, else they get cast to BF16 at the entry.
    graph_input_names = {vi.name for vi in proto.graph.input}
    graph_output_names = {vi.name for vi in proto.graph.output}

    if not keep_io_types:
        # Promote graph inputs to BF16 (if they were FLOAT).
        for vi in proto.graph.input:
            if vi.type.tensor_type.elem_type == TensorProto.FLOAT:
                vi.type.tensor_type.elem_type = TensorProto.BFLOAT16
                bf16_tensors.add(vi.name)

    # Walk nodes in order. For blocked nodes, cast inputs to FP32 and outputs
    # back to BF16. For non-blocked nodes, leave them alone (they pick up
    # BF16 from the converted initializers and casts).
    rewritten_node_count = 0
    for node in list(proto.graph.node):
        is_blocked = node.op_type in block_op_set or node.name in block_node_set
        if is_blocked:
            new_input_names: list[str] = []
            for inp in node.input:
                if not inp:
                    new_input_names.append(inp)
                    continue
                # Only need a Cast if the producing tensor is currently BF16.
                if inp in bf16_tensors:
                    cast_out = f"{inp}_cast_to_fp32_{cast_uid}"
                    cast_uid += 1
                    new_nodes.append(
                        _make_cast(
                            inp,
                            cast_out,
                            TensorProto.FLOAT,
                            f"Cast_pre_{node.name or node.op_type}_{cast_uid}",
                        )
                    )
                    new_input_names.append(cast_out)
                else:
                    new_input_names.append(inp)
            del node.input[:]
            node.input.extend(new_input_names)
            # Outputs of the blocked node are FP32; if they feed into the
            # graph's external outputs we leave them; otherwise downstream
            # nodes expect BF16 unless they're also blocked. Insert Cast.
            for out in list(node.output):
                if not out:
                    continue
                if out in graph_output_names and keep_io_types:
                    # External output is FP32, no cast needed.
                    continue
                # Replace `out` in subsequent nodes with a casted variant.
                cast_in = out
                cast_out = f"{out}_cast_to_bf16_{cast_uid}"
                cast_uid += 1
                new_nodes.append(
                    _make_cast(
                        cast_in,
                        cast_out,
                        TensorProto.BFLOAT16,
                        f"Cast_post_{node.name or node.op_type}_{cast_uid}",
                    )
                )
                # Rewire all later nodes that consume `out` to consume `cast_out`.
                for later in proto.graph.node:
                    if later is node:
                        continue
                    new_inputs = list(later.input)
                    rewired = False
                    for i, name in enumerate(new_inputs):
                        if name == out:
                            new_inputs[i] = cast_out
                            rewired = True
                    if rewired:
                        del later.input[:]
                        later.input.extend(new_inputs)
                bf16_tensors.add(cast_out)
            rewritten_node_count += 1
        else:
            for out in node.output:
                if out:
                    bf16_tensors.add(out)

    # 3) If keep_io_types, insert leading Cast(FP32->BF16) for every FP32
    #    graph input that feeds a non-blocked node, and trailing Cast(BF16
    #    ->FP32) for every BF16 tensor that is a graph output.
    if keep_io_types:
        for vi in proto.graph.input:
            if vi.type.tensor_type.elem_type != TensorProto.FLOAT:
                continue
            inp_name = vi.name
            cast_out = f"{inp_name}_cast_to_bf16_{cast_uid}"
            cast_uid += 1
            new_nodes.insert(
                0,
                _make_cast(
                    inp_name,
                    cast_out,
                    TensorProto.BFLOAT16,
                    f"Cast_input_{inp_name}_{cast_uid}",
                ),
            )
            for n in proto.graph.node:
                new_inputs = list(n.input)
                rewired = False
                for i, name in enumerate(new_inputs):
                    if name == inp_name:
                        new_inputs[i] = cast_out
                        rewired = True
                if rewired:
                    del n.input[:]
                    n.input.extend(new_inputs)
            bf16_tensors.add(cast_out)

        for vi in proto.graph.output:
            if vi.type.tensor_type.elem_type != TensorProto.FLOAT:
                continue
            # Find the producing node and rewire via Cast back to FP32.
            out_name = vi.name
            producing_idx = None
            for i, n in enumerate(proto.graph.node):
                if out_name in n.output:
                    producing_idx = i
                    break
            if producing_idx is None:
                continue
            producing = proto.graph.node[producing_idx]
            # If the producing op is blocked, its output is already FP32
            # via step 2's logic - but only if it's not an external output.
            # For external FP32 outputs we want the producer's output to
            # stay FP32. The simplest unified approach is: rename producer's
            # output to `<name>_pre_cast` and cast that -> `<name>` (FP32).
            internal_name = f"{out_name}_pre_cast_{cast_uid}"
            cast_uid += 1
            for i, o in enumerate(list(producing.output)):
                if o == out_name:
                    producing.output[i] = internal_name
            new_nodes.append(
                _make_cast(
                    internal_name,
                    out_name,
                    TensorProto.FLOAT,
                    f"Cast_output_{out_name}_{cast_uid}",
                )
            )
            # Producer's output dtype is BF16 unless it's a blocked op
            # (already FP32). The Cast handles the rest.
            bf16_tensors.add(internal_name)

    # 4) Splice the new Cast nodes back into the graph in topological order.
    # ONNX requires nodes to be sorted topologically; the simplest correct
    # approach is to append all new nodes and then run a topological sort.
    proto.graph.node.extend(new_nodes)
    proto = _topo_sort(proto)

    print(
        f"  bf16 conversion done: {converted_initializers} initializers converted, "
        f"{rewritten_node_count} nodes blocked at fp32, {len(new_nodes)} cast nodes inserted"
    )
    return proto


def _dedup_inserted_casts(model: onnx.ModelProto) -> int:
    """Merge duplicate Cast nodes inserted by `convert_float_to_float16`.

    The helper inserts a Cast(input -> input_cast_to_*) before every blocked
    node, but when two blocked nodes share an input tensor it inserts two
    Cast nodes with identical (input, output, to, name) - an SSA violation
    that fails `onnx.checker.check_model`. Walk the graph and collapse
    duplicate Casts that share `(input_name, to_attr)` into a single node.
    """
    nodes = list(model.graph.node)
    seen: dict[tuple[str, int], onnx.NodeProto] = {}
    keep: list[onnx.NodeProto] = []
    removed = 0
    for n in nodes:
        if n.op_type != "Cast" or len(n.input) != 1 or len(n.output) != 1:
            keep.append(n)
            continue
        # Only target Casts inserted by the helper (their names end with
        # `_cast_to_fp32_node` or similar). Be conservative: match by output
        # suffix introduced by the helper.
        if not (n.output[0].endswith("_cast_to_fp32") or n.output[0].endswith("_cast_to_fp16")):
            keep.append(n)
            continue
        to_attr = next((a.i for a in n.attribute if a.name == "to"), -1)
        key = (n.input[0], to_attr)
        if key in seen:
            removed += 1
            continue
        seen[key] = n
        keep.append(n)

    if removed > 0:
        del model.graph.node[:]
        model.graph.node.extend(keep)
    return removed


def _topo_sort(model: onnx.ModelProto) -> onnx.ModelProto:
    """Stable topological sort of model.graph.node based on input/output edges."""
    nodes = list(model.graph.node)
    initializer_names = {i.name for i in model.graph.initializer}
    graph_input_names = {vi.name for vi in model.graph.input}
    available: set[str] = set(initializer_names) | set(graph_input_names) | {""}

    sorted_nodes: list[onnx.NodeProto] = []
    remaining = list(nodes)
    while remaining:
        progressed = False
        next_remaining: list[onnx.NodeProto] = []
        for n in remaining:
            if all(inp in available for inp in n.input):
                sorted_nodes.append(n)
                for out in n.output:
                    if out:
                        available.add(out)
                progressed = True
            else:
                next_remaining.append(n)
        if not progressed:
            # Cycle or missing producer. Append remaining in order; onnx.checker
            # will surface the issue.
            sorted_nodes.extend(next_remaining)
            break
        remaining = next_remaining

    del model.graph.node[:]
    model.graph.node.extend(sorted_nodes)
    return model


# ---------------------------------------------------------------------------
# Save + verify
# ---------------------------------------------------------------------------


def assert_pure_ai_onnx(model_path: Path) -> list[str]:
    """Reload the produced graph and verify only `ai.onnx` nodes are present."""
    proto = onnx.load(str(model_path), load_external_data=False)
    domains = sorted({(n.domain or "ai.onnx") for n in proto.graph.node})
    forbidden = [d for d in domains if d not in ("ai.onnx", "")]
    if forbidden:
        raise RuntimeError(
            f"Converted graph contains forbidden op domains {forbidden}. "
            "Re-run with a wider --op-block-list."
        )
    return domains


def save_with_single_sidecar(
    proto: onnx.ModelProto,
    final_out: Path,
    ir_version: int,
) -> None:
    """Write the converted proto to `final_out` with a single .onnx_data sidecar."""
    if proto.ir_version < ir_version:
        proto.ir_version = ir_version
    for tensor in proto.graph.initializer:
        tensor.ClearField("data_location")
        tensor.ClearField("external_data")

    sidecar_name = final_out.name + "_data"
    if (final_out.parent / sidecar_name).exists():
        (final_out.parent / sidecar_name).unlink()
    if final_out.exists():
        final_out.unlink()
    final_out.parent.mkdir(parents=True, exist_ok=True)

    onnx.save_model(
        proto,
        str(final_out),
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=sidecar_name,
        size_threshold=1024,
        convert_attribute=False,
    )
    onnx.checker.check_model(str(final_out), full_check=False)


# ---------------------------------------------------------------------------
# Main entrypoint
# ---------------------------------------------------------------------------


def convert_graph(args: argparse.Namespace) -> None:
    import re

    input_path: Path = args.input.resolve()
    output_path: Path = args.output.resolve()
    if not input_path.exists():
        raise SystemExit(f"input not found: {input_path}")

    op_block_list = [s.strip() for s in args.op_block_list.split(",") if s.strip()]
    node_block_list = [s.strip() for s in args.node_block_list.split(",") if s.strip()]

    # Resolve --node-block-pattern by walking the graph proto (cheap; we only
    # touch the small protobuf, not the external sidecar).
    if args.node_block_pattern:
        pattern = re.compile(args.node_block_pattern)
        proto_only = onnx.load(str(input_path), load_external_data=False)
        for node in proto_only.graph.node:
            if node.name and pattern.search(node.name) and node.name not in node_block_list:
                node_block_list.append(node.name)
        print(
            f"  --node-block-pattern matched {sum(1 for n in proto_only.graph.node if n.name and pattern.search(n.name))} nodes "
            f"(adding to FP32 block list)"
        )

    print(f"input:        {input_path}")
    print(f"output:       {output_path}")
    print(f"precision:    {args.precision}")
    print(f"keep_io:      {args.keep_io_types}")
    print(f"op_block:     {op_block_list}")
    if node_block_list:
        print(f"node_block:   {node_block_list}")

    fp32_size = input_path.stat().st_size
    fp32_sidecar = input_path.with_name(input_path.name + "_data")
    fp32_data_size = fp32_sidecar.stat().st_size if fp32_sidecar.exists() else 0
    print(
        f"  fp32 graph={fp32_size / 1e6:.2f} MB  "
        f"sidecar={fp32_data_size / 1e9:.2f} GB"
    )

    if args.precision == "fp16":
        proto = convert_to_fp16(input_path, op_block_list, node_block_list, args.keep_io_types)
    elif args.precision == "fp16w":
        proto = convert_to_weights_fp16(
            input_path,
            skip_small_elems=args.fp16w_skip_small,
            skip_norm_max_elems=args.fp16w_skip_norm_max_elems,
        )
    else:
        proto = convert_to_bf16(input_path, op_block_list, node_block_list, args.keep_io_types)

    # The FP16 helper occasionally inserts duplicate Cast nodes when two
    # blocked nodes share an input tensor (it doesn't dedup by input name),
    # which fails the SSA check. Collapse those before save.
    removed = _dedup_inserted_casts(proto)
    if removed:
        print(f"  removed {removed} duplicate Cast nodes inserted by the helper")

    # Both paths can leave the graph topologically out-of-order (the FP16
    # converter inserts Cast nodes around blocked ops; the BF16 walker
    # appends new Cast nodes at the end). Re-sort before save so
    # onnx.checker is happy.
    proto = _topo_sort(proto)

    save_with_single_sidecar(proto, output_path, args.ir_version)

    domains = assert_pure_ai_onnx(output_path)
    out_size = output_path.stat().st_size
    out_data = output_path.with_name(output_path.name + "_data")
    out_data_size = out_data.stat().st_size if out_data.exists() else 0
    print(
        f"  saved {output_path} (+ {out_data.name})  "
        f"graph={out_size / 1e6:.2f} MB  sidecar={out_data_size / 1e9:.2f} GB"
    )
    print(f"  node-domains={domains}")
    if fp32_data_size > 0:
        ratio = out_data_size / fp32_data_size
        print(f"  sidecar size ratio ({args.precision} / fp32) = {ratio:.3f}")


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    try:
        convert_graph(args)
    except RuntimeError as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        raise SystemExit(2) from exc


if __name__ == "__main__":
    main()