import hipdnn
import torch


def build_concat_conv_graph(
    hipdnn_handle,
    torch_tensor_x1,
    torch_tensor_x2,
    torch_tensor_w,
    padding,
    stride,
    dilation,
    hipdnn_data_type,
    concat_axis,
):
    # Create graph
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn.data_type.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="concat_conv",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x1 = graph.tensor_like(torch_tensor_x1)
    hipdnn_tensor_x2 = graph.tensor_like(torch_tensor_x2)
    hipdnn_tensor_w = graph.tensor_like(torch_tensor_w)

    # Create concatenate op
    hipdnn_tensor_concat_output = graph.concatenate(
        x=[hipdnn_tensor_x1, hipdnn_tensor_x2], axis=concat_axis, name="concatenate"
    )

    # Create conv op
    hipdnn_tensor_y = graph.conv_fprop(
        image=hipdnn_tensor_concat_output,
        weight=hipdnn_tensor_w,
        padding=padding,
        stride=stride,
        dilation=dilation,
        name="conv2d",
    )
    hipdnn_tensor_y.set_output(True)
    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_x1, hipdnn_tensor_x2, hipdnn_tensor_w, hipdnn_tensor_y)


if __name__ == "__main__":
    # Input dimensions
    n = 1
    c = 32
    h = 128
    w = 128

    # Filter dimensions
    k = 32
    r = 2
    s = 2

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

    hipdnn_data_type = hipdnn.data_type.HALF
    torch_data_type = torch.float16
    concat_axis = 1

    torch_tensor_x1 = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    torch_tensor_x2 = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    torch_tensor_w = torch.rand(k, 2 * c, r, s, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_x1, hipdnn_tensor_x2, hipdnn_tensor_w, hipdnn_tensor_y = (
        build_concat_conv_graph(
            hipdnn_handle,
            torch_tensor_x1,
            torch_tensor_x2,
            torch_tensor_w,
            [pad_h, pad_w],
            [stride_h, stride_w],
            [dil_h, dil_w],
            hipdnn_data_type,
            concat_axis,
        )
    )

    torch_tensor_y = torch.empty(hipdnn_tensor_y.get_dim(), dtype=torch_data_type, device="cuda")
    variant_pack = {
        hipdnn_tensor_x1: torch_tensor_x1.data_ptr(),
        hipdnn_tensor_x2: torch_tensor_x2.data_ptr(),
        hipdnn_tensor_w: torch_tensor_w.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("Concat_conv graph execution complete.")