#include <iostream>

#include "hipdnn_frontend/Types.hpp"
#include "hipdnn_frontend/attributes/PointwiseAttributes.hpp"
#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 n = 4; // Batch size
    const int64_t c = 64; // Number of channels
    const int64_t h = 16; // Height
    const int64_t w = 16; // Width
    const int64_t k = 64;

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

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

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

        auto reluBwdAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                     .set_name("relu_bwd")
                                     .set_mode(hipdnn_frontend::PointwiseMode::RELU_BWD);

        auto dReluOutput = graph->pointwise(dgrad, input, reluBwdAttributes);
        dReluOutput->set_output(true);

        auto reductionAttributes = hipdnn_frontend::graph::ReductionAttributes()
                                       .set_name("reduction")
                                       .set_mode(hipdnn_frontend::ReductionMode_t::ADD);

        auto reduceOutput = graph->reduction(dReluOutput, reductionAttributes);
        reduceOutput->set_output(true)
            .set_dim({1, k, 1, 1})
            .set_stride({k, 1, k, k})
            .set_name("reduce_output");

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

        return std::make_tuple(graph, input, dgrad, dReluOutput, reduceOutput);
    };

    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, input, dgrad, dx, dBias] = buildPointwiseReductionGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> inputTensor(input->get_dim(),
                                                              input->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dgradTensor(dgrad->get_dim(),
                                                              dgrad->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dxTensor(dx->get_dim(), dx->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dBiasTensor(dBias->get_dim(),
                                                              dBias->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[input->get_uid()] = inputTensor.memory().deviceData();
    variantPack[dgrad->get_uid()] = dgradTensor.memory().deviceData();
    variantPack[dx->get_uid()] = dxTensor.memory().deviceData();
    variantPack[dBias->get_uid()] = dBiasTensor.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 << "PointwiseReduction graph execution complete. \n";

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