import hipdnn
import torch


def build_convint8_bias_graph(
    hipdnn_handle,
    torch_tensor_x,
    torch_tensor_w,
    torch_tensor_zero_point_dq,
    torch_tensor_scale_dq,
    torch_tensor_bias,
    torch_tensor_zero_point_q,
    torch_tensor_scale_q,
    padding,
    stride,
    dilation,
    hipdnn_data_type,
):
    # Create graph
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn_data_type,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="convint8_bias",
    )

    # Create hipdnn conv input and filter tensor with NCHWc32 layout
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x).set_vector_count_and_dimension(32, 1)
    hipdnn_tensor_w = graph.tensor_like(torch_tensor_w).set_vector_count_and_dimension(32, 1)

    # Create conv_fprop op
    hipdnn_tensor_conv_output = graph.conv_fprop(
        image=hipdnn_tensor_x,
        weight=hipdnn_tensor_w,
        padding=padding,
        stride=stride,
        dilation=dilation,
        name="conv_fprop_node",
    )

    # Create sub node for dequantize:zero_point_dq
    hipdnn_tensor_zero_point_dq = graph.tensor_like(torch_tensor_zero_point_dq)
    hipdnn_tensor_zero_point_dq.set_value(0.0)

    hipdnn_tensor_conv_deq_sub_output = graph.sub(
        a=hipdnn_tensor_conv_output, b=hipdnn_tensor_zero_point_dq, name="conv_deq_sub_node"
    )

    # Create mul node for dequantize:scale_dq
    hipdnn_tensor_scale_dq = graph.tensor_like(torch_tensor_scale_dq)
    hipdnn_tensor_scale_dq.set_value(1.0)

    hipdnn_tensor_conv_deq_mul_output = graph.mul(
        a=hipdnn_tensor_conv_deq_sub_output, b=hipdnn_tensor_scale_dq, name="conv_deq_mul_node"
    )

    # Create bias node
    hipdnn_tensor_bias = graph.tensor_like(torch_tensor_bias)
    hipdnn_tensor_bias_output = graph.add(
        a=hipdnn_tensor_conv_deq_mul_output, b=hipdnn_tensor_bias, name="bias_node"
    )

    # Create div node for quantize:scale_q
    hipdnn_tensor_scale_q = graph.tensor_like(torch_tensor_scale_q)
    hipdnn_tensor_scale_q.set_value(1.0)
    hipdnn_tensor_quantize_div_output = graph.div(
        a=hipdnn_tensor_bias_output, b=hipdnn_tensor_scale_q, name="quantize_div_node"
    )

    # Create add node for quantize:zero_point_q
    hipdnn_tensor_zero_point_q = graph.tensor_like(torch_tensor_zero_point_q)
    hipdnn_tensor_zero_point_q.set_value(0.0)
    hipdnn_tensor_output = graph.add(
        a=hipdnn_tensor_quantize_div_output, b=hipdnn_tensor_zero_point_q, name="quantize_add_node"
    )

    hipdnn_tensor_output.set_output(True).set_vector_count_and_dimension(32, 1)

    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_w,
        hipdnn_tensor_bias,
        hipdnn_tensor_output,
    )


if __name__ == "__main__":
    # Input dimensions
    n = 2  # Batch size
    c = 32  # Number of input channels
    h = 16  # Height
    w = 8  # Width

    # Filter dimensions
    k = 128  # Number of output channels
    r = 3  # Filter height
    s = 3  # Filter width

    # Convolution parameters
    stride_h = 1  # Height stride
    stride_w = 1  # Width stride
    pad_h = 0  # Height padding
    pad_w = 0  # Width padding
    dil_h = 1  # Height dilation
    dil_w = 1  # Width dilation

    hipdnn_data_type = hipdnn.data_type.INT8
    torch_data_type = torch.int8
    bias_data_type = torch.float32
    quantize_data_type = torch.float32

    torch_tensor_x = torch.randint(
        low=-128,
        high=128,
        size=(n, c, h, w),
        dtype=torch_data_type,
        device="cuda",
    )
    torch_tensor_w = torch.randint(
        low=-128, high=128, size=(k, c, r, s), dtype=torch_data_type, device="cuda"
    )

    torch_tensor_zero_point_dq = torch.rand(1, 1, 1, 1, device="cuda")
    torch_tensor_scale_dq = torch.rand(1, 1, 1, 1, device="cuda")

    torch_tensor_bias = torch.rand(1, k, 1, 1, dtype=bias_data_type, device="cuda")

    torch_tensor_zero_point_q = torch.rand(1, 1, 1, 1, device="cuda")
    torch_tensor_scale_q = torch.rand(1, 1, 1, 1, device="cuda")

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_w,
        hipdnn_tensor_bias,
        hipdnn_tensor_y,
    ) = build_convint8_bias_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_w,
        torch_tensor_zero_point_dq,
        torch_tensor_scale_dq,
        torch_tensor_bias,
        torch_tensor_zero_point_q,
        torch_tensor_scale_q,
        [pad_h, pad_w],
        [stride_h, stride_w],
        [dil_h, dil_w],
        hipdnn_data_type,
    )

    torch_tensor_y = torch.empty(
        hipdnn_tensor_y.get_dim(),
        dtype=torch_data_type,
        device="cuda",
    )
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_w: torch_tensor_w.data_ptr(),
        hipdnn_tensor_bias: torch_tensor_bias.data_ptr(),
        hipdnn_tensor_y: torch_tensor_y.data_ptr(),
    }
    workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda")

    graph.exec(variant_pack=variant_pack, workspace=workspace.data_ptr())
    print("convint8_bias graph execution complete.")