#!/usr/bin/env python3
"""Extract YOLO26 raw detection heads from a standard Ultralytics ONNX export."""

from __future__ import annotations

import argparse
import re
from pathlib import Path

import onnx
from onnx import TensorProto, helper, shape_inference


TERMINAL_HEAD = re.compile(
    r"one2one_cv([23])\.(\d+)/one2one_cv\1\.\2\.2/Conv$"
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    model = shape_inference.infer_shapes(onnx.load(args.input))
    shapes = {
        value.name: [dim.dim_value for dim in value.type.tensor_type.shape.dim]
        for value in list(model.graph.value_info) + list(model.graph.output)
    }
    heads: dict[tuple[int, int], str] = {}
    for node in model.graph.node:
        match = TERMINAL_HEAD.search(node.name)
        if match and node.op_type == "Conv":
            branch, scale = int(match.group(1)), int(match.group(2))
            heads[(branch, scale)] = node.output[0]

    expected = {(branch, scale) for branch in (2, 3) for scale in range(3)}
    if set(heads) != expected:
        missing = sorted(expected - set(heads))
        raise SystemExit(f"could not locate all YOLO26 one-to-one heads; missing {missing}")

    del model.graph.output[:]
    output_names = []
    for scale, stride in enumerate((8, 16, 32)):
        for branch, kind in ((3, "cls"), (2, "box")):
            source = heads[(branch, scale)]
            source_shape = shapes[source]
            if len(source_shape) != 4:
                raise SystemExit(f"unexpected head shape for {source}: {source_shape}")
            output_name = f"{kind}_s{stride}"
            model.graph.node.append(
                helper.make_node(
                    "Transpose",
                    [source],
                    [output_name],
                    name=f"bpu_output_{kind}_s{stride}",
                    perm=[0, 2, 3, 1],
                )
            )
            model.graph.output.append(
                helper.make_tensor_value_info(
                    output_name,
                    TensorProto.FLOAT,
                    [source_shape[0], source_shape[2], source_shape[3], source_shape[1]],
                )
            )
            output_names.append(output_name)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_suffix(".unpruned.onnx")
    onnx.save(model, temporary)
    onnx.utils.extract_model(str(temporary), str(args.output), [model.graph.input[0].name], output_names)
    temporary.unlink()
    extracted = onnx.load(args.output)
    onnx.checker.check_model(extracted)
    print(f"wrote {args.output} with outputs:")
    for output in extracted.graph.output:
        dims = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
        print(f"  {output.name}: {dims}")


if __name__ == "__main__":
    main()