#  Copyright (c) 2023, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause


from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
from coremltools.converters.mil.mil.passes.helper import block_context_manager
from coremltools.converters.mil.mil.passes.pass_registry import register_pass


@register_pass(namespace="common")
class topological_reorder(AbstractGraphPass):
    """
    Topologically re-orders the list of operations in a program by places each operation closer to its
    first use, or at the end if it's not consumed by any other operation.

    Currently, This pass re-orders only Transpose and Cast operations.

    .. code-block::

        # Example: input program
        main(x: (2, 4, fp32)) {
            x = mb.cast(x=x, dtype="fp16")
            x1 = mb.square(x=x)
            x1_t = mb.transpose(x=x1, perm=[1, 0])
            x2 = mb.cast(x=x1_t, dtype="fp32")
            x3 = mb.log(x=x)
            x3_t = mb.transpose(x=x3, perm=[1, 0])
            x4 = mb.cast(x=x3_t, dtype="fp32")
            x5 = mb.relu(x=x)
            x6 = mb.cast(x=x5, dtype="fp32")
            x7 = mb.relu(x=x6)
            x8 = mb.relu(x=x)
        } -> x2, x4, x7, x8

        # After moving `cast` ops becomes
        main(x: (2, 4, fp32)) {
            x = mb.cast(x=x, dtype="fp16")
            x1 = mb.square(x=x)
            x1_t = mb.transpose(x=x1, perm=[1, 0])
            x3 = mb.log(x=x)
            x3_t = mb.transpose(x=x3, perm=[1, 0])
            x5 = mb.relu(x=x)
            x6 = mb.cast(x=x5, dtype="fp32")
            x7 = mb.relu(x=x6)
            x8 = mb.relu(x=x)
            x4 = mb.cast(x=x3_t, dtype="fp32")
            x2 = mb.cast(x=x1_t, dtype="fp32")
        } -> x2, x4, x7, x8

        # After moving `transpose` ops becomes
        main(x: (2, 4, fp32)) {
            x = mb.cast(x=x, dtype="fp16")
            x1 = mb.square(x=x)
            x3 = mb.log(x=x)
            x5 = mb.relu(x=x)
            x6 = mb.cast(x=x5, dtype="fp32")
            x7 = mb.relu(x=x6)
            x8 = mb.relu(x=x)
            x3_t = mb.transpose(x=x3, perm=[1, 0])
            x4 = mb.cast(x=x3_t, dtype="fp32")
            x1_t = mb.transpose(x=x1, perm=[1, 0])
            x2 = mb.cast(x=x1_t, dtype="fp32")
        } -> x2, x4, x7, x8
    """

    def METHOD_NAME(self, prog):
        for f_name, f in prog.functions.items():
            self._move_operations_to_the_end_block(f, ["cast", "transpose"])

    @staticmethod
    @block_context_manager
    def _move_operations_to_the_end_block(block, op_type_to_move):
        # Moves ops with `op_type_to_move` in `block.operations` (list) to the end of the program.
        # Note: ops with `op_type_to_move` and is dead code are moved toward end, which can be eliminated
        # later with dead-code-elimination pass.
        #
        # Inputs:
        #  - block (mil.Block): block to be modified in-place
        #  - op_type_to_move (List[str])
        # Returns:
        #  - set[Var]: Set of vars consumed in block (or returned as block output)

        # first_use maps var to (index, op) representing the first op in block.operation that consumes this var.
        first_use = {}  # var -> op
        ops_to_remove = []  # list of ops to be deleted at the end of pass
        for index, op in enumerate(reversed(block.operations[:])):
            current_op = op

            if op.op_type in op_type_to_move:
                # Mark op for deletion
                ops_to_remove.append(op)

                # Create list of operations consuming each output of current operation
                first_consumers = [first_use[v] for v in op.outputs if v in first_use]

                before_op = None  # None means adding at the end of block
                if len(first_consumers) > 0:
                    # Current op should be moved right before this first consumer of one of it's output.
                    #  1. Find indices for all the consumer ops of outputs
                    #  2. Move current op right before first consumer i.e. smallest index in block.operations
                    first_use_indices = [
                        block.operations.index(first_use_op) for first_use_op in first_consumers
                    ]
                    before_op = block.operations[min(first_use_indices)]

                # Create new copy of current operation
                new_var = getattr(mb, op.op_type)(**op.inputs, before_op=before_op)

                if not isinstance(new_var, (list, tuple)):
                    new_var = [new_var]

                # Override current_op to be newly created op to ensure `first_use`
                # points to newly created op instead of old one.
                current_op = new_var[0].op

                for old_output_var, new_output_var in zip(op.outputs, new_var):
                    block.replace_uses_of_var_after_op(
                        anchor_op=None, old_var=old_output_var, new_var=new_output_var
                    )

            # Collect input vars from sub-block if present
            relevant_inputs = set()
            for b in current_op.blocks:
                relevant_inputs |= topological_reorder._move_operations_to_the_end_block(
                    b, op_type_to_move
                )

            # Collect vars from operation input
            for v in current_op.inputs.values():
                if isinstance(v, (tuple, list)):
                    relevant_inputs |= set(v)
                    continue
                relevant_inputs.add(v)

            # Mark current op as first use for all the input vars
            #  a) of it's sub-block
            #  b) of current op
            for v in relevant_inputs:
                # input is seen for the first time or
                # current_op is first_use i.e. appears before earlier recorded first_use.
                # Note: since ops are moved to the end, it's possible that an op is moved right after
                # earlier recorded first_use and in such cases, first_use should not be modified.
                #
                # == Example ==
                #     main( %x: (10, 20, fp32)(Tensor)) {
                #      block0() {
                #         %cast_0: (10, 20, fp16)(Tensor) = cast(x= %x, dtype = "fp16", name = "cast_0")
                #         %cast_1: (10, 20, fp32)(Tensor) = cast(x= %cast_0, dtype = "fp32", name = "cast_1")
                #         %transpose_0: (20, 10, fp16)(Tensor) = transpose(x= %cast_0, perm = [1, 0], name = "transpose_0")
                #         %transpose_1: (10, 20, fp16)(Tensor) = transpose(x= %transpose_0, perm = [1, 0], name = "transpose_1")
                #       } -> (% cast_1, % transpose_1)
                #     }
                # In above example, `%cast_1` will be moved to the end of the block and first_use info for `%cast_0`
                # should point to `%transpose_0` and not to `%cast_1`
                if v not in first_use or block.operations.index(
                    first_use[v]
                ) > block.operations.index(current_op):
                    first_use[v] = current_op

        # Remove ops that are reordered
        block.remove_ops(ops_to_remove)

        # Returns set of vars consumed in current block
        vars_consumed_in_block = set([v for v in first_use])
        vars_consumed_in_block.update(block.outputs)
        return vars_consumed_in_block