import hipdnn
import torch


def build_batchnorm_training_graph(
    hipdnn_handle,
    torch_tensor_x,
    torch_tensor_scale,
    torch_tensor_bias,
    torch_tensor_prev_running_mean,
    torch_tensor_prev_running_var,
    torch_tensor_momentum,
    torch_tensor_epsilon,
    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="batchnorm_training",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_scale = graph.tensor_like(torch_tensor_scale)
    hipdnn_tensor_bias = graph.tensor_like(torch_tensor_bias)
    hipdnn_tensor_prev_running_mean = graph.tensor_like(torch_tensor_prev_running_mean)
    hipdnn_tensor_prev_running_var = graph.tensor_like(torch_tensor_prev_running_var)
    hipdnn_tensor_momentum = graph.tensor_like(torch_tensor_momentum)
    hipdnn_tensor_momentum.set_value(0.1)
    hipdnn_tensor_epsilon = graph.tensor_like(torch_tensor_epsilon)
    hipdnn_tensor_epsilon.set_value(1e-5)

    # Create batchnorm op
    (
        hipdnn_tensor_y,
        hipdnn_tensor_saved_mean,
        hipdnn_tensor_saved_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_var,
    ) = graph.batchnorm(
        input=hipdnn_tensor_x,
        scale=hipdnn_tensor_scale,
        bias=hipdnn_tensor_bias,
        in_running_mean=hipdnn_tensor_prev_running_mean,
        in_running_var=hipdnn_tensor_prev_running_var,
        epsilon=hipdnn_tensor_epsilon,
        momentum=hipdnn_tensor_momentum,
    )
    hipdnn_tensor_y.set_output(True)
    hipdnn_tensor_saved_mean.set_output(True)
    hipdnn_tensor_saved_inv_variance.set_output(True)
    hipdnn_tensor_next_running_mean.set_output(True)
    hipdnn_tensor_next_running_var.set_output(True)

    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_scale,
        hipdnn_tensor_bias,
        hipdnn_tensor_prev_running_mean,
        hipdnn_tensor_prev_running_var,
        hipdnn_tensor_epsilon,
        hipdnn_tensor_momentum,
        hipdnn_tensor_y,
        hipdnn_tensor_saved_mean,
        hipdnn_tensor_saved_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_var,
    )


if __name__ == "__main__":
    # Input dimensions
    n = 4  # Batch size
    c = 16  # Number of input channels
    h = 56  # Height
    w = 56  # 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_scale = torch.rand(1, c, 1, 1, dtype=torch.float32, device="cuda")
    torch_tensor_bias = torch.rand(1, c, 1, 1, dtype=torch.float32, device="cuda")
    torch_tensor_prev_running_mean = torch.rand(1, c, 1, 1, dtype=torch.float32, device="cuda")
    torch_tensor_prev_running_var = torch.rand(1, c, 1, 1, dtype=torch.float32, device="cuda")
    torch_tensor_epsilon = torch.full(
        (1, 1, 1, 1), 1e-5, dtype=torch.float32, requires_grad=False, device="cuda"
    )
    torch_tensor_momentum = torch.full(
        (1, 1, 1, 1), 1e-5, dtype=torch.float32, requires_grad=False, device="cuda"
    )

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_scale,
        hipdnn_tensor_bias,
        hipdnn_tensor_prev_running_mean,
        hipdnn_tensor_prev_running_var,
        hipdnn_tensor_epsilon,
        hipdnn_tensor_momentum,
        hipdnn_tensor_y,
        hipdnn_tensor_saved_mean,
        hipdnn_tensor_saved_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_var,
    ) = build_batchnorm_training_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_scale,
        torch_tensor_bias,
        torch_tensor_prev_running_mean,
        torch_tensor_prev_running_var,
        torch_tensor_momentum,
        torch_tensor_epsilon,
        hipdnn_data_type,
    )

    torch_tensor_y = torch.empty(hipdnn_tensor_y.get_dim(), dtype=torch_data_type, device="cuda")
    torch_tensor_saved_mean = torch.empty(
        hipdnn_tensor_saved_mean.get_dim(), dtype=torch.float32, device="cuda"
    )
    torch_tensor_saved_inv_variance = torch.empty(
        hipdnn_tensor_saved_inv_variance.get_dim(), dtype=torch.float32, device="cuda"
    )
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_prev_running_mean: torch_tensor_prev_running_mean.data_ptr(),
        hipdnn_tensor_prev_running_var: torch_tensor_prev_running_var.data_ptr(),
        hipdnn_tensor_scale: torch_tensor_scale.data_ptr(),
        hipdnn_tensor_bias: torch_tensor_bias.data_ptr(),
        hipdnn_tensor_epsilon: torch_tensor_epsilon.data_ptr(),
        hipdnn_tensor_momentum: torch_tensor_momentum.data_ptr(),
        hipdnn_tensor_y: torch_tensor_y.data_ptr(),
        hipdnn_tensor_next_running_mean: torch_tensor_prev_running_mean.data_ptr(),
        hipdnn_tensor_next_running_var: torch_tensor_prev_running_var.data_ptr(),
        hipdnn_tensor_saved_mean: torch_tensor_saved_mean.data_ptr(),
        hipdnn_tensor_saved_inv_variance: torch_tensor_saved_inv_variance.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("BatchNorm training graph execution complete.")