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

#include "convert_reduce_no_keep_dims.hpp"

#include <memory>

#include "openvino/core/graph_util.hpp"
#include "openvino/core/node.hpp"
#include "openvino/core/rt_info.hpp"
#include "openvino/op/constant.hpp"
#include "openvino/op/squeeze.hpp"
#include "openvino/op/util/arithmetic_reductions_keep_dims.hpp"
#include "openvino/op/util/logical_reduction_keep_dims.hpp"
#include "openvino/pass/matcher_pass.hpp"
#include "openvino/pass/pattern/matcher.hpp"
#include "openvino/pass/pattern/op/label.hpp"
#include "openvino/pass/pattern/op/wrap_type.hpp"

template <class T>
ov::matcher_pass_callback ov::intel_cpu::ConvertReduceNoKeepDimsBase::convert_reduce() {
    return [&](ov::pass::pattern::Matcher& m) {
        auto reduce = std::dynamic_pointer_cast<T>(m.get_match_root());
        if (!reduce || reduce->get_keep_dims()) {
            return false;
        }

        reduce->set_keep_dims(true);
        const auto reduce_new = reduce->clone_with_new_inputs({reduce->input_value(0), reduce->input_value(1)});
        std::shared_ptr<ov::Node> squeeze = std::make_shared<ov::op::v0::Squeeze>(reduce_new, reduce->input_value(1));
        squeeze->set_friendly_name(reduce_new->get_friendly_name());
        ov::copy_runtime_info(reduce, {reduce_new, squeeze});
        ov::replace_node(reduce, squeeze);

        return true;
    };
}

template <typename ReductionType>
ov::intel_cpu::ConvertReduction<ReductionType>::ConvertReduction() {
    auto m = std::make_shared<ov::pass::pattern::Matcher>(
        ov::pass::pattern::wrap_type<ReductionType>(
            {ov::pass::pattern::any_input(), ov::pass::pattern::wrap_type<ov::op::v0::Constant>()}),
        "ConvertReduction");
    register_matcher(m, convert_reduce<ReductionType>());
}

template class ov::intel_cpu::ConvertReduction<ov::op::util::LogicalReductionKeepDims>;
template class ov::intel_cpu::ConvertReduction<ov::op::util::ArithmeticReductionKeepDims>;
