// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT

#include <iostream>

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

#include "../utils.hpp"

int main()
{
    using InputType = float;

    const int64_t n = 1;
    const int64_t c = 4;
    const int64_t h = 32;
    const int64_t w = 32;

    auto buildMulMulAddAddGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();
        graph->set_name("mul_mul_add_add_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 a = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("a")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, 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 b = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("b")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));
        auto y = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("y")
                .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({1, c, 1, 1})
                .set_stride({c, 1, c, c}));

        auto mulAttrs0 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("mul0_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::MUL);
        auto mulOut0 = graph->pointwise(a, x, mulAttrs0);

        auto mulAttrs1 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("mul1_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::MUL);
        auto mulOut1 = graph->pointwise(b, y, mulAttrs1);

        auto addAttrs0 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("add0_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::ADD);
        auto addOut0 = graph->pointwise(mulOut0, mulOut1, addAttrs0);

        auto addAttrs1 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("add1_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::ADD);
        auto z = graph->pointwise(addOut0, bias, addAttrs1);
        z->set_output(true);

        HIPDNN_FE_CHECK(graph->build(handle));

        return std::make_tuple(graph, a, x, b, y, bias, z);
    };

    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, a, x, b, y, bias, z] = buildMulMulAddAddGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> aTensor(a->get_dim(), a->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> xTensor(x->get_dim(), x->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> bTensor(b->get_dim(), b->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> yTensor(y->get_dim(), y->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> biasTensor(bias->get_dim(), bias->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> zTensor(z->get_dim(), z->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[a->get_uid()] = aTensor.memory().deviceData();
    variantPack[x->get_uid()] = xTensor.memory().deviceData();
    variantPack[b->get_uid()] = bTensor.memory().deviceData();
    variantPack[y->get_uid()] = yTensor.memory().deviceData();
    variantPack[bias->get_uid()] = biasTensor.memory().deviceData();
    variantPack[z->get_uid()] = zTensor.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 << "MulMulAddAdd graph execution complete. \n";

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