import hipdnn
import torch


def build_rng_graph(
    hipdnn_handle,
    torch_tensor_seed,
    torch_tensor_offset,
    rng_distribution,
    dim,
    stride,
    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="rng",
    )

    # Create hipdnn tensors
    hipdnn_tensor_seed = graph.tensor_like(torch_tensor_seed)
    hipdnn_tensor_offset = graph.tensor_like(torch_tensor_offset)

    # Create rng op
    hipdnn_tensor_y = graph.rng(
        seed=hipdnn_tensor_seed,
        offset=hipdnn_tensor_offset,
        rng_distribution=rng_distribution,
        dim=dim,
        stride=stride,
        name="rng",
    )
    hipdnn_tensor_y.set_output(True)
    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_seed, hipdnn_tensor_offset, hipdnn_tensor_y)


if __name__ == "__main__":
    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32
    rng_distribution = hipdnn.rng_distribution.UNIFORM
    dim = [2, 2]
    stride = [1, 1]

    torch_tensor_seed = torch.rand(1, dtype=torch_data_type, device="cpu")
    torch_tensor_offset = torch.rand(1, dtype=torch_data_type, device="cpu")

    hipdnn_handle = hipdnn.create_handle()

    # graph, hipdnn_tensor_seed, hipdnn_tensor_offset, hipdnn_tensor_y = build_rng_graph(
    #     hipdnn_handle, torch_tensor_seed, torch_tensor_offset, rng_distribution, dim, stride, hipdnn_data_type
    # )

    # torch_tensor_y = torch.empty(hipdnn_tensor_y.get_dim(), dtype=torch_data_type, device="cuda")
    # variant_pack = {
    #     hipdnn_tensor_seed: torch_tensor_seed.data_ptr(),
    #     hipdnn_tensor_offset: torch_tensor_offset.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("Rng graph execution complete.")