#include <iostream>

#include "utils.hpp"

#include <hipdnn_data_sdk/utilities/Tensor.hpp>
#include <hipdnn_data_sdk/utilities/Workspace.hpp>
#include <hipdnn_frontend.hpp>

int main()
{
    using InputType = hipdnn_data_sdk::types::half;

    const int64_t n = 16; // Batch size
    // Input
    const int64_t c = 16; // Number of channels
    const int64_t h = 16; // Height
    const int64_t w = 16; // Width

    auto buildBnBackwardGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();
        graph->set_name("bn_backward_graph")
            .set_io_data_type(hipdnn_frontend::getDataTypeEnumFromType<InputType>())
            .set_intermediate_data_type(hipdnn_frontend::getDataTypeEnumFromType<InputType>())
            .set_compute_data_type(hipdnn_frontend::DataType::FLOAT);

        auto dy = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("dy")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));

        auto x = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("x")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));

        auto scale = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("scale")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));

        auto savedMean = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("save_mean")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));

        auto savedInvVariance = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("save_inv_variance")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));

        auto bnBwdAttributes = hipdnn_frontend::graph::BatchnormBackwardAttributes()
                                   .set_name("bn_backward_node")
                                   .set_saved_mean_and_inv_variance(savedMean, savedInvVariance);

        auto [dx, dscale, dbias] = graph->batchnorm_backward(dy, x, scale, bnBwdAttributes);
        dx->set_output(true);
        dscale->set_output(true);
        dbias->set_output(true);

        // build graph
        HIPDNN_FE_CHECK(graph->build(handle));

        return std::make_tuple(graph, dy, x, scale, savedMean, savedInvVariance, dx, dscale, dbias);
    };

    auto backend = hipdnn_frontend::detail::hipdnnBackend();
    if(!backend)
    {
        std::cout << "Creat backend failed. \n";
        return 1;
    }

    hipdnnHandle_t handle;
    HIPDNN_CHECK(backend->create(&handle));

    auto [graph, dy, x, scale, savedMean, savedInvVariance, dx, dscale, dbias]
        = buildBnBackwardGraph(handle);

    // Allocate DCU memory
    hipdnn_data_sdk::utilities::Tensor<InputType> dyTensor(dy->get_dim(), dy->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> xTensor(x->get_dim(), x->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> scaleTensor(scale->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> savedMeanTensor(savedMean->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> savedInvVarTensor(savedInvVariance->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> dxTensor(dx->get_dim(), dx->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dscaleTensor(dscale->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> dbiasTensor(dbias->get_dim());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[dy->get_uid()] = dyTensor.memory().deviceData();
    variantPack[x->get_uid()] = xTensor.memory().deviceData();
    variantPack[scale->get_uid()] = scaleTensor.memory().deviceData();
    variantPack[savedMean->get_uid()] = savedMeanTensor.memory().deviceData();
    variantPack[savedInvVariance->get_uid()] = savedInvVarTensor.memory().deviceData();
    variantPack[dx->get_uid()] = dxTensor.memory().deviceData();
    variantPack[dscale->get_uid()] = dscaleTensor.memory().deviceData();
    variantPack[dbias->get_uid()] = dbiasTensor.memory().deviceData();

    int64_t workspaceSize = 0;
    HIPDNN_FE_CHECK(graph->get_workspace_size(workspaceSize));
    const hipdnn_data_sdk::utilities::Workspace workspace(static_cast<size_t>(workspaceSize));

    HIPDNN_FE_CHECK(graph->execute(handle, variantPack, workspace.get()));

    std::cout << "Batch normalization backward graph execution complete. \n";

    HIPDNN_CHECK(backend->destroy(handle));
    return 0;
}