import sys
import numpy as np
from onnx_modifier import ONNXModifier


def change_inf_to_value(om: ONNXModifier):
    records = set()
    for where_node in om.get_nodes("Where"):
        for input_name in where_node.inputs[1:]:
            init = om.get_initializer(input_name)
            if init is None:
                continue
            assert input_name == init.name
            init_name = input_name
            if init_name in records:
                continue

            # info = np.finfo(np.float32)
            info = np.finfo(np.float16)
            data = om.get_initializer_value(init.name)
            if data.size > 1:
                continue
            if data == np.inf:
                om.set_initializer_value(init_name, np.array(info.max, dtype=np.float32))
            elif data == -np.inf:
                om.set_initializer_value(init_name, np.array(info.min, dtype=np.float32))
            else:
                continue
            # print("Changed value:", init_name)
            records.add(init_name)


# def optimize_where_ndoes(om: ONNXModifier):
#     """Where节点等价替换
#     (1) condition为initializer, X为0, Y为输入数据:
#         Where(cond, X, Y) ==> Mul(Y, ~cond)
#     (2) condition为initializer, X为负无穷, Y为输入数据
#         Where(cond, X, Y) ==>  Sub(Y, Where(cond, np.inf, 0))
#     (3) condition为真实输入, X为负无穷, Y为输入数据
#         Where(cond, X, Y) ==> Sub(Y, Mul(Cast(cond, to=float32), np.inf))

#     cases:
#         1. Where(cond, -inf, input)
#             a. /transformer/encoder/fusion_layers.*/attn/Where 
#             b. /transformer/encoder/fusion_layers.*/attn/Where_1
#             c. /class_embed.0_*/Where: Where(cond, -inf, input)
#         2. Where(cond, 0, input):
#             a. /transformer/encoder/layers.*/self_attn/Where
#             b. /transformer/decoder/layers.*/cross_attn/Where
#     """

#     for where_node in om.get_nodes("Where"):
#         where_name = where_node.name
#         # print("Process where node:", where_name)
#         x_value = om.get_initializer_value(where_node.inputs[1])
#         assert x_value.size == 1
#         assert x_value == np.array(0.0, dtype=np.float32) or \
#             x_value == np.array(-np.inf, dtype=np.float32)

#         cond_init = om.get_initializer(where_node.inputs[0])
#         if cond_init is not None:
#             cond_value = om.get_initializer_value(where_node.inputs[0])
#             if x_value == np.array(0.0, dtype=np.float32):
#                 # Where(cond, X, Y) ==> Mul(Y, ~cond)
#                 mul_name = where_name.replace("Where", "NewMul")
#                 mul_b_init = om.create_initializer(mul_name + "_B", 
#                                                    (~cond_value).astype(np.float32))
#                 mul_node = om.create_node("Mul", 
#                                           mul_name, 
#                                           [where_node.inputs[2], mul_b_init.name], 
#                                           [mul_name+"_output_0"], 
#                                           index=where_node.index)
#                 next_nodes = where_node.next_nodes
#                 for next_node in next_nodes:
#                     next_node.replace_input(where_node.outputs[0], mul_node.outputs[0])
#             elif x_value == np.array(-np.inf, dtype=np.float32):
#                 # Where(cond, X, Y) ==>  Sub(Y, Where(cond, np.inf, 0))
#                 sub_name = where_name.replace("Where", "NewSub")
#                 sub_b_init = om.create_initializer(
#                     sub_name + "_B", 
#                     np.where(cond_value.astype(np.float32), 
#                              np.finfo(np.float16).max, 0.0).astype(np.float32)
#                 )
#                 sub_node = om.create_node("Sub", 
#                                           sub_name, 
#                                           [where_node.inputs[2], sub_b_init.name], 
#                                           [sub_name+"_output_0"], 
#                                           index=where_node.index)
#                 next_nodes = where_node.next_nodes
#                 for next_node in next_nodes:
#                     next_node.replace_input(where_node.outputs[0], sub_node.outputs[0])
#         else:
#             # Where(cond, X, Y) ==> Sub(Y, Mul(Cast(cond, to=float32), np.inf))
#             assert x_value == np.array(-np.inf, dtype=np.float32)
#             cast_name = where_name.replace("Where", "NewCast")
#             mul_name = where_name.replace("Where", "NewMul")
#             sub_name = where_name.replace("Where", "NewSub")
#             cast_node = om.create_node("Cast", 
#                                        cast_name, 
#                                        [where_node.inputs[0]], 
#                                        [cast_name+"_output_0"], 
#                                        to=1, 
#                                        index=where_node.index)
#             mul_b_init = om.create_initializer(mul_name + "_B", 
#                                                np.array([np.finfo(np.float16).max], np.float32))
#             mul_node = om.create_node("Mul", 
#                                       mul_name, 
#                                       [cast_node.outputs[0], mul_b_init.name], 
#                                       [mul_name+"_output_0"], 
#                                       index=cast_node.index+1)
#             sub_node = om.create_node("Sub", 
#                                       sub_name, 
#                                       [where_node.inputs[2], mul_node.outputs[0]], 
#                                       [sub_name+"_output_0"], 
#                                       index=mul_node.index+1)
#             next_nodes = where_node.next_nodes
#             for next_node in next_nodes:
#                 next_node.replace_input(where_node.outputs[0], sub_node.outputs[0])
    
#     om.update_map()

def optimize_where_ndoes(om: ONNXModifier):
    """Where节点等价替换 (加入安全校验版本)"""

    for where_node in om.get_nodes("Where"):
        where_name = where_node.name
        
        # 1. 安全获取 X 的值，如果 X 不是常量（initializer），直接跳过不优化
        x_init = om.get_initializer(where_node.inputs[1])
        if x_init is None:
            continue
            
        x_value = om.get_initializer_value(where_node.inputs[1])
        
        # 2. 避免 assert 崩溃：如果 size 不为 1，说明不是我们要找的 Attention Mask 节点，跳过
        if x_value.size != 1:
            continue
            
        # 3. 判断是否符合优化条件（0.0 或 -inf），不符合直接跳过
        is_zero = (x_value == np.array(0.0, dtype=np.float32))
        is_neg_inf = (x_value == np.array(-np.inf, dtype=np.float32))
        
        if not (is_zero or is_neg_inf):
            continue

        cond_init = om.get_initializer(where_node.inputs[0])
        if cond_init is not None:
            cond_value = om.get_initializer_value(where_node.inputs[0])
            if is_zero:
                # Where(cond, X, Y) ==> Mul(Y, ~cond)
                mul_name = where_name.replace("Where", "NewMul")
                mul_b_init = om.create_initializer(mul_name + "_B", 
                                                   (~cond_value).astype(np.float32))
                mul_node = om.create_node("Mul", 
                                          mul_name, 
                                          [where_node.inputs[2], mul_b_init.name], 
                                          [mul_name+"_output_0"], 
                                          index=where_node.index)
                next_nodes = where_node.next_nodes
                for next_node in next_nodes:
                    next_node.replace_input(where_node.outputs[0], mul_node.outputs[0])
            elif is_neg_inf:
                # Where(cond, X, Y) ==>  Sub(Y, Where(cond, np.inf, 0))
                sub_name = where_name.replace("Where", "NewSub")
                sub_b_init = om.create_initializer(
                    sub_name + "_B", 
                    np.where(cond_value.astype(np.float32), 
                             np.finfo(np.float16).max, 0.0).astype(np.float32)
                )
                sub_node = om.create_node("Sub", 
                                          sub_name, 
                                          [where_node.inputs[2], sub_b_init.name], 
                                          [sub_name+"_output_0"], 
                                          index=where_node.index)
                next_nodes = where_node.next_nodes
                for next_node in next_nodes:
                    next_node.replace_input(where_node.outputs[0], sub_node.outputs[0])
        else:
            # Where(cond, X, Y) ==> Sub(Y, Mul(Cast(cond, to=float32), np.inf))
            # 当 condition 不是 initializer 时，只处理 -inf 的情况
            if not is_neg_inf:
                continue
                
            cast_name = where_name.replace("Where", "NewCast")
            mul_name = where_name.replace("Where", "NewMul")
            sub_name = where_name.replace("Where", "NewSub")
            cast_node = om.create_node("Cast", 
                                       cast_name, 
                                       [where_node.inputs[0]], 
                                       [cast_name+"_output_0"], 
                                       to=1, 
                                       index=where_node.index)
            mul_b_init = om.create_initializer(mul_name + "_B", 
                                               np.array([np.finfo(np.float16).max], np.float32))
            mul_node = om.create_node("Mul", 
                                      mul_name, 
                                      [cast_node.outputs[0], mul_b_init.name], 
                                      [mul_name+"_output_0"], 
                                      index=cast_node.index+1)
            sub_node = om.create_node("Sub", 
                                      sub_name, 
                                      [where_node.inputs[2], mul_node.outputs[0]], 
                                      [sub_name+"_output_0"], 
                                      index=mul_node.index+1)
            next_nodes = where_node.next_nodes
            for next_node in next_nodes:
                next_node.replace_input(where_node.outputs[0], sub_node.outputs[0])
    
    om.update_map()

def optimize_transpose_nodes(om: ONNXModifier):
    transpose_list = [
        "/transformer/encoder/Transpose",
        "/transformer/encoder/Transpose_1",
        "/transformer/encoder/Transpose_2",
        "/transformer/encoder/Transpose_3",
        "/transformer/encoder/Transpose_4",
        "/transformer/encoder/Transpose_5",
        "/transformer/encoder/Transpose_6",
        "/transformer/encoder/Transpose_7",
        "/transformer/encoder/Transpose_8",
        "/transformer/encoder/Transpose_9",
        "/transformer/encoder/Transpose_10",
        "/transformer/encoder/Transpose_11",
        "/transformer/decoder/layers.0/Transpose",
        "/transformer/decoder/layers.0/Transpose_1",
        "/transformer/decoder/layers.0/Transpose_2",
        "/transformer/decoder/layers.1/Transpose",
        "/transformer/decoder/layers.1/Transpose_1",
        "/transformer/decoder/layers.1/Transpose_2",
        "/transformer/decoder/layers.2/Transpose",
        "/transformer/decoder/layers.2/Transpose_1",
        "/transformer/decoder/layers.2/Transpose_2",
        "/transformer/decoder/layers.3/Transpose",
        "/transformer/decoder/layers.3/Transpose_1",
        "/transformer/decoder/layers.3/Transpose_2",
        "/transformer/decoder/layers.4/Transpose",
        "/transformer/decoder/layers.4/Transpose_1",
        "/transformer/decoder/layers.4/Transpose_2",
        "/transformer/decoder/layers.5/Transpose",
        "/transformer/decoder/layers.5/Transpose_1",
        "/transformer/decoder/layers.5/Transpose_2",
        "/transformer/Transpose_8",
        "/transformer/decoder/Transpose",
        "/transformer/decoder/Transpose_1",
        "/transformer/decoder/Transpose_2",
        "/transformer/decoder/Transpose_3",
        "/transformer/decoder/Transpose_4",
        "/transformer/decoder/Transpose_5",
        "/transformer/decoder/Transpose_6",
        "/transformer/decoder/Transpose_7",
        "/transformer/decoder/Transpose_8",
        "/transformer/decoder/Transpose_9",
        "/transformer/decoder/Transpose_10",
        "/transformer/decoder/Transpose_11"        
    ]
    for name in transpose_list:
        node = om.get_node(name)
        # 安全校验：如果找不到这个节点，说明当前模型不需要优化这个点，跳过
        if node is None:
            continue
        
        if 'perm' in node.attrs and (node.attrs['perm'] == [1, 0 , 2] or node.attrs['perm'] == [1, 0 , 2, 3]):
            next_nodes = om.get_next_nodes(node)
            for node_ in next_nodes:
                node_.replace_input(node.outputs[0], node.inputs[0])

    # modify /transformer/encoder/text_layers.*/self_attn/Reshape_4
    shape_init1 = om.create_initializer(
        "/transformer/encoder/text_layers.x/self_attn/des_shape", 
        np.array([1, 4, 256], np.int64)
    )
    for i in range(6):
        reshape_node = om.get_node(f"/transformer/encoder/text_layers.{i}/self_attn/Reshape_4")
        if reshape_node is not None:
            reshape_node.set_input(1, shape_init1.name)

    # modify /transformer/enc_out_class_embed/Transpose
    trans_node = om.get_node("/transformer/enc_out_class_embed/Transpose")
    if trans_node is not None:
        trans_node.set_attribute("perm", [0, 2, 1])

    # modify /transformer/decoder/Reshape_* # 安全校验：避免写死的随机变量名 _v_5525 引发崩溃
    init_5525 = om.get_initializer("_v_5525")
    if init_5525 is not None:
        om.set_initializer_value("_v_5525", np.array([1, 900, -1], np.int64))

    # modify /transformer/decoder/layers.*/self_attn/Reshape_4
    # modify /transformer/decoder/layers.*/ca_text/Reshape_6
    shape_init3 = om.create_initializer(
        "/transformer/decoder/layers.x/self_attn_ca_text/des_shape", 
        np.array([1, 900, 256], np.int64)
    )
    for i in range(6):
        reshape_node1 = om.get_node(f"/transformer/decoder/layers.{i}/self_attn/Reshape_4")
        if reshape_node1 is not None:
            reshape_node1.set_input(1, shape_init3.name)
        reshape_node2 = om.get_node(f"/transformer/decoder/layers.{i}/ca_text/Reshape_6")
        if reshape_node2 is not None:
            reshape_node2.set_input(1, shape_init3.name)
    
    # modify /transformer/decoder/layers.0/Add
    # modify /transformer/decoder/layers.0/Add_1
    init_name = "/transformer/Tile_1_output_0"
    tile_init = om.get_initializer(init_name)
    if tile_init is not None:
        add_value = om.get_initializer_value(init_name)
        om.set_initializer_value(init_name, np.ascontiguousarray(add_value.transpose(1, 0, 2)))
    
    om.update_map()
    
    # 将形状推断包起来，防止自定义算子（MSDeformAttn）导致推理失败崩溃
    try:
        om.infer_shape(strict_mode=False)
    except Exception as e:
        print(f"[Warning] infer_shape 跳过 (可能由于自定义算子引起). 详细信息: {e}")

def optmize_sin_cos_block(om: ONNXModifier):
    node_pairs = [
        ("/transformer/decoder/Gather_1",  "/transformer/decoder/ref_point_head/layers.0/MatMul"),
        ("/transformer/decoder/Gather_6",  "/transformer/decoder/ref_point_head/layers.0_1/MatMul"),
        ("/transformer/decoder/Gather_11", "/transformer/decoder/ref_point_head/layers.0_2/MatMul"),
        ("/transformer/decoder/Gather_16", "/transformer/decoder/ref_point_head/layers.0_3/MatMul"),
        ("/transformer/decoder/Gather_21", "/transformer/decoder/ref_point_head/layers.0_4/MatMul"),
        ("/transformer/decoder/Gather_26", "/transformer/decoder/ref_point_head/layers.0_5/MatMul"),
    ]
    
    # 提前创建一些公用的 initializer
    unsqueeze_axes_init1 = om.create_initializer("/transformer/decoder/sin_cos_block/unsqueeze_axes1", np.array([3, 4], np.int64))
    slice_axes_init = om.create_initializer("/transformer/decoder/sin_cos_block/slice_axes", np.array([4], np.int64))
    slice_steps_init = om.create_initializer("/transformer/decoder/sin_cos_block/slice_steps", np.array([1], np.int64))
    slice_starts_init1 = om.create_initializer("/transformer/decoder/sin_cos_block/slice_starts1", np.array([0], np.int64))
    slice_ends_init1 = om.create_initializer("/transformer/decoder/sin_cos_block/slice_ends1", np.array([1], np.int64))
    slice_starts_init2 = om.create_initializer("/transformer/decoder/sin_cos_block/slice_steps2", np.array([1], np.int64))
    slice_ends_init2 = om.create_initializer("/transformer/decoder/sin_cos_block/slice_ends2", np.array([2], np.int64))
    reshape_init = om.create_initializer("/transformer/decoder/sin_cos_block/reshape_dst_shape", np.array([1, 900, -1], np.int64))
    
    for i, (gather_name, matmul_name) in enumerate(node_pairs):
        gather_node = om.get_node(gather_name)
        matmul_node = om.get_node(matmul_name)
        
        # 【安全校验】：如果找不到这一对节点，说明不需要/无法优化这个 block，直接跳过
        if gather_node is None or matmul_node is None:
            continue
            
        try:
            next_node = om.get_next_nodes(gather_node)[0]
            if next_node.op_type != "Mul": continue
            mul_init_value = om.get_initializer_value(next_node.inputs[1])
            if mul_init_value.size != 1: continue
            
            next_node = om.get_next_nodes(next_node)[0]
            if next_node.op_type != "Unsqueeze": continue
            next_node.set_inputs([gather_node.inputs[0], unsqueeze_axes_init1.name])
            
            next_node = om.get_next_nodes(next_node)[0]
            if next_node.op_type != "Div": continue
            div_init_value = om.get_initializer_value(next_node.inputs[1])
            new_value = (div_init_value / mul_init_value).reshape(1, 1, 1, 64, 2)
            new_init = om.create_initializer(next_node.name + "_B", new_value)
            next_node.set_input(1, new_init.name)

            next_nodes = om.get_next_nodes(next_node)
            if len(next_nodes) != 2 or not all(x.op_type == 'Slice' for x in next_nodes): continue
            
            sin_node, cos_node = None, None
            for j, slice_node in enumerate(next_nodes):
                slice_node.set_inputs([slice_node.inputs[0], 
                                       slice_starts_init1.name if j == 0 else slice_starts_init2.name,
                                       slice_ends_init1.name if j == 0 else slice_ends_init2.name,
                                       slice_axes_init.name,
                                       slice_steps_init.name])
                n_node = om.get_next_nodes(slice_node)[0]
                if n_node.op_type == "Sin":
                    sin_node = n_node
                elif n_node.op_type == "Cos":
                    cos_node = n_node
                else:
                    raise RuntimeError("match fail!")
                n_node = om.get_next_nodes(n_node)[0]
                n_node = om.get_next_nodes(n_node)[0]

            next_node = n_node # Concat node
            if next_node.op_type != "Concat": continue
            next_node.set_inputs([sin_node.outputs[0], cos_node.outputs[0]])
            next_node.set_attribute("axis", 4)
            
            next_node = om.get_next_nodes(next_node)[0]
            if next_node.op_type != "Reshape": continue
            next_node.set_input(1, reshape_init.name)

            matmul_node.set_input(0, next_node.outputs[0])
            if i == 0:
                mm_b_value = om.get_initializer_value(matmul_node.inputs[1])
                mm_b_value = np.concatenate([mm_b_value[128:256, ...], 
                                             mm_b_value[0:128, ...], 
                                             mm_b_value[256:, ...]], 
                                            axis=0)
                om.set_initializer_value(matmul_node.inputs[1], mm_b_value)
        except Exception as e:
            # 如果匹配过程中发生任何形状或节点断层的意外，静默跳过这个 block
            continue
            
    om.update_map()
    try:
        om.infer_shape(strict_mode=False)
    except:
        pass


def fuse_one_attention(om: ONNXModifier, softmax_name: str, new_mask: bool = None, num_heads: int = 12):
    softmax_node = om.get_node(softmax_name)
    tmp_node = om.get_prev_nodes(softmax_node)[0]
    assert tmp_node.op_type in ["MatMul", "Add"]

    mask = None
    if tmp_node.op_type == "Add":
        mask_node = tmp_node
        tmp_node = om.get_from_node(mask_node.inputs[0])
        if tmp_node.op_type == "Div":
            tmp_node = om.get_from_node(tmp_node.inputs[0])
        assert tmp_node.op_type == "MatMul"
        mask = mask_node.inputs[1]
        assert new_mask is not None

    tmp_node1 = om.get_from_node(tmp_node.inputs[0])
    if tmp_node1.op_type == "Mul":
        tmp_node1 = om.get_prev_nodes(tmp_node1)[0]
    tmp_node2 = om.get_from_node(tmp_node.inputs[1])
    assert tmp_node1.op_type == tmp_node2.op_type == "Transpose"
    tmp_node1 = om.get_prev_nodes(tmp_node1)[0]
    tmp_node2 = om.get_prev_nodes(tmp_node2)[0]
    assert tmp_node1.op_type == tmp_node2.op_type == "Reshape"
    q, k = tmp_node1.inputs[0], tmp_node2.inputs[0]

    tmp_node = om.get_next_nodes(softmax_node)[0]
    assert tmp_node.op_type == "MatMul"
    tmp_node3 = om.get_from_node(tmp_node.inputs[1])
    if tmp_node3 is not None:
        assert tmp_node3.op_type == "Transpose"
        tmp_node3 = om.get_prev_nodes(tmp_node3)[0]
        assert tmp_node3.op_type == "Reshape"
        v = tmp_node3.inputs[0]
    else:
        v_init = om.get_initializer(tmp_node.inputs[1])
        v_init_value = om.get_initializer_value(tmp_node.inputs[1])
        v_init_value = v_init_value[None, ...].transpose(0, 2, 1, 3)
        B, S, H, D = v_init_value.shape
        v_init_value = np.ascontiguousarray(v_init_value.reshape(B, S, H*D))
        om.set_initializer_value(tmp_node.inputs[1], v_init_value)
        v = v_init.name

    tmp_node = om.get_next_nodes(tmp_node)[0]
    assert tmp_node.op_type == "Transpose"
    tmp_node = om.get_next_nodes(tmp_node)[0]
    assert tmp_node.op_type == "Reshape"
    mha_next_node = om.get_next_nodes(tmp_node)[0]
    if mha_next_node.op_type == "Gemm":
        gemm_next_node = om.get_next_nodes(mha_next_node)[0]
        assert gemm_next_node.op_type == "Reshape"
        reshape_next_node = om.get_next_nodes(gemm_next_node)[0]
        assert reshape_next_node.op_type == "Add"
    else:
        assert mha_next_node.op_type == "MatMul"

    name_prefix = '/'.join(softmax_name.split('/')[:-1])
    mha_name = f"{name_prefix}/MultiHeadAttention"
    mha_node = om.create_node("MultiHeadAttention",
                                mha_name,
                                [q, k, v] if mask is None else [q, k, v, new_mask],
                                [mha_name+'_output_0'],
                                num_heads=num_heads,
                                domain="com.microsoft",
                                index=mha_next_node.index-1)
    mha_next_node.replace_input(mha_next_node.inputs[0], mha_node.outputs[0])

    if mha_next_node.op_type == "Gemm":
        weights = om.get_initializer_value(mha_next_node.inputs[1])
        transB = mha_next_node.attrs["transB"]
        assert transB == 1
        weights = np.ascontiguousarray(weights.transpose(1, 0))
        om.set_initializer_value(mha_next_node.inputs[1], weights)
        new_matmul_name = mha_next_node.name.replace("Gemm", "MatMul(Gemm)")
        new_matmul_node = om.create_node("MatMul",
                                            new_matmul_name,
                                            [mha_node.outputs[0], mha_next_node.inputs[1]],
                                            [new_matmul_name + "_output_0"],
                                            index=mha_next_node.index)
        new_bias_name = mha_next_node.name.replace("Gemm", "Add(Gemm)")
        new_add_node = om.create_node("Add",
                                        new_bias_name,
                                        [new_matmul_node.outputs[0], mha_next_node.inputs[2]],
                                        [new_bias_name + "_output_0"],
                                        index=new_matmul_node.index+1)
        reshape_next_node.replace_input(gemm_next_node.outputs[0], new_add_node.outputs[0])


def optimize_normal_attention(om: ONNXModifier):
    def create_new_attention_mask():
        mask_next_node = om.get_to_nodes("attention_mask")[0]
        cast_node = om.create_node("Cast",
                                   "Cast_for_attention_mask",
                                   ["attention_mask"],
                                   ["Cast_for_attention_mask_output_0"],
                                   to=1,
                                   index=mask_next_node.index)
        reducesum_node = om.create_node("ReduceSum", 
                                        "ReduceSum_for_mask", 
                                        [cast_node.outputs[0]], 
                                        ["ReduceSum_for_mask_output_0"], 
                                        axes=1,
                                        keepdims=0, 
                                        index=cast_node.index+1)
        return reducesum_node.outputs[0]

    # bert
    # for i in range(12):
    #     fuse_one_attention(om, f"/bert/encoder/layer.{i}/attention/self/Softmax", "text_token_mask", num_heads=12)

    new_mask = create_new_attention_mask()
    for i in range(6):
        # /transformer/encoder
        # fuse_one_attention(om, f"/transformer/encoder/text_layers.{i}/self_attn/Softmax", "text_token_mask", num_heads=4)
        # /transformer/decoder
        fuse_one_attention(om, f"/transformer/decoder/layers.{i}/self_attn/Softmax", new_mask, num_heads=8)
        fuse_one_attention(om, f"/transformer/decoder/layers.{i}/ca_text/Softmax", new_mask, num_heads=8)

    om.update_map()


def optimize_backbone_attention(om: ONNXModifier):
    def get_original_mask(mask_name, name_prefix):
        mask_value = om.get_initializer_value(mask_name)
        orig_mask = np.where(mask_value==0, 1, 0).astype(np.bool_)
        orig_mask_init = om.create_initializer(f"{name_prefix}/mask", orig_mask)
        return orig_mask_init.name

    def _fuse_one_attention(softmax_name: str):
        name_prefix = '/'.join(softmax_name.split('/')[:-1])

        softmax_node = om.get_node(softmax_name)
        tmp_node = om.get_prev_nodes(softmax_node)[0]
        pos_bias_init = None
        if tmp_node.op_type == "Reshape":
            tmp_node = om.get_prev_nodes(tmp_node)[0]
            assert tmp_node.op_type == "Add"
            pos_bias_init = om.get_initializer(tmp_node.inputs[1])
            tmp_node = om.get_prev_nodes(tmp_node)[0]
            assert tmp_node.op_type == "Reshape"
            tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Add"
        mask = get_original_mask(tmp_node.inputs[1], name_prefix)

        tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "MatMul"
        qk_matmul = tmp_node
        tmp_node = om.get_from_node(qk_matmul.inputs[0])
        assert tmp_node.op_type == "Mul"
        tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Gather"
        q_gather_node = tmp_node

        tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Transpose"
        tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Reshape"
        reshape_node = tmp_node

        tmp_node = om.get_from_node(qk_matmul.inputs[1])
        assert tmp_node.op_type == "Transpose"
        tmp_node = om.get_prev_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Gather"
        k_gather_node = tmp_node

        tmp_node = om.get_next_nodes(softmax_node)[0]
        assert tmp_node.op_type == "MatMul"
        v_gather_node = om.get_from_node(tmp_node.inputs[1])
        assert v_gather_node.op_type == "Gather"
        tmp_node = om.get_next_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Transpose"
        tmp_node = om.get_next_nodes(tmp_node)[0]
        assert tmp_node.op_type == "Reshape"
        mha_out = tmp_node.outputs[0]

        old_dst_shape = om.get_initializer_value(reshape_node.inputs[1])
        b, s, _, h, d = old_dst_shape
        new_dst_shape = [b, s, _, h*d]
        new_dst_shape_init = om.create_initializer(f"{name_prefix}/qkv_hidden_states_shape", 
                                                   np.array(new_dst_shape, np.int64))
        reshape_node.set_input(1, new_dst_shape_init.name)
        for node in [q_gather_node, k_gather_node, v_gather_node]:
            node.set_input(0, reshape_node.outputs[0])
            node.set_attribute("axis", 2)

        mha_name = f"{name_prefix}/MultiHeadAttention"
        inputs = [q_gather_node.outputs[0], 
                  k_gather_node.outputs[0], 
                  v_gather_node.outputs[0],
                  mask]
        if pos_bias_init is not None:
            inputs.append(pos_bias_init.name)
        mha_node = om.create_node("MultiHeadAttention",
                                    mha_name,
                                    inputs,
                                    [mha_name+'_output_0'],
                                    num_heads=h,
                                    domain="com.microsoft",
                                    index=softmax_node.index)
        mha_next_node = om.get_to_nodes(mha_out)[0]
        mha_next_node.replace_input(mha_out, mha_node.outputs[0])

    num_layers = 4
    for l in range(num_layers):  
        num_blocks = 18 if l == 2 else 2
        for b in range(num_blocks):
            _fuse_one_attention(f"/backbone/backbone.0/layers.{l}/blocks.{b}/attn/softmax/Softmax")


def optimize_bidirect_attention(om: ONNXModifier):
    for i in range(6):
        reduce_max_name = f"/transformer/encoder/fusion_layers.{i}/attn/ReduceMax_1"
        reduce_max_node = om.get_node(reduce_max_name)
        
        # 【安全校验】
        if reduce_max_node is None:
            continue
            
        next_nodes = om.get_next_nodes(reduce_max_node)
        if not next_nodes:
            continue
            
        next_node = next_nodes[0]
        if next_node.op_type != "Sub":
            continue

        name_prefix = '/'.join(reduce_max_name.split('/')[:-1])
        matmul_name = f"{name_prefix}/identity_MatMul"
        matmul_init = om.create_initializer(matmul_name + "_B", 
                                            np.diag(np.array([1] * 1)).astype(np.float32))
        matmul_node = om.create_node(
            "MatMul",
            matmul_name,
            [reduce_max_node.outputs[0], matmul_init.name],
            [f"{matmul_name}_output_0"],
            index = reduce_max_node.index+1
        )
        next_node.set_input(1, matmul_node.outputs[0])

# def main():
#     input_onnx_path = sys.argv[1]
#     output_onnx_path = sys.argv[2]
#     # input_onnx_path = "ground_sim.onnx"
#     # output_onnx_path = "ground_sim_0424_2nd.onnx"

#     om = ONNXModifier(input_onnx_path)
#     optimize_where_ndoes(om)         # 1. 替换where节点
#     optimize_transpose_nodes(om)     # 2. 优化transpose节点
#     optmize_sin_cos_block(om)        # 3. 优化位置编码
#     # om.add_opset_import("com.microsoft", 1)
#     # optimize_normal_attention(om)    # 4. 融合bert、transformer中的mha
#     # optimize_ms_deform_attn(om)      # 5. 融合多尺度可变形注意力
#     # optimize_backbone_attention(om)  # 6. 融合backbone中的注意力
#     optimize_bidirect_attention(om)  # 7. 优化双向注意力
#     om.save(output_onnx_path, save_as_external_data=False)

def main():
    # 假设你的原始模型路径
    input_onnx_path = "../weights/ground_deform_sim.onnx"
    # 优化后的模型输出路径
    output_onnx_path = "../weights_opt/ground_deform_opt.onnx"

    print(f"Loading ONNX model from {input_onnx_path}...")
    om = ONNXModifier(input_onnx_path)
    
    print("1. Optimizing Where nodes (Crucial for FP16 & MIGraphX)...")
    optimize_where_ndoes(om)         
    
    print("2. Optimizing Transpose nodes...")
    optimize_transpose_nodes(om)     
    
    # print("3. Optimizing Sin/Cos positional encoding...")
    # optmize_sin_cos_block(om)        
    
    # print("4. Optimizing Bidirectional attention...")
    # optimize_bidirect_attention(om)  


    print(f"Saving optimized model to {output_onnx_path}...")
    om.save(output_onnx_path, save_as_external_data=False)
    print("Optimization Done!")

if __name__ == "__main__":
    main()
