import hipdnn
import torch


def build_prelu_backward_graph(
    hipdnn_handle,
    torch_tensor_x,
    torch_tensor_dy,
    torch_tensor_weight,
    negative_slope,
    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="convolution_forward",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_dy = graph.tensor_like(torch_tensor_dy)
    hipdnn_tensor_weight = graph.tensor_like(torch_tensor_weight)

    # Create prelu op
    hipdnn_tensor_dx, hipdnn_tensor_dweight = graph.prelu_backward(
        input=hipdnn_tensor_x,
        weight=hipdnn_tensor_weight,
        loss=hipdnn_tensor_dy,
        negative_slope=negative_slope,
        name="prelu_backward",
    )
    hipdnn_tensor_dx.set_output(True)
    hipdnn_tensor_dweight.set_output(True)
    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_dy,
        hipdnn_tensor_weight,
        hipdnn_tensor_dx,
        hipdnn_tensor_dweight,
    )


if __name__ == "__main__":
    # Input dimensions
    batch, channels, height, width = 128, 64, 112, 112

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

    torch_tensor_x = torch.rand(
        batch, channels, height, width, dtype=torch_data_type, device="cuda"
    )
    torch_tensor_dy = torch.rand(
        batch, channels, height, width, dtype=torch_data_type, device="cuda"
    )
    torch_tensor_weight = torch.rand(channels, dtype=torch_data_type, device="cuda")
    negative_slope = 0.1

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_dy,
        hipdnn_tensor_weight,
        hipdnn_tensor_dx,
        hipdnn_tensor_dweight,
    ) = build_prelu_backward_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_dy,
        torch_tensor_weight,
        negative_slope,
        hipdnn_data_type,
    )

    torch_tensor_dx = torch.empty(hipdnn_tensor_dx.get_dim(), dtype=torch_data_type, device="cuda")
    torch_tensor_dweight = torch.empty(
        hipdnn_tensor_dweight.get_dim(), dtype=torch_data_type, device="cuda"
    )
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_dy: torch_tensor_dy.data_ptr(),
        hipdnn_tensor_weight: torch_tensor_weight.data_ptr(),
        hipdnn_tensor_dx: torch_tensor_dx.data_ptr(),
        hipdnn_tensor_dweight: torch_tensor_dweight.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("Prelu backward graph execution complete.")