import hipdnn
import torch


def build_genstats_graph(hipdnn_handle, torch_tensor_x, 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="genstats_graph",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)

    # Create op
    hipdnn_tensor_sum, hipdnn_tensor_sq_sum = graph.genstats(
        hipdnn_tensor_x, 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_sum, hipdnn_tensor_sq_sum)


if __name__ == "__main__":
    # Input dimensions
    n = 2  # Batch size
    c = 3  # Number of input channels
    h = 4  # Height
    w = 5  # 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")

    # hipdnn_handle = hipdnn.create_handle()

    # graph, hipdnn_tensor_x, hipdnn_tensor_sum, hipdnn_tensor_sq_sum = build_genstats_graph(hipdnn_handle,torch_tensor_x,hipdnn_data_type)

    # 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_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("genstats graph execution complete.")