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

#include <iostream>
#include <memory>
#include <tuple>
#include <unordered_map>

#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;

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

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

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

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

        auto pointwiseAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                       .set_name("add_node")
                                       .set_mode(hipdnn_frontend::PointwiseMode::ADD);

        auto out = graph->pointwise(in0, in1, pointwiseAttributes);
        out->set_output(true);

        HIPDNN_FE_CHECK(graph->build(handle));

        return std::make_tuple(graph, in0, in1, out);
    };

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

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

    auto [graph, in0, in1, out] = buildBinaryPointwiseGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> in0Tensor(in0->get_dim(), in0->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> in1Tensor(in1->get_dim(), in1->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> outTensor(out->get_dim(), out->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[in0->get_uid()] = in0Tensor.memory().deviceData();
    variantPack[in1->get_uid()] = in1Tensor.memory().deviceData();
    variantPack[out->get_uid()] = outTensor.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 << "Binary pointwise ADD sample execution complete. \n";

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