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

    const int64_t n0 = 700; // time_step
    const int64_t n1 = 4; // batch_size
    const int64_t n2 = 10; // class_number

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

        graph->set_name("ctc_loss_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 probs = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("probs")
                .set_dim({n0, n1, n2})
                .set_stride({n1 * n2, n1, 1}));

        auto ctcLossAttributes = hipdnn_frontend::graph::CtcLossAttributes()
                                     .set_blank_label_id(0)
                                     .set_apply_softmax(true)
                                     .set_algo(0)
                                     .set_labels({1, 2, 3, 3, 3, 5, 6})
                                     .set_label_lengths({1, 1, 2, 3})
                                     .set_input_lengths({500, 500, 600})
                                     .set_name("ctc_loss");

        auto [losses, gradients] = graph->ctc_loss(probs, ctcLossAttributes);
        losses->set_output(true);
        gradients->set_output(true);

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

        return std::make_tuple(graph, probs, losses, gradients);
    };

    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, probs, losses, gradients] = buildCtcLossGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> probsTensor(probs->get_dim(),
                                                              probs->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> lossesTensor(losses->get_dim(),
                                                               losses->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> gradientsTensor(gradients->get_dim(),
                                                                  gradients->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[probs->get_uid()] = probsTensor.memory().deviceData();
    variantPack[losses->get_uid()] = lossesTensor.memory().deviceData();
    variantPack[gradients->get_uid()] = gradientsTensor.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 << "CtcLoss graph execution complete.\n";

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