"""
Rebuild BiRefNet-lite ONNX so it runs on any WebGPU device.

The official export has Concat nodes with up to 1024 inputs and Split nodes
with up to 32 outputs. ONNX Runtime's WebGPU backend binds every input and
output of a node as a storage buffer, and GPUs cap that per shader
(maxStorageBuffersPerShaderStage: 8 minimum, 10 or 16 on many devices), so
those nodes fail with "Too many storage buffers in shader".

This replaces them with trees of small nodes that compute exactly the same
values (concatenation and splitting are associative along one axis):
  - Concat: at most CONCAT_FAN_IN inputs per node (+1 output)
  - Split:  at most SPLIT_FAN_OUT outputs per node (+ data + sizes inputs)
Every node then uses at most MAX_BUFFERS storage buffers.

It also replaces Sum with a chain of Add nodes: ONNX Runtime's WebGPU backend
has no Sum kernel, so Sum runs on the CPU in WebAssembly memory, and the
deformable convolutions' Sums of four 392 MiB tensors exhaust its 4 GB heap
("std::bad_alloc").

Usage: python rebuild.py <in.onnx> <out.onnx>
"""

import sys

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

CONCAT_FAN_IN = 6
SPLIT_FAN_OUT = 5
MAX_BUFFERS = 7


def buffers(node: onnx.NodeProto) -> int:
    return len([i for i in node.input if i]) + len([o for o in node.output if o])


class Rewriter:
    def __init__(self, model: onnx.ModelProto):
        self.model = model
        self.graph = model.graph
        self.counter = 0
        self.initializers = {t.name: t for t in self.graph.initializer}
        self.constants = {
            n.output[0]: n for n in self.graph.node if n.op_type == "Constant" and len(n.output) == 1
        }
        self.names = {o for n in self.graph.node for o in n.output} | set(self.initializers)
        self.names |= {i.name for i in self.graph.input}

    def fresh(self, base: str) -> str:
        while True:
            self.counter += 1
            name = f"{base}/wgpu_{self.counter}"
            if name not in self.names:
                self.names.add(name)
                return name

    def const_value(self, name: str) -> np.ndarray:
        if name in self.initializers:
            return numpy_helper.to_array(self.initializers[name])
        if name in self.constants:
            attr = next(a for a in self.constants[name].attribute if a.name == "value")
            return numpy_helper.to_array(attr.t)
        raise ValueError(f"split sizes '{name}' are not constant")

    def int64_initializer(self, base: str, values: list[int]) -> str:
        name = self.fresh(base + "/sizes")
        tensor = numpy_helper.from_array(np.array(values, dtype=np.int64), name)
        self.graph.initializer.append(tensor)
        self.initializers[name] = tensor
        return name

    def concat_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
        axis = next(a.i for a in node.attribute if a.name == "axis")
        inputs = list(node.input)
        made: list[onnx.NodeProto] = []
        while len(inputs) > CONCAT_FAN_IN:
            grouped = []
            for start in range(0, len(inputs), CONCAT_FAN_IN):
                chunk = inputs[start : start + CONCAT_FAN_IN]
                if len(chunk) == 1:
                    grouped.append(chunk[0])
                    continue
                out = self.fresh(node.name)
                made.append(helper.make_node("Concat", chunk, [out], name=self.fresh(node.name), axis=axis))
                grouped.append(out)
            inputs = grouped
        made.append(helper.make_node("Concat", inputs, list(node.output), name=node.name, axis=axis))
        return made

    def split_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
        axis = next((a.i for a in node.attribute if a.name == "axis"), 0)
        if len(node.input) < 2 or not node.input[1]:
            raise ValueError(f"{node.name}: Split without explicit sizes isn't handled")
        sizes = [int(v) for v in self.const_value(node.input[1]).tolist()]
        outputs = list(node.output)
        assert len(sizes) == len(outputs)
        made: list[onnx.NodeProto] = []

        def split(data: str, parts: list[tuple[int, str]], base: str) -> None:
            # parts: (size, final output name)
            if len(parts) <= SPLIT_FAN_OUT:
                made.append(
                    helper.make_node(
                        "Split",
                        [data, self.int64_initializer(base, [s for s, _ in parts])],
                        [o for _, o in parts],
                        name=self.fresh(base),
                        axis=axis,
                    )
                )
                return
            # Groups of roughly equal count, each split again below.
            per_group = -(-len(parts) // SPLIT_FAN_OUT)
            groups = [parts[i : i + per_group] for i in range(0, len(parts), per_group)]
            group_outputs = []
            for group in groups:
                group_outputs.append(group[0][1] if len(group) == 1 else self.fresh(base))
            made.append(
                helper.make_node(
                    "Split",
                    [data, self.int64_initializer(base, [sum(s for s, _ in g) for g in groups])],
                    group_outputs,
                    name=self.fresh(base),
                    axis=axis,
                )
            )
            for group, out in zip(groups, group_outputs):
                if len(group) > 1:
                    split(out, group, base)

        split(node.input[0], list(zip(sizes, outputs)), node.name)
        return made

    def add_chain(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
        # ((a + b) + c) + d: the same order Sum adds in, but Add has a WebGPU kernel.
        inputs = [i for i in node.input if i]
        if len(inputs) == 1:
            return [helper.make_node("Identity", inputs, list(node.output), name=node.name)]
        made: list[onnx.NodeProto] = []
        acc = inputs[0]
        for index, other in enumerate(inputs[1:], start=2):
            out = node.output[0] if index == len(inputs) else self.fresh(node.name)
            made.append(helper.make_node("Add", [acc, other], [out], name=self.fresh(node.name)))
            acc = out
        return made

    def run(self) -> dict[str, int]:
        stats = {"Concat": 0, "Split": 0, "Sum": 0}
        new_nodes: list[onnx.NodeProto] = []
        for node in self.graph.node:
            if node.op_type == "Sum":
                # ONNX Runtime's WebGPU backend has no Sum kernel, so Sum runs on the
                # CPU, inside WebAssembly's 4 GB heap. In the deformable convolutions it
                # adds four 392 MiB tensors (784 MiB each after the fp16 → fp32 casts
                # the CPU needs): "std::bad_alloc".
                new_nodes.extend(self.add_chain(node))
                stats["Sum"] += 1
            elif buffers(node) > MAX_BUFFERS and node.op_type == "Concat":
                new_nodes.extend(self.concat_tree(node))
                stats["Concat"] += 1
            elif buffers(node) > MAX_BUFFERS and node.op_type == "Split":
                new_nodes.extend(self.split_tree(node))
                stats["Split"] += 1
            else:
                new_nodes.append(node)
        del self.graph.node[:]
        self.graph.node.extend(new_nodes)
        return stats


def main() -> None:
    source, target = sys.argv[1], sys.argv[2]
    model = onnx.load(source)
    stats = Rewriter(model).run()
    too_wide = [(n.op_type, n.name, buffers(n)) for n in model.graph.node if buffers(n) > MAX_BUFFERS]
    if too_wide:
        raise SystemExit(f"still too wide: {too_wide[:10]}")
    onnx.checker.check_model(model)
    onnx.save(model, target)
    widest = max(buffers(n) for n in model.graph.node)
    print(f"{source} -> {target}: rewrote {stats}, nodes {len(model.graph.node)}, widest node {widest} buffers")


if __name__ == "__main__":
    main()