// Copyright (C) 2018-2021 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once

#include "ngraph/op/op.hpp"

namespace ngraph
{
    namespace op
    {
        namespace v1
        {
            class NGRAPH_API DeformablePSROIPooling : public Op
            {
            public:
                static constexpr NodeTypeInfo type_info{"DeformablePSROIPooling", 1};
                const NodeTypeInfo& get_type_info() const override { return type_info; }
                DeformablePSROIPooling() = default;
                /// \brief Constructs a DeformablePSROIPooling operation
                ///
                /// \param input           Input tensor with position sensitive score maps
                /// \param coords          Input tensor with list of five element tuples
                ///                        describing ROI coordinates
                /// \param offsets         Input tensor with transformation values
                /// \param output_dim      Pooled output channel number
                /// \param group_size      Number of horizontal bins per row to divide ROI area,
                ///                        it defines output width and height
                /// \param spatial_scale   Multiplicative spatial scale factor to translate ROI
                ///                        coordinates from their input scale to the scale used when
                ///                        pooling
                /// \param mode            Specifies mode for pooling.
                /// \param spatial_bins_x  Specifies numbers of bins to divide ROI single
                ///                        bin over width
                /// \param spatial_bins_y  Specifies numbers of bins to divide ROI single
                ///                        bin over height
                /// \param no_trans        The flag that specifies whenever third input exists
                ///                        and contains transformation (offset) values
                /// \param trans_std       The value that all transformation (offset) values are
                ///                        multiplied with
                /// \param part_size       The number of parts the output tensor spatial dimensions
                ///                        are divided into. Basically it is the height
                ///                        and width of the third input
                DeformablePSROIPooling(const Output<Node>& input,
                                       const Output<Node>& coords,
                                       const Output<Node>& offsets,
                                       const int64_t output_dim,
                                       const float spatial_scale,
                                       const int64_t group_size = 1,
                                       const std::string mode = "bilinear_deformable",
                                       int64_t spatial_bins_x = 1,
                                       int64_t spatial_bins_y = 1,
                                       float trans_std = 1,
                                       int64_t part_size = 1);

                DeformablePSROIPooling(const Output<Node>& input,
                                       const Output<Node>& coords,
                                       const int64_t output_dim,
                                       const float spatial_scale,
                                       const int64_t group_size = 1,
                                       const std::string mode = "bilinear_deformable",
                                       int64_t spatial_bins_x = 1,
                                       int64_t spatial_bins_y = 1,
                                       float trans_std = 1,
                                       int64_t part_size = 1);

                bool visit_attributes(AttributeVisitor& visitor) override;

                void validate_and_infer_types() override;

                virtual std::shared_ptr<Node>
                    clone_with_new_inputs(const OutputVector& new_args) const override;

                int64_t get_output_dim() const { return m_output_dim; }
                int64_t get_group_size() const { return m_group_size; }
                float get_spatial_scale() const { return m_spatial_scale; }
                const std::string& get_mode() const { return m_mode; }
                int64_t get_spatial_bins_x() const { return m_spatial_bins_x; }
                int64_t get_spatial_bins_y() const { return m_spatial_bins_y; }
                float get_trans_std() const { return m_trans_std; }
                int64_t get_part_size() const { return m_part_size; }

            private:
                int64_t m_output_dim;
                float m_spatial_scale;
                int64_t m_group_size = 1;
                std::string m_mode = "bilinear_deformable";
                int64_t m_spatial_bins_x = 1;
                int64_t m_spatial_bins_y = 1;
                float m_trans_std = 1.f;
                int64_t m_part_size = 1;
            };
        } // namespace v1
    }     // namespace op
} // namespace ngraph
