#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 = 2; // Batch size
    // Input
    const int64_t c = 3; // Number of channels
    const int64_t h = 4; // Height
    const int64_t w = 5; // Width

    auto buildAdamwGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();

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

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

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

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

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

        auto adamwAttributes = hipdnn_frontend::graph::AdamwAttributes()
                                   .set_name("adamw_node")
                                   .set_transformeradamw(false)
                                   .set_max_exp_avg_sqs(maxExpAvgSqs);

        graph->adamw(params, grads, expAvgs, expAvgSqs, adamwAttributes);

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

        return std::make_tuple(graph, params, grads, expAvgs, expAvgSqs, maxExpAvgSqs);
    };

    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, params, grads, expAvgs, expAvgSqs, maxExpAvgSqs] = buildAdamwGraph(handle);

    // Allocate DCU memory
    hipdnn_data_sdk::utilities::Tensor<InputType> paramsTensor(params->get_dim(),
                                                               params->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> gradsTensor(grads->get_dim(),
                                                              grads->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> expAvgsTensor(expAvgs->get_dim(),
                                                                expAvgs->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> expAvgSqsTensor(expAvgSqs->get_dim(),
                                                                  expAvgSqs->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> maxExpAvgSqsTensor(maxExpAvgSqs->get_dim(),
                                                                     maxExpAvgSqs->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[params->get_uid()] = paramsTensor.memory().deviceData();
    variantPack[grads->get_uid()] = gradsTensor.memory().deviceData();
    variantPack[expAvgs->get_uid()] = expAvgsTensor.memory().deviceData();
    variantPack[expAvgSqs->get_uid()] = expAvgSqsTensor.memory().deviceData();
    variantPack[maxExpAvgSqs->get_uid()] = maxExpAvgSqsTensor.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 << "Adamw graph execution complete. \n";

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