import hipdnn
import torch


def build_deform_convolution_graph(
    hipdnn_handle,
    torch_tensor_dy,
    torch_tensor_x,
    torch_tensor_w,
    torch_tensor_offset,
    torch_tensor_mask,
    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.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="deform_convolution",
    )

    # Create hipdnn tensors
    hipdnn_tensor_dy = graph.tensor_like(torch_tensor_dy)
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_w = graph.tensor_like(torch_tensor_w)
    hipdnn_tensor_offset = graph.tensor_like(torch_tensor_offset)
    hipdnn_tensor_mask = graph.tensor_like(torch_tensor_mask)

    # Create op
    hipdnn_tensor_dx, hipdnn_tensor_doffset, hipdnn_tensor_dmask = graph.deform_conv_dgrad(
        loss=hipdnn_tensor_dy,
        filter=hipdnn_tensor_w,
        offset=hipdnn_tensor_offset,
        image=hipdnn_tensor_x,
        mask=hipdnn_tensor_mask,
        padding=padding,
        stride=stride,
        dilation=dilation,
        name="deform_conv_bwd",
    )
    hipdnn_tensor_dx.set_output(True)
    hipdnn_tensor_doffset.set_output(True)
    hipdnn_tensor_dmask.set_output(True)
    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_dy,
        hipdnn_tensor_w,
        hipdnn_tensor_offset,
        hipdnn_tensor_x,
        hipdnn_tensor_mask,
        hipdnn_tensor_dx,
        hipdnn_tensor_doffset,
        hipdnn_tensor_dmask,
    )


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

    # Filter dimensions
    k = 1  # 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

    h_out = int((h + 2 * pad_h - (dil_h * (r - 1) + 1)) / stride_h + 1)
    w_out = int((w + 2 * pad_w - (dil_w * (s - 1) + 1)) / stride_w + 1)

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

    torch_tensor_dy = torch.rand(n, k, h_out, w_out, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    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
    )
    torch_tensor_offset = torch.rand(
        n, 2 * r * s, h_out, w_out, dtype=torch_data_type, device="cuda"
    ).to(memory_format=torch.channels_last)
    torch_tensor_mask = torch.rand(n, r * s, h_out, w_out, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_dy,
        hipdnn_tensor_w,
        hipdnn_tensor_offset,
        hipdnn_tensor_x,
        hipdnn_tensor_mask,
        hipdnn_tensor_dx,
        hipdnn_tensor_doffset,
        hipdnn_tensor_dmask,
    ) = build_deform_convolution_graph(
        hipdnn_handle,
        torch_tensor_dy,
        torch_tensor_x,
        torch_tensor_w,
        torch_tensor_offset,
        torch_tensor_mask,
        [pad_h, pad_w],
        [stride_h, stride_w],
        [dil_h, dil_w],
        hipdnn_data_type,
    )

    torch_tensor_dx = torch.empty(
        hipdnn_tensor_dx.get_dim(),
        dtype=torch_data_type,
        memory_format=torch.channels_last,
        device="cuda",
    )
    torch_tensor_doffset = torch.empty(
        hipdnn_tensor_doffset.get_dim(),
        dtype=torch_data_type,
        memory_format=torch.channels_last,
        device="cuda",
    )
    torch_tensor_dmask = torch.empty(
        hipdnn_tensor_dmask.get_dim(),
        dtype=torch_data_type,
        memory_format=torch.channels_last,
        device="cuda",
    )
    variant_pack = {
        hipdnn_tensor_dy: torch_tensor_dy.data_ptr(),
        hipdnn_tensor_w: torch_tensor_w.data_ptr(),
        hipdnn_tensor_offset: torch_tensor_offset.data_ptr(),
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_mask: torch_tensor_mask.data_ptr(),
        hipdnn_tensor_dx: torch_tensor_dx.data_ptr(),
        hipdnn_tensor_doffset: torch_tensor_doffset.data_ptr(),
        hipdnn_tensor_dmask: torch_tensor_dmask.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("deform conv bwd graph execution complete.")