#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 n = 1; // Batch size
    // Input
    const int64_t c = 32; // Number of channels
    const int64_t outH = 270; // Height
    const int64_t outW = 480; // Width

    // Filter
    const int64_t k = 32; // Number of filters
    const int64_t r = 3; // Height
    const int64_t s = 3; // Width

    // Conv param
    const int64_t strideH = 1; // Height stride
    const int64_t strideW = 1; // Width stride
    const int64_t padH = 0; // Height padding
    const int64_t padW = 0; // Width padding
    const int64_t dilH = 1; // Height dilation
    const int64_t dilW = 1; // Width dilation
    const int64_t outputPadH = 0; // Output height padding
    const int64_t outputPadW = 0; // Output width padding

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

        graph->set_name("convbwd_bias_relu_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); //

        // create conv
        auto loss = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("loss")
                .set_dim({n, k, outH, outW})
                .set_stride({k * outH * outW, 1, k * outW, k}));
        auto filter = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("filter")
                .set_dim({k, c, r, s})
                .set_stride({c * r * s, 1, c * s, c}));
        auto convDgradAttributes = hipdnn_frontend::graph::ConvDgradAttributes()
                                       .set_name("conv_backward_node")
                                       .set_padding({padH, padW})
                                       .set_stride({strideH, strideW})
                                       .set_dilation({dilH, dilW})
                                       .set_output_padding({outputPadH, outputPadW});
        auto convOutput = graph->conv_dgrad(loss, filter, convDgradAttributes);

        // create bias
        auto biasIn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("bias_in")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));
        auto biasAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                  .set_name("bias_node")
                                  .set_mode(hipdnn_frontend::PointwiseMode_t::ADD);
        auto biasOutput = graph->pointwise(convOutput, biasIn, biasAttributes);

        // create relu
        auto reluAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                  .set_name("relu_node")
                                  .set_mode(hipdnn_frontend::PointwiseMode_t::RELU_FWD);
        auto output = graph->pointwise(biasOutput, reluAttributes);
        output->set_output(true);

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

        return std::make_tuple(graph, loss, filter, biasIn, 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, loss, filter, bias, output] = buildConvbwdBiasReluGraph(handle);

    // Allocate DCU memory
    hipdnn_data_sdk::utilities::Tensor<InputType> lossTensor(loss->get_dim(), loss->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> wTensor(filter->get_dim(), filter->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> biasTensor(bias->get_dim(), bias->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> outTensor(output->get_dim(),
                                                            output->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[loss->get_uid()] = lossTensor.memory().deviceData();
    variantPack[filter->get_uid()] = wTensor.memory().deviceData();
    variantPack[bias->get_uid()] = biasTensor.memory().deviceData();
    variantPack[output->get_uid()] = outTensor.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 << "Convolutionbwd_bias_relu graph execution complete. \n";

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