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

    const int64_t dyN0 = 64;
    const int64_t dyN1 = 32;

    const int64_t dxN0 = 128;
    const int64_t dxN1 = 64;

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

        graph->set_name("getitem_backward_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 dy = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("dy")
                .set_dim({dyN0, dyN1})
                .set_stride({dyN1, 1}));

        auto xIndeices = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("x_indeices")
                .set_dim({64, 32})
                .set_stride({32, 1})
                .set_data_type(hipdnn_frontend::DataType::INT32));
        auto yIndeices = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("y_indeices")
                .set_dim({64, 32})
                .set_stride({32, 1})
                .set_data_type(hipdnn_frontend::DataType::INT32));

        auto getitemBackwardAttributes = hipdnn_frontend::graph::GetitemBackwardAttributes()
                                             .set_dims({0, 1})
                                             .set_indices({xIndeices, yIndeices})
                                             .set_offset(0)
                                             .set_name("getitem_backward");

        auto [dx, error] = graph->getitem_backward(dy, getitemBackwardAttributes);
        dx->set_output(true).set_dim({dxN0, dxN1}).set_stride({dxN1, 1});
        error->set_output(true);

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

        return std::make_tuple(graph, dy, xIndeices, yIndeices, dx, error);
    };

    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, dy, xindex, yindex, dx, error] = buildGetitemBackwardGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> dyTensor(dy->get_dim(), dy->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dxTensor(dx->get_dim(), dx->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> xIndexTensor(xindex->get_dim(),
                                                               xindex->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> yIndexTensor(yindex->get_dim(),
                                                               yindex->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> errorTensor(error->get_dim(),
                                                              error->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[dy->get_uid()] = dyTensor.memory().deviceData();
    variantPack[dx->get_uid()] = dxTensor.memory().deviceData();
    variantPack[xindex->get_uid()] = xIndexTensor.memory().deviceData();
    variantPack[yindex->get_uid()] = yIndexTensor.memory().deviceData();
    variantPack[error->get_uid()] = errorTensor.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 << "GetitemBackwardIndices graph execution complete.\n";

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