#include <iostream>

#include "hipdnn_frontend/Types.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 = int8_t;
    using BiasType = float;

    const int64_t n = 2; // Batch size
    // Input
    const int64_t c = 64; // Number of channels
    const int64_t h = 16; // Height
    const int64_t w = 8; // Width

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

    // Conv param
    const std::vector<int64_t> strides = {1, 1};
    const std::vector<int64_t> padding = {1, 1};
    const std::vector<int64_t> dilation = {1, 1};

    const int64_t outH = ((h + 2 * padding[0] - (dilation[0] * (r - 1) + 1)) / strides[0]) + 1;
    const int64_t outW = ((w + 2 * padding[1] - (dilation[1] * (s - 1) + 1)) / strides[1]) + 1;

    const int64_t vectorCount = 32;

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

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

        // create filter with NCHWc32
        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, r * s, s, 1})
                .set_vector_count(vectorCount));
        auto convFpropAttributes = hipdnn_frontend::graph::ConvFpropAttributes()
                                       .set_name("conv_fprop_node")
                                       .set_padding(padding)
                                       .set_stride(strides)
                                       .set_dilation(dilation);
        auto convOutput = graph->conv_fprop(input, filter, convFpropAttributes);

        // create sub node for dequantize:zero_point_dq
        auto zeroPointDq = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("zero_point_dq").set_value(0));
        auto convDeqSubAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                        .set_name("conv_deq_sub_node")
                                        .set_mode(hipdnn_frontend::PointwiseMode_t::SUB);
        auto convDeqSubOutput = graph->pointwise(convOutput, zeroPointDq, convDeqSubAttributes);

        // create mul node for dequantize:scale_dq
        auto scaleDq = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("scale_dq").set_value(2.0));
        auto convDeqMulAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                        .set_name("conv_deq_mul_node")
                                        .set_mode(hipdnn_frontend::PointwiseMode_t::MUL);
        auto convDeqMulOutput = graph->pointwise(convDeqSubOutput, scaleDq, convDeqMulAttributes);

        // create bias
        auto bias = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("bias")
                .set_dim({1, k, 1, 1})
                .set_stride({k, 1, 1, 1})
                .set_data_type(hipdnn_frontend::getDataTypeEnumFromType<BiasType>()));
        auto biasAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                  .set_name("bias_node")
                                  .set_mode(hipdnn_frontend::PointwiseMode_t::ADD);
        auto biasOutput = graph->pointwise(convDeqMulOutput, bias, biasAttributes);

        // create add
        auto add = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("add")
                .set_dim({n, k, outH, outW})
                .set_stride({outH * outW * k, outH * outW, outW, 1})
                .set_vector_count(vectorCount));

        // create sub node for dequantize:zero_point_dq_add
        auto zeroPointDqAdd = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("zero_point_dq_add").set_value(0));
        auto addDeqSubAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                       .set_name("add_deq_sub_node")
                                       .set_mode(hipdnn_frontend::PointwiseMode_t::SUB);
        auto addDeqSubOutput = graph->pointwise(add, zeroPointDqAdd, addDeqSubAttributes);

        // create mul node for dequantize:scale_dq_add
        auto scaleDqAdd = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("scale_dq_add").set_value(1.0));
        auto addDeqMulAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                       .set_name("add_deq_mul_node")
                                       .set_mode(hipdnn_frontend::PointwiseMode_t::MUL);
        auto addDeqMulOutput = graph->pointwise(addDeqSubOutput, scaleDqAdd, addDeqMulAttributes);

        //  create add
        auto addAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                 .set_name("add_node")
                                 .set_mode(hipdnn_frontend::PointwiseMode_t::ADD);
        auto addOutput = graph->pointwise(biasOutput, addDeqMulOutput, addAttributes);

        // create div node for quantize:scale_q
        auto scaleQ = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("scale_q").set_value(1));
        auto quantizeDivAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                         .set_name("quantize_div_node")
                                         .set_mode(hipdnn_frontend::PointwiseMode_t::DIV);
        auto quantizeDivOutput = graph->pointwise(addOutput, scaleQ, quantizeDivAttributes);

        // cretate  add node for quantize:zero_point_q.
        auto zeroPointQ = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes().set_name("zero_point_q").set_value(0));
        auto quantizeAddAttributes = hipdnn_frontend::graph::PointwiseAttributes()
                                         .set_name("quantize_add_node")
                                         .set_mode(hipdnn_frontend::PointwiseMode_t::ADD);
        auto quantizeOutput
            = graph->pointwise(quantizeDivOutput, zeroPointQ, quantizeAddAttributes);
        quantizeOutput->set_output(true).set_vector_count(vectorCount);

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

        return std::make_tuple(graph, input, filter, bias, add, quantizeOutput);
    };

    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, filter, bias, add, output] = buildConvBiasGraph(handle);

    // Allocate DCU memory
    hipdnn_data_sdk::utilities::Tensor<InputType> inputTensor(input->get_dim(),
                                                              input->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> addTensor(add->get_dim(), add->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> outTensor(output->get_dim(),
                                                            output->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[input->get_uid()] = inputTensor.memory().deviceData();
    variantPack[filter->get_uid()] = wTensor.memory().deviceData();
    variantPack[bias->get_uid()] = biasTensor.memory().deviceData();
    variantPack[add->get_uid()] = addTensor.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 << "int8_convolution_bias_add graph execution complete. \n";

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