import hipdnn
import torch


def build_binary_pointwise_graph(
    hipdnn_handle,
    torch_tensor_in0,
    torch_tensor_in1,
    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="add_graph",
    )

    # Create hipdnn tensors
    hipdnn_tensor_in0 = graph.tensor_like(torch_tensor_in0)
    hipdnn_tensor_in1 = graph.tensor_like(torch_tensor_in1)

    # Using the add op as an example, other binary pointwise ops can be used similarly.
    # Create binary pointwise ADD op
    hipdnn_tensor_out = graph.add(
        hipdnn_tensor_in0,
        hipdnn_tensor_in1,
        hipdnn.data_type.FLOAT,
        "add_node",
    )
    hipdnn_tensor_out.set_output(True)
    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_in0, hipdnn_tensor_in1, hipdnn_tensor_out)


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

    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")
    torch_tensor_b = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_in0, hipdnn_tensor_in1, hipdnn_tensor_out = build_binary_pointwise_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_b,
        hipdnn_data_type,
    )

    torch_tensor_y = torch.empty(hipdnn_tensor_out.get_dim(), dtype=torch_data_type, device="cuda")
    variant_pack = {
        hipdnn_tensor_in0: torch_tensor_x.data_ptr(),
        hipdnn_tensor_in1: torch_tensor_b.data_ptr(),
        hipdnn_tensor_out: 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("Binary pointwise ADD graph execution complete.")