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

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

    // params
    const int64_t n = 1; // Batch size
    const int64_t c = 32; // Number of channels
    const int64_t h = 128; // Height
    const int64_t w = 128; // Width
    const int64_t k = 32; // Number of filters
    const int64_t r = 2; // Height
    const int64_t s = 2; // Width

    const int64_t axis = 1;

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

        const auto inputType = hipdnn_frontend::getDataTypeEnumFromType<InputType>();
        graph->set_name("concat_conv_graph")
            .set_io_data_type(inputType)
            .set_intermediate_data_type(inputType)
            .set_compute_data_type(hipdnn_frontend::DataType::FLOAT);

        // create concat
        auto x1 = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("x1")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c})
                .set_data_type(inputType));
        auto x2 = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("x2")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c})
                .set_data_type(inputType));
        auto concatenateAttributes = hipdnn_frontend::graph::ConcatenateAttributes().set_axis(axis);
        auto concatOutput = graph->concatenate({x1, x2}, concatenateAttributes);

        // create conv
        const int64_t c2 = c * 2;
        auto filter = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("filter")
                .set_dim({k, c2, r, s})
                .set_stride({c2 * r * s, 1, c2 * s, c2}));
        auto convFpropAttributes = hipdnn_frontend::graph::ConvFpropAttributes()
                                       .set_name("conv_fprop_node")
                                       .set_padding({1, 1})
                                       .set_stride({1, 1})
                                       .set_dilation({1, 1});
        auto y = graph->conv_fprop(concatOutput, filter, convFpropAttributes);
        y->set_output(true);

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

        return std::make_tuple(graph, x1, x2, filter, y);
    };

    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, x1, x2, filter, y] = buildConcatConvGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> x1Tensor(x1->get_dim(), x1->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> x2Tensor(x2->get_dim(), x2->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> filterTensor(filter->get_dim(),
                                                               filter->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> yTensor(y->get_dim(), y->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[x1->get_uid()] = x1Tensor.memory().deviceData();
    variantPack[x2->get_uid()] = x2Tensor.memory().deviceData();
    variantPack[filter->get_uid()] = filterTensor.memory().deviceData();
    variantPack[y->get_uid()] = yTensor.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 << "Concatenate graph execution complete. \n";

    HIPDNN_CHECK(backend->destroy(handle));

    return 0;
}