import hipdnn
import torch


def build_conv_genstats_graph(
    hipdnn_handle, torch_tensor_x, torch_tensor_w, padding, stride, dilation, hipdnn_data_type
):
    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="conv_genstats",
    )

    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_w = graph.tensor_like(torch_tensor_w)

    hipdnn_tensor_y = graph.conv_fprop(
        image=hipdnn_tensor_x,
        weight=hipdnn_tensor_w,
        padding=padding,
        stride=stride,
        dilation=dilation,
        name="conv",
    )
    hipdnn_tensor_y.set_output(True)

    hipdnn_tensor_sum, hipdnn_tensor_sq_sum = graph.genstats(
        hipdnn_tensor_y, hipdnn.data_type.FLOAT, name="genstats"
    )
    hipdnn_tensor_sum.set_output(True)
    hipdnn_tensor_sq_sum.set_output(True)

    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_w,
        hipdnn_tensor_y,
        hipdnn_tensor_sum,
        hipdnn_tensor_sq_sum,
    )


if __name__ == "__main__":
    n = 4
    c = 64
    h = 16
    w = 16
    k = 32
    r = 3
    s = 3

    stride_h = 1
    stride_w = 1
    pad_h = 1
    pad_w = 1
    dil_h = 1
    dil_w = 1

    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32

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

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_w,
        hipdnn_tensor_y,
        hipdnn_tensor_sum,
        hipdnn_tensor_sq_sum,
    ) = build_conv_genstats_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_w,
        [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,
        memory_format=torch.channels_last,
        device="cuda",
    )
    torch_tensor_sum = torch.empty(
        hipdnn_tensor_sum.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_sq_sum = torch.empty(
        hipdnn_tensor_sq_sum.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_y: torch_tensor_y.data_ptr(),
        hipdnn_tensor_sum: torch_tensor_sum.data_ptr(),
        hipdnn_tensor_sq_sum: torch_tensor_sq_sum.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("conv_genstats graph execution complete.")