// Tencent is pleased to support the open source community by making TNN available.
//
// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the 
// specific language governing permissions and limitations under the License.

#include <math.h>

#include "onnx2tnn.h"

int Onnx2TNN::FuseSignedMul(onnx::GraphProto* mutable_graph,
                                 std::vector<IndexNode> & index_nodes,
                                 std::map<std::string, onnx::TensorProto>& weights,
                                 std::map<std::string, int>& node_reference,
                                 std::set<std::string>& blob_names) {
    auto const node_count = index_nodes.size();

    for (int i = 0; i < node_count; i++) {
        auto node = index_nodes[i].node;

        // SignedMul <= Sub - Sign - Add - Div - Gather - Slice - Mul
        do {
            if (node->op_type() == "Sub" && i + 6 < node_count) {
                if (node_reference.find(node->output(0)) == node_reference.end() || node_reference[node->output(0)] != 1)
                    break;

                auto node_sub = node;
                auto node_sign = index_nodes[i+1].node;
                auto node_add = index_nodes[i+2].node;
                auto node_div = index_nodes[i+3].node;
                auto node_gather = index_nodes[i+4].node;
                auto node_slice = index_nodes[i+5].node;
                auto node_mul = index_nodes[i+6].node;

                if (node_sign->op_type() != "Sign" || node_add->op_type() != "Add" ||
                    node_div->op_type() != "Div" || node_gather->op_type() != "Gather" ||
                    node_slice->op_type() != "Slice" || node_mul->op_type() != "Mul")
                    break;

                // reduce
                node_sub->set_op_type(k_tnn_noop_type);
                node_sign->set_op_type(k_tnn_noop_type);
                node_add->set_op_type(k_tnn_noop_type);
                node_div->set_op_type(k_tnn_noop_type);
                node_gather->set_op_type(k_tnn_noop_type);
                node_slice->set_op_type(k_tnn_noop_type);

                node_reference.erase(node_reference.find(node_sub->output(0)));
                node_reference.erase(node_reference.find(node_sign->output(0)));
                node_reference.erase(node_reference.find(node_add->output(0)));
                node_reference.erase(node_reference.find(node_div->output(0)));
                node_reference.erase(node_reference.find(node_gather->output(0)));
                node_reference.erase(node_reference.find(node_slice->output(0)));
                blob_names.erase(node_sub->output(0));
                blob_names.erase(node_sign->output(0));
                blob_names.erase(node_add->output(0));
                blob_names.erase(node_div->output(0));
                blob_names.erase(node_gather->output(0));
                blob_names.erase(node_slice->output(0));

                node_mul->set_op_type("SignedMul");
                node_mul->clear_input();
                node_mul->add_input(node_sub->input(0));

                onnx::AttributeProto* attr_group = node_mul->add_attribute();
                attr_group->set_name("alpha");
                attr_group->set_f(1.0f);

                attr_group = node_mul->add_attribute();
                attr_group->set_name("beta");
                attr_group->set_f(1.0f);

                attr_group = node_mul->add_attribute();
                attr_group->set_name("gamma");
                attr_group->set_f(2.0f);

                i += 6;
            }
        } while (0);
    }

    ClearEmptyNode(index_nodes);
    return 0;
}
