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

#include "ngraph/op/detection_output.hpp"

#include "itt.hpp"

using namespace std;

BWDCMP_RTTI_DEFINITION(ov::op::v0::DetectionOutput);

ov::op::v0::DetectionOutput::DetectionOutput(const Output<Node>& box_logits,
                                             const Output<Node>& class_preds,
                                             const Output<Node>& proposals,
                                             const Output<Node>& aux_class_preds,
                                             const Output<Node>& aux_box_preds,
                                             const Attributes& attrs)
    : Op({box_logits, class_preds, proposals, aux_class_preds, aux_box_preds}),
      m_attrs(attrs) {
    constructor_validate_and_infer_types();
}

ov::op::v0::DetectionOutput::DetectionOutput(const Output<Node>& box_logits,
                                             const Output<Node>& class_preds,
                                             const Output<Node>& proposals,
                                             const Attributes& attrs)
    : Op({box_logits, class_preds, proposals}),
      m_attrs(attrs) {
    constructor_validate_and_infer_types();
}

void ov::op::v0::DetectionOutput::validate_and_infer_types() {
    NGRAPH_OP_SCOPE(v0_DetectionOutput_validate_and_infer_types);
    NODE_VALIDATION_CHECK(this, m_attrs.num_classes > 0, "Number of classes must be greater than zero");

    NODE_VALIDATION_CHECK(this, m_attrs.keep_top_k.size() > 0, "keep_top_k attribute must be provided");

    NODE_VALIDATION_CHECK(this,
                          m_attrs.code_type == "caffe.PriorBoxParameter.CORNER" ||
                              m_attrs.code_type == "caffe.PriorBoxParameter.CENTER_SIZE",
                          "code_type must be either \"caffe.PriorBoxParameter.CORNER\" or "
                          "\"caffe.PriorBoxParameter.CENTER_SIZE\"");

    auto box_logits_et = get_input_element_type(0);
    NODE_VALIDATION_CHECK(this,
                          box_logits_et.is_real(),
                          "Box logits' data type must be floating point. Got " + box_logits_et.get_type_name());
    auto class_preds_et = get_input_element_type(1);
    NODE_VALIDATION_CHECK(this,
                          class_preds_et == box_logits_et,
                          "Class predictions' data type must be the same as box logits type (" +
                              box_logits_et.get_type_name() + "). Got " + class_preds_et.get_type_name());
    auto proposals_et = get_input_element_type(2);
    NODE_VALIDATION_CHECK(this,
                          proposals_et.is_real(),
                          "Proposals' data type must be floating point. Got " + proposals_et.get_type_name());

    const ov::PartialShape& box_logits_pshape = get_input_partial_shape(0);
    const ov::PartialShape& class_preds_pshape = get_input_partial_shape(1);
    const ov::PartialShape& proposals_pshape = get_input_partial_shape(2);

    int num_loc_classes = m_attrs.share_location ? 1 : m_attrs.num_classes;
    int prior_box_size = m_attrs.normalized ? 4 : 5;

    Dimension num_images = Dimension::dynamic();
    Dimension num_prior_boxes = Dimension::dynamic();
    if (box_logits_pshape.rank().is_static()) {
        NODE_VALIDATION_CHECK(
            this,
            box_logits_pshape.rank().get_length() == 2,
            "Box logits rank must be 2. Got " + std::to_string(box_logits_pshape.rank().get_length()));
        num_images = box_logits_pshape[0];
        if (box_logits_pshape[1].is_static()) {
            NODE_VALIDATION_CHECK(this,
                                  (box_logits_pshape[1].get_length() % (num_loc_classes * 4)) == 0,
                                  "Box logits' second dimension must be a multiply of num_loc_classes * 4 (" +
                                      std::to_string(num_loc_classes * 4) + "). Current value is: ",
                                  box_logits_pshape[1].get_length(),
                                  ".");
            num_prior_boxes = box_logits_pshape[1].get_length() / (num_loc_classes * 4);
        }
    }
    if (class_preds_pshape.rank().is_static()) {
        NODE_VALIDATION_CHECK(
            this,
            class_preds_pshape.rank().get_length() == 2,
            "Class predictions rank must be 2. Got " + std::to_string(class_preds_pshape.rank().get_length()));
        if (num_images.is_dynamic() && class_preds_pshape[0].is_static()) {
            num_images = class_preds_pshape[0];
        } else {
            NODE_VALIDATION_CHECK(this,
                                  class_preds_pshape[0].compatible(num_images),
                                  "Class predictions' first dimension is not compatible with batch size.");
        }
        if (class_preds_pshape[1].is_static()) {
            if (num_prior_boxes.is_dynamic()) {
                NODE_VALIDATION_CHECK(this,
                                      class_preds_pshape[1].get_length() % m_attrs.num_classes == 0,
                                      "Class predictions' second dimension must be a multiply of num_classes (" +
                                          std::to_string(m_attrs.num_classes) + "). Current value is: ",
                                      class_preds_pshape[1].get_length(),
                                      ".");
                num_prior_boxes = class_preds_pshape[1].get_length() / m_attrs.num_classes;
            } else {
                int num_prior_boxes_val = num_prior_boxes.get_length();
                NODE_VALIDATION_CHECK(this,
                                      class_preds_pshape[1].get_length() == num_prior_boxes_val * m_attrs.num_classes,
                                      "Class predictions' second dimension must be equal to num_prior_boxes * "
                                      "num_classes (" +
                                          std::to_string(num_prior_boxes_val * m_attrs.num_classes) +
                                          "). Current value is: ",
                                      class_preds_pshape[1].get_length(),
                                      ".");
            }
        }
    }
    if (proposals_pshape.rank().is_static()) {
        NODE_VALIDATION_CHECK(this,
                              proposals_pshape.rank().get_length() == 3,
                              "Proposals rank must be 3. Got " + std::to_string(proposals_pshape.rank().get_length()));
        if (num_images.is_static() && proposals_pshape[0].is_static()) {
            int64_t proposals_1st_dim = proposals_pshape[0].get_length();
            int64_t num_images_val = num_images.get_length();
            NODE_VALIDATION_CHECK(this,
                                  proposals_1st_dim == 1 || proposals_1st_dim == num_images_val,
                                  "Proposals' first dimension is must be equal to either batch size (" +
                                      std::to_string(num_images_val) +
                                      ") or 1. Got: " + std::to_string(proposals_1st_dim) + ".");
        }
        if (proposals_pshape[1].is_static()) {
            size_t proposals_expected_2nd_dim = m_attrs.variance_encoded_in_target ? 1 : 2;
            NODE_VALIDATION_CHECK(this,
                                  proposals_pshape[1].compatible(proposals_expected_2nd_dim),
                                  "Proposals' second dimension is mismatched. Current value is: ",
                                  proposals_pshape[1].get_length(),
                                  ", expected: ",
                                  proposals_expected_2nd_dim,
                                  ".");
        }
        if (proposals_pshape[2].is_static()) {
            if (num_prior_boxes.is_dynamic()) {
                NODE_VALIDATION_CHECK(this,
                                      proposals_pshape[2].get_length() % prior_box_size == 0,
                                      "Proposals' third dimension must be a multiply of prior_box_size (" +
                                          std::to_string(prior_box_size) + "). Current value is: ",
                                      proposals_pshape[2].get_length(),
                                      ".");
                num_prior_boxes = proposals_pshape[2].get_length() / prior_box_size;
            } else {
                int num_prior_boxes_val = num_prior_boxes.get_length();
                NODE_VALIDATION_CHECK(this,
                                      proposals_pshape[2].get_length() == num_prior_boxes_val * prior_box_size,
                                      "Proposals' third dimension must be equal to num_prior_boxes "
                                      "* prior_box_size (" +
                                          std::to_string(num_prior_boxes_val * prior_box_size) +
                                          "). Current value is: ",
                                      proposals_pshape[2].get_length(),
                                      ".");
            }
        }
    }

    if (get_input_size() > 3) {
        auto aux_class_preds_et = get_input_element_type(3);
        NODE_VALIDATION_CHECK(this,
                              aux_class_preds_et == class_preds_et,
                              "Additional class predictions' data type must be the same as class "
                              "predictions data type (" +
                                  class_preds_et.get_type_name() + "). Got " + aux_class_preds_et.get_type_name());
        auto aux_box_preds_et = get_input_element_type(4);
        NODE_VALIDATION_CHECK(this,
                              aux_box_preds_et == box_logits_et,
                              "Additional box predictions' data type must be the same as box logits data type (" +
                                  box_logits_et.get_type_name() + "). Got " + aux_box_preds_et.get_type_name());

        const ov::PartialShape& aux_class_preds_pshape = get_input_partial_shape(3);
        const ov::PartialShape& aux_box_preds_pshape = get_input_partial_shape(4);
        if (aux_class_preds_pshape.rank().is_static()) {
            NODE_VALIDATION_CHECK(this,
                                  aux_class_preds_pshape[0].compatible(num_images),
                                  "Additional class predictions' first dimension must be "
                                  "compatible with batch size.");
            if (num_prior_boxes.is_static()) {
                int num_prior_boxes_val = num_prior_boxes.get_length();
                NODE_VALIDATION_CHECK(this,
                                      aux_class_preds_pshape[1].get_length() == num_prior_boxes_val * 2,
                                      "Additional class predictions' second dimension must be equal to "
                                      "num_prior_boxes * 2 (" +
                                          std::to_string(num_prior_boxes_val * 2) + "). Got " +
                                          std::to_string(aux_class_preds_pshape[1].get_length()) + ".");
            }
        }
        NODE_VALIDATION_CHECK(this,
                              aux_box_preds_pshape.compatible(box_logits_pshape),
                              "Additional box predictions' shape must be compatible with box logits shape.");
    }

    std::vector<Dimension> output_shape{1, 1};
    if (m_attrs.keep_top_k[0] > 0) {
        output_shape.push_back(num_images * m_attrs.keep_top_k[0]);
    } else if (m_attrs.top_k > 0) {
        output_shape.push_back(num_images * m_attrs.top_k * m_attrs.num_classes);
    } else {
        output_shape.push_back(num_images * num_prior_boxes * m_attrs.num_classes);
    }
    output_shape.emplace_back(7);

    set_output_type(0, box_logits_et, output_shape);
}

shared_ptr<ov::Node> ov::op::v0::DetectionOutput::clone_with_new_inputs(const OutputVector& new_args) const {
    NGRAPH_OP_SCOPE(v0_DetectionOutput_clone_with_new_inputs);
    check_new_args_count(this, new_args);

    auto num_args = new_args.size();

    NODE_VALIDATION_CHECK(this, num_args == 3 || num_args == 5, "DetectionOutput accepts 3 or 5 inputs.");

    if (num_args == 3) {
        return make_shared<DetectionOutput>(new_args.at(0), new_args.at(1), new_args.at(2), m_attrs);
    } else {
        return make_shared<DetectionOutput>(new_args.at(0),
                                            new_args.at(1),
                                            new_args.at(2),
                                            new_args.at(3),
                                            new_args.at(4),
                                            m_attrs);
    }
}

bool ov::op::v0::DetectionOutput::visit_attributes(AttributeVisitor& visitor) {
    NGRAPH_OP_SCOPE(v0_DetectionOutput_visit_attributes);
    visitor.on_attribute("num_classes", m_attrs.num_classes);
    visitor.on_attribute("background_label_id", m_attrs.background_label_id);
    visitor.on_attribute("top_k", m_attrs.top_k);
    visitor.on_attribute("variance_encoded_in_target", m_attrs.variance_encoded_in_target);
    visitor.on_attribute("keep_top_k", m_attrs.keep_top_k);
    visitor.on_attribute("code_type", m_attrs.code_type);
    visitor.on_attribute("share_location", m_attrs.share_location);
    visitor.on_attribute("nms_threshold", m_attrs.nms_threshold);
    visitor.on_attribute("confidence_threshold", m_attrs.confidence_threshold);
    visitor.on_attribute("clip_after_nms", m_attrs.clip_after_nms);
    visitor.on_attribute("clip_before_nms", m_attrs.clip_before_nms);
    visitor.on_attribute("decrease_label_id", m_attrs.decrease_label_id);
    visitor.on_attribute("normalized", m_attrs.normalized);
    visitor.on_attribute("input_height", m_attrs.input_height);
    visitor.on_attribute("input_width", m_attrs.input_width);
    visitor.on_attribute("objectness_score", m_attrs.objectness_score);
    return true;
}
