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

# ==================================================================================================

base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/"
pose_model_path = base_path + "rtmpose-m_384x288.onnx"
det_model_path = base_path + "rtmdet-nano_320x320.onnx"

norm_mean = -1 * np.array([103.53, 116.28, 123.675])
norm_std = 1.0 / np.array([57.375, 57.12, 58.395])


# ==================================================================================================


def add_steps_to_onnx(model_path):

    # Load existing model
    model = onnx.load(model_path)
    graph = model.graph

    mean = norm_mean.astype(np.float32)
    std = norm_std.astype(np.float32)

    use_bgr = bool("rtmpose" in model_path)
    if use_bgr:
        mean = mean[::-1]
        std = std[::-1]

    mean = np.reshape(mean, (1, 3, 1, 1)).astype(np.float32)
    std = np.reshape(std, (1, 3, 1, 1)).astype(np.float32)

    use_fp16 = bool("fp16" in model_path)
    if use_fp16:
        mean = mean.astype(np.float16)
        std = std.astype(np.float16)

    # Add the initializers to the graph
    mean_initializer = numpy_helper.from_array(mean, name="norm_mean")
    std_initializer = numpy_helper.from_array(std, name="norm_std")
    graph.initializer.extend([mean_initializer, std_initializer])

    # Define layer names, assuming the first input is the image tensor
    input_name = graph.input[0].name
    mean_added_output = "mean_added_output"
    normalized_output = "normalized_output"

    # Node to add mean
    mean_add_node = helper.make_node(
        "Add",
        inputs=[input_name, "norm_mean"],
        outputs=[mean_added_output],
        name="Mean_Addition",
    )

    # Node to multiply by std
    std_mul_node = helper.make_node(
        "Mul",
        inputs=[mean_added_output, "norm_std"],
        outputs=[normalized_output],
        name="Std_Multiplication",
    )

    # Replace original input of the model with the output of normalization
    for node in graph.node:
        for idx, input_name_in_node in enumerate(node.input):
            if input_name_in_node == input_name:
                node.input[idx] = normalized_output

    # Add the new nodes to the graph
    graph.node.insert(0, mean_add_node)
    graph.node.insert(1, std_mul_node)

    path = model_path.replace(".onnx", "_with-norm.onnx")
    onnx.save(model, path)


def main():
    add_steps_to_onnx(pose_model_path)
    add_steps_to_onnx(det_model_path)
    add_steps_to_onnx(det_model_path.replace(".onnx", "_fp16.onnx"))
    add_steps_to_onnx(pose_model_path.replace(".onnx", "_fp16.onnx"))


# ==================================================================================================

if __name__ == "__main__":
    main()