"""Rewrites a constant-folded BiRefNet graph so every WebGPU shader fits in 8 storage buffers and nothing falls back to the CPU."""
import math
import sys
from collections import Counter, defaultdict

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

FAN = 6

source, target = sys.argv[1], sys.argv[2]
model = shape_inference.infer_shapes(onnx.load(source))
graph = model.graph
shapes = {}
for value in list(graph.value_info) + list(graph.input) + list(graph.output):
    t = value.type.tensor_type
    shapes[value.name] = ([d.dim_value for d in t.shape.dim], t.elem_type)
inits = {i.name: i for i in graph.initializer}
for name, init in inits.items():
    shapes[name] = (list(init.dims), init.data_type)

counter = Counter()


def fresh(base):
    counter[base] += 1
    return f"wonfy/{base}/{counter[base]}"


new_inits = []


def const(base, array):
    name = fresh(base)
    new_inits.append(numpy_helper.from_array(np.asarray(array), name))
    return name


def attr(node, name, default=None):
    for a in node.attribute:
        if a.name == name:
            return helper.get_attribute_value(a)
    return default


consumers = defaultdict(list)
for n in graph.node:
    for x in n.input:
        consumers[x].append(n)
producers = {out: n for n in graph.node for out in n.output}

# The int64 index chain of each deformable-conv sample, done in float instead (the values are whole numbers).
cast_nodes = [n for n in graph.node if n.op_type == "Cast" and attr(n, "to") == TensorProto.INT64]
region = set()
frontier = [n.output[0] for n in cast_nodes]
while frontier:
    name = frontier.pop()
    if name in region:
        continue
    region.add(name)
    for c in consumers[name]:
        if c.op_type in ("Slice", "Add", "Clip", "Reshape"):
            frontier.append(c.output[0])
        elif c.op_type == "Concat":
            region.add(c.output[0])
        elif c.op_type != "GatherND":
            raise SystemExit(f"unexpected consumer {c.op_type} of int64 tensor {name}")
cast_source = {n.output[0]: n.input[0] for n in cast_nodes}

drop = set(id(n) for n in cast_nodes)
renames = {}
for n in graph.node:
    if n.op_type in ("Add", "Clip") and n.output[0] in region:
        for i, x in enumerate(n.input):
            if x in inits and inits[x].data_type == TensorProto.INT64:
                n.input[i] = const("float_const", numpy_helper.to_array(inits[x]).astype(np.float32))
for n in graph.node:
    for i, x in enumerate(n.input):
        if x in cast_source:
            n.input[i] = cast_source[x]

out_nodes = []
stats = Counter()


def emit(op, inputs, outputs, **attrs):
    out_nodes.append(helper.make_node(op, inputs, outputs, name=fresh(op.lower()), **attrs))


def concat_tree(inputs, output, axis):
    while len(inputs) > FAN:
        grouped = []
        for start in range(0, len(inputs), FAN):
            group = inputs[start:start + FAN]
            if len(group) == 1:
                grouped.append(group[0])
                continue
            name = fresh("concat")
            emit("Concat", group, [name], axis=axis)
            grouped.append(name)
        inputs = grouped
    emit("Concat", inputs, [output], axis=axis)


def split_tree(source_name, sizes, outputs, axis):
    if len(outputs) <= FAN:
        emit("Split", [source_name, const("split_sizes", np.array(sizes, dtype=np.int64))], outputs, axis=axis)
        return
    size = math.ceil(len(outputs) / FAN)
    parts, part_sizes = [], []
    for start in range(0, len(outputs), size):
        group_outputs, group_sizes = outputs[start:start + size], sizes[start:start + size]
        name = group_outputs[0] if len(group_outputs) == 1 else fresh("split")
        parts.append((name, group_sizes, group_outputs))
        part_sizes.append(sum(group_sizes))
    split_tree(source_name, part_sizes, [p[0] for p in parts], axis)
    for name, group_sizes, group_outputs in parts:
        if len(group_outputs) > 1:
            split_tree(name, group_sizes, group_outputs, axis)


concat_inputs = {n.output[0]: list(n.input) for n in graph.node if n.op_type == "Concat" and n.output[0] in region}

for n in graph.node:
    if id(n) in drop:
        continue
    if n.op_type == "Concat" and n.output[0] in region:
        stats["index concat removed"] += 1
        continue
    if n.op_type == "Concat" and len(n.input) > FAN:
        concat_tree(list(n.input), n.output[0], attr(n, "axis"))
        stats["concat"] += 1
        continue
    if n.op_type == "Split" and len(n.output) > FAN:
        sizes = numpy_helper.to_array(inits[n.input[1]]).tolist()
        axis = attr(n, "axis", 0)
        split_tree(n.input[0], sizes, list(n.output), axis)
        stats["split"] += 1
        continue
    if n.op_type == "Sum":
        total = n.input[0]
        for i, x in enumerate(n.input[1:], start=1):
            name = n.output[0] if i == len(n.input) - 1 else fresh("sum")
            emit("Add", [total, x], [name])
            total = name
        stats["sum"] += 1
        continue
    if n.op_type == "GatherND":
        assert attr(n, "batch_dims") == 2, "only batch_dims 2 is handled"
        data, index = n.input
        data_shape = shapes[data][0]
        assert data_shape[:2] == [1, 1], data_shape
        _, _, height, width, channels = data_shape
        y, x = concat_inputs[index]
        count = shapes[n.output[0]][0][2]
        y32, x32, row, flat, flat1, table, picked = (fresh(k) for k in ("y32", "x32", "row", "flat", "flat1", "table", "picked"))
        emit("Cast", [y], [y32], to=TensorProto.INT32)
        emit("Cast", [x], [x32], to=TensorProto.INT32)
        emit("Mul", [y32, const("width", np.array(width, dtype=np.int32))], [row])
        emit("Add", [row, x32], [flat])
        emit("Reshape", [flat, const("shape", np.array([count], dtype=np.int64))], [flat1])
        emit("Reshape", [data, const("shape", np.array([height * width, channels], dtype=np.int64))], [table])
        emit("Gather", [table, flat1], [picked], axis=0)
        emit("Reshape", [picked, const("shape", np.array([1, 1, count, channels], dtype=np.int64))], [n.output[0]])
        stats["gathernd"] += 1
        continue
    out_nodes.append(n)

del graph.node[:]
graph.node.extend(out_nodes)
graph.initializer.extend(new_inits)
del graph.value_info[:]
used = {x for n in graph.node for x in n.input}
kept = [i for i in graph.initializer if i.name in used]
del graph.initializer[:]
graph.initializer.extend(kept)
del model.opset_import[:]
model.opset_import.extend([helper.make_opsetid("", 17)])

model = shape_inference.infer_shapes(model, strict_mode=True)
onnx.checker.check_model(model)
onnx.save(model, target)
print("rewrote", dict(stats))

# Audit: buffers per node, and anything WebGPU would run on the CPU.
graph = model.graph
types = {v.name: v.type.tensor_type.elem_type for v in list(graph.value_info) + list(graph.output) + list(graph.input)}
cpu_inputs = {("Reshape", 1), ("Split", 1), ("Slice", 1), ("Slice", 2), ("Slice", 3), ("Slice", 4), ("Clip", 1), ("Clip", 2), ("Pad", 1), ("Pad", 2), ("Resize", 1), ("Resize", 2), ("Resize", 3)}
worst = Counter()
for n in graph.node:
    buffers = sum(1 for i, x in enumerate(n.input) if x and (n.op_type, i) not in cpu_inputs) + len(n.output)
    worst[n.op_type] = max(worst[n.op_type], buffers)
print("most buffers per op:", dict(sorted(worst.items(), key=lambda kv: -kv[1])))
int64 = Counter(n.op_type for n in graph.node if any(types.get(o) == TensorProto.INT64 for o in n.output))
print("ops still making int64:", dict(int64) or "none")
print("ops:", dict(Counter(n.op_type for n in graph.node).most_common()))