import hipdnn
import torch


def build_matmul_bias_graph(
    hipdnn_handle,
    torch_tensor_inputA,
    torch_tensor_inputB,
    torch_tensor_bias,
    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="matmul_bias",
    )

    # Create hipdnn tensors
    hipdnn_tensor_inputA = graph.tensor_like(torch_tensor_inputA)
    hipdnn_tensor_inputB = graph.tensor_like(torch_tensor_inputB)
    hipdnn_tensor_bias = graph.tensor_like(torch_tensor_bias)

    # Create matmul
    hipdnn_tensor_matmul_output = graph.matmul(
        a=hipdnn_tensor_inputA,
        b=hipdnn_tensor_inputB,
        name="matmul",
    )

    # Create bias
    hipdnn_tensor_bias_out = graph.add(
        a=hipdnn_tensor_matmul_output, b=hipdnn_tensor_bias, name="bias"
    )

    # Create relu
    hipdnn_tensor_y = graph.relu(input=hipdnn_tensor_bias_out, lower_clip=0.0, name="relu")
    hipdnn_tensor_y.set_output(True)

    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_inputA, hipdnn_tensor_inputB, hipdnn_tensor_bias, hipdnn_tensor_y)


if __name__ == "__main__":
    # Input dimensions
    b = 2  # Batch size
    n = 16  # Height
    m = 32  # Width

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

    torch_tensor_inputA = torch.rand(b, n, m, dtype=torch_data_type, device="cuda")
    torch_tensor_inputB = torch.rand(b, m, n, dtype=torch_data_type, device="cuda")
    torch_tensor_bias = torch.rand(1, 1, n, dtype=torch_data_type, device="cuda")

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_inputA, hipdnn_tensor_inputB, hipdnn_tensor_bias, hipdnn_tensor_y = (
        build_matmul_bias_graph(
            hipdnn_handle,
            torch_tensor_inputA,
            torch_tensor_inputB,
            torch_tensor_bias,
            hipdnn_data_type,
        )
    )

    torch_tensor_y = torch.empty(
        hipdnn_tensor_y.get_dim(),
        dtype=torch_data_type,
        device="cuda",
    )
    variant_pack = {
        hipdnn_tensor_inputA: torch_tensor_inputA.data_ptr(),
        hipdnn_tensor_inputB: torch_tensor_inputB.data_ptr(),
        hipdnn_tensor_bias: torch_tensor_bias.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("matmul_bias_relu graph execution complete.")