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

#include "openvino/frontend/paddle/node_context.hpp"
#include "openvino/opsets/opset6.hpp"

namespace ov {
namespace frontend {
namespace paddle {
namespace op {
NamedOutputs argmax(const NodeContext& node) {
    auto data = node.get_input("X");
    bool flatten = node.get_attribute<bool>("flatten");
    const element::Type& index_element_type = element::i64;
    const Output<ov::Node> k = ov::opset6::Constant::create(ov::element::i64, {}, {1});

    if (!flatten) {
        auto axis = node.get_attribute<int64_t>("axis");
        const auto axis_to_remove = ov::opset6::Constant::create(element::u64, Shape{}, {axis});
        auto node_topk = std::make_shared<ov::opset6::TopK>(data, k, axis, "max", "index", index_element_type);
        const auto reshaped_indices = std::make_shared<ov::opset6::Squeeze>(node_topk->output(1), axis_to_remove);
        return node.default_single_output_mapping(
            {std::make_shared<ov::opset6::Convert>(reshaped_indices, element::i64)},
            {"Out"});
    } else {
        int64_t axis = 0;
        const Output<ov::Node> reshape_flatten = ov::opset6::Constant::create(ov::element::i64, {1}, {-1});
        auto node_reshape = std::make_shared<ov::opset6::Reshape>(data, reshape_flatten, true);
        auto node_topk = std::make_shared<ov::opset6::TopK>(node_reshape, k, axis, "max", "index", index_element_type);
        return node.default_single_output_mapping(
            {std::make_shared<ov::opset6::Convert>(node_topk->output(1), element::i64)},
            {"Out"});
    }
}

}  // namespace op
}  // namespace paddle
}  // namespace frontend
}  // namespace ov
