#!/usr/bin/env python3
"""Replace the policy's grouped 1x1 Conv expert layer with an equivalent Gemm.

The RobotLab MoE expert layer is exported as:

    [1, 2048] -> Unsqueeze -> grouped Conv1d(group=8, kernel=1)
              -> Squeeze -> [1, 256]

For this fixed shape, the grouped Conv is exactly a block-diagonal linear layer.
Using Gemm lets hb_mapper keep the expert path on BPU instead of cutting through
CPU float execution.
"""

import argparse
from pathlib import Path

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


CONV_NAME = "/student_moe_encoder/moe/experts/experts/Conv"


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


def initializer_map(model):
    return {initializer.name: initializer for initializer in model.graph.initializer}


def producer_map(model):
    producers = {}
    for node in model.graph.node:
        for output in node.output:
            producers[output] = node
    return producers


def replace_node(model):
    initializers = initializer_map(model)
    producers = producer_map(model)
    nodes = list(model.graph.node)

    conv = next((node for node in nodes if node.name == CONV_NAME), None)
    if conv is None:
        candidates = [
            node for node in nodes
            if node.op_type == "Conv"
            and int(attr_value(node, "group", 1)) > 1
            and attr_value(node, "kernel_shape") == [1]
        ]
        if len(candidates) == 1:
            conv = candidates[0]
            print(f"[WARN] {CONV_NAME!r} not found; using grouped Conv {conv.name!r}")
    if conv is None:
        raise ValueError(f"Cannot find unique grouped 1x1 Conv node; fixed name {CONV_NAME!r} not found")
    if conv.op_type != "Conv":
        raise ValueError(f"{CONV_NAME!r} is {conv.op_type}, expected Conv")

    if attr_value(conv, "kernel_shape") != [1]:
        raise ValueError("Only kernel_shape=[1] is supported")
    if attr_value(conv, "strides") not in (None, [1]):
        raise ValueError("Only stride=1 is supported")
    if attr_value(conv, "pads") not in (None, [0, 0]):
        raise ValueError("Only no padding is supported")

    groups = int(attr_value(conv, "group", 1))
    weight = numpy_helper.to_array(initializers[conv.input[1]])
    bias = numpy_helper.to_array(initializers[conv.input[2]]) if len(conv.input) >= 3 else None
    if weight.ndim != 3 or weight.shape[2] != 1:
        raise ValueError(f"Expected Conv weight [out, in_per_group, 1], got {weight.shape}")

    out_channels, in_per_group, _ = weight.shape
    if out_channels % groups != 0:
        raise ValueError(f"out_channels={out_channels} is not divisible by groups={groups}")
    out_per_group = out_channels // groups
    in_channels = in_per_group * groups

    full_weight = np.zeros((out_channels, in_channels), dtype=np.float32)
    flat_weight = weight[:, :, 0].astype(np.float32, copy=False)
    for group in range(groups):
        out_begin = group * out_per_group
        out_end = out_begin + out_per_group
        in_begin = group * in_per_group
        in_end = in_begin + in_per_group
        full_weight[out_begin:out_end, in_begin:in_end] = flat_weight[out_begin:out_end]

    unsqueeze = producers.get(conv.input[0])
    if unsqueeze is None or unsqueeze.op_type != "Unsqueeze":
        raise ValueError("Expected Conv input to be produced by Unsqueeze")
    squeeze_users = [node for node in nodes if conv.output[0] in node.input]
    if len(squeeze_users) != 1 or squeeze_users[0].op_type != "Squeeze":
        raise ValueError("Expected Conv output to feed exactly one Squeeze")
    squeeze = squeeze_users[0]

    gemm_weight_name = conv.input[1] + "_as_gemm"
    model.graph.initializer.append(
        numpy_helper.from_array(full_weight, name=gemm_weight_name)
    )

    gemm_inputs = [unsqueeze.input[0], gemm_weight_name]
    if bias is not None:
        gemm_inputs.append(conv.input[2])

    gemm = helper.make_node(
        "Gemm",
        inputs=gemm_inputs,
        outputs=[squeeze.output[0]],
        name=CONV_NAME + "_as_Gemm",
        alpha=1.0,
        beta=1.0,
        transB=1,
    )

    conv_index = nodes.index(conv)
    remove_names = {unsqueeze.name, conv.name, squeeze.name}
    new_nodes = []
    inserted = False
    for index, node in enumerate(nodes):
        if node.name in remove_names:
            if index >= conv_index and not inserted:
                new_nodes.append(gemm)
                inserted = True
            continue
        new_nodes.append(node)

    del model.graph.node[:]
    model.graph.node.extend(new_nodes)
    return {
        "groups": groups,
        "conv_weight_shape": list(weight.shape),
        "gemm_weight_shape": list(full_weight.shape),
        "removed": sorted(remove_names),
        "added": gemm.name,
    }


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

    model = onnx.load(str(args.input))
    info = replace_node(model)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(args.output))
    print(f"Wrote {args.output}")
    for key, value in info.items():
        print(f"{key}={value}")


if __name__ == "__main__":
    main()
