#include <iostream>

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

#include "hipdnn_data_sdk/utilities/Workspace.hpp"
#include "utils.hpp"

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

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

    auto buildBnFinalizeGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();
        graph->set_name("bn_finalize_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 sum = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("sum")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));

        auto sqSum = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("sq_sum")
                .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({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));

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

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

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

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

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

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

        epsilon->set_value(1e-5);
        momentum->set_value(0.001f);
        accumCount->set_value(static_cast<int32_t>(n * h * w));

        auto bnFinalizeAttributes
            = hipdnn_frontend::graph::BatchnormFinalizeAttributes()
                  .set_name("bn_finalize_node")
                  .set_previous_running_stats(prevRunningMean, prevRunningVar, momentum);

        auto [eqScale, eqBias, mean, invVariance, nextRunningMean, nextRunningVar]
            = graph->bn_finalize(
                sum, sqSum, scale, bias, epsilon, accumCount, bnFinalizeAttributes);
        eqScale->set_output(true);
        eqBias->set_output(true);
        mean->set_output(true);
        invVariance->set_output(true);
        nextRunningMean->set_output(true);
        nextRunningVar->set_output(true);

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

        return std::make_tuple(graph,
                               sum,
                               sqSum,
                               scale,
                               bias,
                               epsilon,
                               prevRunningMean,
                               prevRunningVar,
                               momentum,
                               accumCount,
                               eqScale,
                               eqBias,
                               mean,
                               invVariance,
                               nextRunningMean,
                               nextRunningVar);
    };

    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,
          sum,
          sqSum,
          scale,
          bias,
          epsilon,
          prevRunningMean,
          prevRunningVar,
          momentum,
          accumCount,
          eqScale,
          eqBias,
          mean,
          invVariance,
          nextRunningMean,
          nextRunningVar]
        = buildBnFinalizeGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> sumTensor(sum->get_dim(), sum->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> sqSumTensor(sqSum->get_dim(),
                                                              sqSum->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> scaleTensor(scale->get_dim(),
                                                              scale->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> biasTensor(bias->get_dim(), bias->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> prevMeanTensor(prevRunningMean->get_dim(),
                                                                 prevRunningMean->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> prevVarTensor(prevRunningVar->get_dim(),
                                                                prevRunningVar->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> momentumTensor(momentum->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> epsilonTensor(epsilon->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> accumCountTensor(accumCount->get_dim());
    hipdnn_data_sdk::utilities::Tensor<InputType> eqScaleTensor(eqScale->get_dim(),
                                                                eqScale->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> eqBiasTensor(eqBias->get_dim(),
                                                               eqBias->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> nextMeanTensor(nextRunningMean->get_dim(),
                                                                 nextRunningMean->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> nextVarTensor(nextRunningVar->get_dim(),
                                                                nextRunningVar->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> meanTensor(mean->get_dim(), mean->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> invVarTensor(invVariance->get_dim(),
                                                               invVariance->get_stride());

    std::unordered_map<int64_t, void*> variantPack;

    variantPack[sum->get_uid()] = sumTensor.memory().deviceData();
    variantPack[sqSum->get_uid()] = sqSumTensor.memory().deviceData();
    variantPack[scale->get_uid()] = scaleTensor.memory().deviceData();
    variantPack[bias->get_uid()] = biasTensor.memory().deviceData();
    variantPack[prevRunningMean->get_uid()] = prevMeanTensor.memory().deviceData();
    variantPack[prevRunningVar->get_uid()] = prevVarTensor.memory().deviceData();
    variantPack[momentum->get_uid()] = momentumTensor.memory().deviceData();
    variantPack[epsilon->get_uid()] = epsilonTensor.memory().deviceData();
    variantPack[accumCount->get_uid()] = accumCountTensor.memory().deviceData();
    variantPack[eqScale->get_uid()] = eqScaleTensor.memory().deviceData();
    variantPack[eqBias->get_uid()] = eqBiasTensor.memory().deviceData();
    variantPack[nextRunningMean->get_uid()] = nextMeanTensor.memory().deviceData();
    variantPack[nextRunningVar->get_uid()] = nextVarTensor.memory().deviceData();
    variantPack[mean->get_uid()] = meanTensor.memory().deviceData();
    variantPack[invVariance->get_uid()] = invVarTensor.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 finalize graph execution complete. \n";

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