"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

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.
"""
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.

# Loop Optimizer.
# op in loop's body graph can be moved out to the loop

from oneflow.python.framework import id_util
from oneflow.python.onnx.util import MakeSure
from .optimizer_base import GraphOptimizerBase


# pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ


class LoopOptimizer(GraphOptimizerBase):
    """Loop Optimizer."""

    # a lot of terms used here come from loop's onnx spec
    # https://github.com/onnx/onnx/blob/master/docs/Operators.md#Loop
    def __init__(self):  # pylint: disable=useless-super-delegation
        super(LoopOptimizer, self).__init__()

    def _Optimize(self, graph):
        return self._ApplyOptimization(graph, self._OptimizeAtCurrentGraphLevel)

    def _OptimizeAtCurrentGraphLevel(self, g):
        has_update = True
        while has_update:
            has_update = False
            nodes = [n for n in g.get_nodes() if n.type == "Loop"]
            for n in nodes:
                has_update_tmp = self._TryMoveTransposeOutOfBodyGraph(n)
                if has_update_tmp:
                    has_update = True
                    self.graph_been_opt = True
        return g

    @staticmethod
    def ConsumerNodesNum(graph, node):
        MakeSure(len(node.output) == 1, "only consider node with only one output")
        res = len(graph.FindOutputConsumers(node.output[0]))
        return res

    def _TryMoveTransposeOutOfBodyGraph(self, loop_node):
        # output node of body graph can be loop-carried-dependent, if so it can't be move out of the body graph
        # return True if moving some nodes successfully
        # for now, we only consider moving transpose
        body_graph = loop_node.get_body_graphs()["body"]
        parent_graph = loop_node.graph
        scan_nodes_name_in_body, scan_node_in_parent = self._ScanOutputs(loop_node)
        scan_nodes = [
            body_graph.get_node_by_output(name) for name in scan_nodes_name_in_body
        ]
        graph_is_changed = False
        for node, name_in_parent in zip(scan_nodes, scan_node_in_parent):
            # 1 delete node in body graph if possible
            # only consider two case: trans is output, or transpose > identity > output
            need_process = False
            if (
                node.type == "Transpose"
                and self.ConsumerNodesNum(body_graph, node) <= 1
            ):
                trans = node
                new_output = node.input[0]
                body_graph.RemoveNode(node.name)
                need_process = True
            elif (
                node.type == "Identity"
                and node.inputs[0].type == "Transpose"
                and self.ConsumerNodesNum(body_graph, node) <= 1
                and self.ConsumerNodesNum(body_graph, node.inputs[0]) <= 1
            ):
                trans = node.inputs[0]
                new_output = node.inputs[0].input[0]
                body_graph.RemoveNode(node.inputs[0].name)
                body_graph.RemoveNode(node.name)
                need_process = True

            if need_process:
                # 2 correct body graph's output
                body_outputs = body_graph.outputs
                body_outputs[body_outputs.index(node.output[0])] = new_output
                # 3 insert new node in parent graph
                ori_perm = list(trans.get_attr("perm").ints)
                new_perm = [0] + [
                    i + 1 for i in ori_perm
                ]  # body output's rank is m > rank of loop's output is m+1
                name = id_util.UniqueStr("trans_moved_from_loop_body")
                _ = parent_graph.InsertNewNodeOnOutput(
                    "Transpose", name_in_parent, name, perm=new_perm
                )
                graph_is_changed = True

        return graph_is_changed

    @classmethod
    def _ScanOutputs(cls, loop):
        # loop has 2+N inputs; loop has N+K outputs;
        # loop's body graph has 1+N+K outputs
        loop_carried = len(loop.input) - 2
        body_graph = loop.get_body_graphs()["body"]
        return body_graph.outputs[loop_carried + 1 :], loop.output[loop_carried:]
