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

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

        graph->set_name("rng_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 seed = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("seed").set_dim({1}).set_stride(
                {1}));

        auto offset = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("offset").set_dim({1}).set_stride(
                {1}));

        auto rngAttributes = hipdnn_frontend::graph::RngAttributes()
                                 .set_name("rng_node")
                                 .set_seed(seed)
                                 .set_offset(offset)
                                 .set_dim({2, 2})
                                 .set_stride({1, 1})
                                 .set_distribution(hipdnn_frontend::RngDistribution_t::UNIFORM);

        auto output = graph->rng(seed, offset, rngAttributes);
        output->set_output(true);

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

        return std::make_tuple(graph, seed, offset, output);
    };

    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, seed, offset, output] = buildRngGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> seedTensor(seed->get_dim(), seed->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> offsetTensor(offset->get_dim(),
                                                               offset->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> outputTensor(output->get_dim(),
                                                               output->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[seed->get_uid()] = seedTensor.memory().deviceData();
    variantPack[offset->get_uid()] = offsetTensor.memory().deviceData();
    variantPack[output->get_uid()] = outputTensor.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 << "Rng graph execution complete. \n";

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