import hipdnn
import torch


def build_scale_bias_graph(
    hipdnn_handle, torch_tensor_x, torch_tensor_scale, torch_tensor_bias, hipdnn_data_type
):
    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="scale_bias",
    )

    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_scale_out = graph.mul(a=hipdnn_tensor_x, b=hipdnn_tensor_scale, name="scale")
    hipdnn_tensor_y = graph.add(a=hipdnn_tensor_scale_out, b=hipdnn_tensor_bias, name="bias")
    hipdnn_tensor_y.set_output(True)

    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_x, hipdnn_tensor_scale, hipdnn_tensor_bias, hipdnn_tensor_y)


if __name__ == "__main__":
    n = 1
    c = 4
    h = 32
    w = 32

    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").to(
        memory_format=torch.channels_last
    )
    torch_tensor_scale = torch.rand(1, c, 1, 1, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    torch_tensor_bias = torch.rand(1, c, 1, 1, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_x, hipdnn_tensor_scale, hipdnn_tensor_bias, hipdnn_tensor_y = (
        build_scale_bias_graph(
            hipdnn_handle,
            torch_tensor_x,
            torch_tensor_scale,
            torch_tensor_bias,
            hipdnn_data_type,
        )
    )

    torch_tensor_y = torch.empty(
        hipdnn_tensor_y.get_dim(), dtype=torch_data_type, device="cuda"
    ).to(memory_format=torch.channels_last)
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_scale: torch_tensor_scale.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("scale_bias graph execution complete.")