#include <iostream>

#include <hipdnn_data_sdk/utilities/Tensor.hpp>
#include <hipdnn_data_sdk/utilities/Workspace.hpp>
#include <hipdnn_frontend.hpp>

#include "utils.hpp"

int main()
{
    using InputType = float;

    const int64_t n = 4; // Batch size

    // Input
    const int64_t c = 64; // Number of channels
    const int64_t h = 16; // Height
    const int64_t w = 16; // Width

    // Filter
    const int64_t k = 64; // Number of filters
    const int64_t r = 1; // Height
    const int64_t s = 1; // 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 outH = ((h + 2 * padH - (dilH * (r - 1) + 1)) / strideH) + 1;
    const int64_t outW = ((w + 2 * padW - (dilW * (s - 1) + 1)) / strideW) + 1;

    const int64_t g = 1; // Number of groups

    auto buildDeformConvBackwardGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();
        graph->set_name("deform_conv_backward_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 image = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("image")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));

        auto filter = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("filter")
                .set_dim({k, c / g, r, s})
                .set_stride({c / g * r * s, 1, c / g * s, c / g}));

        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 offset = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("offset")
                .set_dim({n, 2 * g * r * s, outH, outW})
                .set_stride({2 * g * r * s * outH * outW, 1, 2 * g * r * s * outW, 2 * g * r * s}));

        auto mask = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("mask")
                .set_dim({n, g * r * s, outH, outW})
                .set_stride({g * r * s * outH * outW, 1, g * r * s * outW, g * r * s}));

        auto deformConvBwdAttributes = hipdnn_frontend::graph::DeformConvDgradAttributes()
                                           .set_name("deform_conv_backward_node")
                                           .set_padding({padH, padW})
                                           .set_stride({strideH, strideW})
                                           .set_dilation({dilH, dilW})
                                           .set_x(image)
                                           .set_mask(mask);

        auto [dx, doffset, dmask]
            = graph->deform_conv_dgrad(loss, filter, offset, deformConvBwdAttributes);
        dx->set_output(true);
        doffset->set_output(true);
        dmask->set_output(true);

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

        return std::make_tuple(graph, image, filter, loss, offset, mask, dx, doffset, dmask);
    };

    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, image, filter, loss, offset, mask, dx, doffset, dmask]
        = buildDeformConvBackwardGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> lossTensor(loss->get_dim(), loss->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> inputTensor(image->get_dim(),
                                                              image->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> filterTensor(filter->get_dim(),
                                                               filter->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> offsetTensor(offset->get_dim(),
                                                               offset->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> maskTensor(mask->get_dim(), mask->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dxTensor(dx->get_dim(), dx->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> doffsetTensor(doffset->get_dim(),
                                                                doffset->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dmaskTensor(dmask->get_dim(),
                                                              dmask->get_stride());

    // Pixel-level offset values for each sampling point of the convolution kernel
    offsetTensor.fillWithRandomValues(static_cast<InputType>(0.0f), static_cast<InputType>(1.0f));

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[loss->get_uid()] = lossTensor.memory().deviceData();
    variantPack[image->get_uid()] = inputTensor.memory().deviceData();
    variantPack[filter->get_uid()] = filterTensor.memory().deviceData();
    variantPack[offset->get_uid()] = offsetTensor.memory().deviceData();
    variantPack[mask->get_uid()] = maskTensor.memory().deviceData();
    variantPack[dx->get_uid()] = dxTensor.memory().deviceData();
    variantPack[doffset->get_uid()] = doffsetTensor.memory().deviceData();
    variantPack[dmask->get_uid()] = dmaskTensor.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 << "Deformable convolution backward graph execution complete.\n\n";

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