// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT
#pragma once

#include "Node.hpp"
#include <hipdnn_data_sdk/utilities/ShapeUtilities.hpp>
#include <hipdnn_frontend/Error.hpp>
#include <hipdnn_frontend/Logging.hpp>
#include <hipdnn_frontend/attributes/ConvolutionDgradAttributes.hpp>
#include <hipdnn_frontend/attributes/GraphAttributes.hpp>
#include <hipdnn_frontend/detail/ConvolutionDgradPacker.hpp>
#include <hipdnn_frontend/detail/ConvolutionDgradUnpacker.hpp>

namespace hipdnn_frontend::graph
{
class ConvolutionDgradNode : public BaseNode<ConvolutionDgradNode, NodeType::CONVOLUTION_DGRAD>
{
public:
    ConvDgradAttributes attributes;

    ConvolutionDgradNode(ConvDgradAttributes&& convAttrs, const GraphAttributes& graphAttrs)
        : BaseNode(graphAttrs)
        , attributes(std::move(convAttrs))
    {
    }

    Error unpack_from_descriptor(
        hipdnnBackendDescriptor_t opDesc,
        std::unordered_map<int64_t, std::shared_ptr<TensorAttributes>>& tensorMap) override
    {
        ConvDgradAttributes attrs;
        HIPDNN_CHECK_ERROR(detail::unpackConvBpropOperation(opDesc, tensorMap, attrs));
        attributes = std::move(attrs);
        return {};
    }

    Error pre_validate_node() const override
    {
        HIPDNN_RETURN_IF_FALSE(
            attributes.get_dy(),
            ErrorCode::ATTRIBUTE_NOT_SET,
            "ConvolutionDgradNode missing dy (gradient of output) for pre-validation");

        HIPDNN_RETURN_IF_FALSE(attributes.get_w(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "ConvolutionDgradNode missing w (weights) for pre-validation");

        HIPDNN_RETURN_IF_FALSE(
            attributes.get_dx(),
            ErrorCode::ATTRIBUTE_NOT_SET,
            "ConvolutionDgradNode missing dx (gradient of input) for pre-validation");

        HIPDNN_RETURN_IF_TRUE(attributes.get_pre_padding().empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode missing pre_padding for pre-validation");

        HIPDNN_RETURN_IF_TRUE(attributes.get_post_padding().empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode missing post_padding for pre-validation");

        HIPDNN_RETURN_IF_TRUE(attributes.get_stride().empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode missing stride for pre-validation");

        HIPDNN_RETURN_IF_TRUE(attributes.get_dilation().empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode missing dilation for pre-validation");

        auto dy = attributes.get_dy();
        auto w = attributes.get_w();
        auto dx = attributes.get_dx();

        auto& dyDims = dy->get_dim();

        HIPDNN_RETURN_IF_LT(
            dyDims.size(),
            3,
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: dy tensor must have at least 3 dimensions (N, C, spatial)");

        auto& wDims = w->get_dim();

        HIPDNN_RETURN_IF_NE(
            wDims.size(),
            dyDims.size(),
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: Weight tensor dimension count must match dy tensor "
            "dimension count");

        // Validate output channels match between dy and w tensors
        HIPDNN_RETURN_IF_NE(
            dyDims[1],
            wDims[0],
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: dy tensor channels must match weight tensor output channels");

        auto& dxDims = dx->get_dim();

        HIPDNN_RETURN_IF_TRUE(dxDims.empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode: output dimension inference is not possible; "
                              "set dx dimensions explicitly.");

        HIPDNN_RETURN_IF_NE(dxDims.size(),
                            dyDims.size(),
                            ErrorCode::INVALID_VALUE,
                            "ConvolutionDgradNode: dx tensor dimension count must match dy tensor "
                            "dimension count");

        // Validate batch size matches
        HIPDNN_RETURN_IF_NE(dxDims[0],
                            dyDims[0],
                            ErrorCode::INVALID_VALUE,
                            "ConvolutionDgradNode: dx tensor batch size must match dy "
                            "tensor batch size");

        HIPDNN_RETURN_IF_NE(dxDims[1] % wDims[1],
                            0,
                            ErrorCode::INVALID_VALUE,
                            "ConvolutionDgradNode: dx tensor channels must be divisible by weight "
                            "tensor input channels");

        // dxDims[1] / wDims[1] is group count
        // weightChannels = inputChannels / groups
        // groups = inputChannels / weightChannels
        auto groupCount = dxDims[1] / wDims[1];
        HIPDNN_RETURN_IF_NE(
            wDims[0] % groupCount,
            0,
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: Weight tensor output channels must be divisible by "
            "the number of groups");

        // Validate spatial parameter counts match spatial dimensions
        auto spatialDims = dyDims.size() - 2; // Skip N and C dimensions
        auto& prePadding = attributes.get_pre_padding();
        auto& postPadding = attributes.get_post_padding();
        auto& stride = attributes.get_stride();
        auto& dilation = attributes.get_dilation();

        HIPDNN_RETURN_IF_NE(
            prePadding.size(),
            spatialDims,
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: pre_padding parameter count must match spatial dimension count");

        HIPDNN_RETURN_IF_NE(postPadding.size(),
                            spatialDims,
                            ErrorCode::INVALID_VALUE,
                            "ConvolutionDgradNode: post_padding parameter count must match spatial "
                            "dimension count");

        HIPDNN_RETURN_IF_NE(
            stride.size(),
            spatialDims,
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: stride parameter count must match spatial dimension count");

        HIPDNN_RETURN_IF_NE(
            dilation.size(),
            spatialDims,
            ErrorCode::INVALID_VALUE,
            "ConvolutionDgradNode: dilation parameter count must match spatial dimension count");

        // Check spatial parameters for each dimension
        for(size_t i = 0; i < spatialDims; ++i)
        {
            auto prePad = prePadding[i];
            auto postPad = postPadding[i];
            auto strideVal = stride[i];
            auto dilationVal = dilation[i];

            HIPDNN_RETURN_IF_LT(
                strideVal, 1, ErrorCode::INVALID_VALUE, "ConvolutionDgradNode: Stride must be > 0");

            HIPDNN_RETURN_IF_LT(dilationVal,
                                1,
                                ErrorCode::INVALID_VALUE,
                                "ConvolutionDgradNode: Dilation must > 0");

            HIPDNN_RETURN_IF_LT(prePad,
                                0,
                                ErrorCode::INVALID_VALUE,
                                "ConvolutionDgradNode: Pre-padding must be non-negative");

            HIPDNN_RETURN_IF_LT(postPad,
                                0,
                                ErrorCode::INVALID_VALUE,
                                "ConvolutionDgradNode: Post-padding must be non-negative");

            auto outputSize = dyDims[i + 2];
            auto kernelSize = wDims[i + 2];
            auto inputSize = dxDims[i + 2];

            auto dilatedKernelSize = (dilationVal * (kernelSize - 1)) + 1;
            auto numerator = inputSize + prePad + postPad - dilatedKernelSize;

            HIPDNN_RETURN_IF_LT(numerator,
                                0,
                                ErrorCode::INVALID_VALUE,
                                "ConvolutionDgradNode: Input spatial dimension at index "
                                    + std::to_string(i) + " (" + std::to_string(inputSize)
                                    + ") is too small for the kernel size ("
                                    + std::to_string(kernelSize) + ") and dilation ("
                                    + std::to_string(dilationVal) + ")");

            const int64_t expectedOutputSize = (numerator / strideVal) + 1;

            HIPDNN_RETURN_IF_NE(
                outputSize,
                expectedOutputSize,
                ErrorCode::INVALID_VALUE,
                "ConvolutionDgradNode: dy tensor spatial dimension at index " + std::to_string(i)
                    + " (" + std::to_string(outputSize) + ") does not match expected dimension ("
                    + std::to_string(expectedOutputSize)
                    + ") given dx dimensions, kernel size, padding, stride, and dilation");
        }

        return {};
    }

    Error infer_properties_node() override
    {
        auto dy = attributes.get_dy();
        auto w = attributes.get_w();
        auto dx = attributes.get_dx();

        HIPDNN_RETURN_IF_FALSE(dy,
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "ConvolutionDgradNode missing dy for setting properties");

        HIPDNN_RETURN_IF_FALSE(w,
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "ConvolutionDgradNode missing w for setting properties");

        HIPDNN_RETURN_IF_FALSE(dx,
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "ConvolutionDgradNode missing dx for setting properties");

        HIPDNN_CHECK_ERROR(attributes.fill_from_context(graph_attributes));

        auto& dxDims = dx->get_dim();

        HIPDNN_RETURN_IF_TRUE(dxDims.empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "ConvolutionDgradNode: output dimension inference is not possible; "
                              "set dx dimensions explicitly.");

        if(dx->get_stride().empty())
        {
            auto& dyStrides = dy->get_stride();
            auto& dxDimsFinal = dx->get_dim();

            HIPDNN_RETURN_IF_TRUE(
                dyStrides.empty(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "ConvolutionDgradNode: Cannot infer dx strides - missing dy strides");

            HIPDNN_RETURN_IF_NE(
                dyStrides.size(),
                dxDimsFinal.size(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "ConvolutionDgradNode: Stride dimension mismatch between dy and dx tensors");

            // Extract stride order from dy tensor and apply to dx tensor
            auto strideOrder = hipdnn_data_sdk::utilities::extractStrideOrder(dyStrides);

            // Generate dx strides using the extracted stride order and dx dimensions
            auto dxStrides = hipdnn_data_sdk::utilities::generateStrides(dxDimsFinal, strideOrder);

            dx->set_stride(dxStrides);
        }

        return {};
    }

    Error create_operation(
        std::unordered_map<int64_t, detail::ScopedHipdnnBackendDescriptor>& tensorDescs,
        std::vector<detail::ScopedHipdnnBackendDescriptor>& operations) const override
    {
        return detail::createConvDgradOperation(attributes, tensorDescs, operations);
    }
};

typedef ConvolutionDgradNode DgradNode;
}