#################################################################################################
#
# Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#################################################################################################

"""
Eliminate layout manipulation nodes
"""

from copy import deepcopy

from cutlass.backend.evt.ir import DAGIR, LayoutNode
from cutlass.backend.evt.passes.pass_manager import EVTPassBase
from cutlass.backend.evt.passes.pass_shape_type_propagation import PassShapeTypePropagation


class PassLayoutManipulateElimination(EVTPassBase):
    """
    Eliminate layout manipulation nodes
    """
    dependencies = [PassShapeTypePropagation]

    def __init__(self, dag_ir: DAGIR) -> None:
        super().__init__(dag_ir)
        self.copy_cnt = 0

    def call(self):
        self.layout_nodes_worklist = self.get_all_layout_nodes()
        # Run while loop utill all layout nodes are eliminated
        while(len(self.layout_nodes_worklist) > 0):
            node = self.layout_nodes_worklist.pop(0)
            # for node in layout_nodes:
            # Step 1: get the propagation direction
            direction = self.get_propagation_direction(node)
            self.visited = []
            getattr(self, f"propagate_to_{direction}")(self.dag_ir.get_node_meta(node), node)
            # Eliminate the current node
            input_node = self.dag_ir.get_all_inputs(node)[0]
            self.dag_ir.replace_all_uses_with(node, input_node)
            # layout_nodes = self.get_all_layout_nodes()

    def get_all_layout_nodes(self):
        layout_nodes = []
        for node_meta in reversed(self.dag_ir.node_metas_topological_order()):
            if isinstance(node_meta, LayoutNode):
                layout_nodes.append(node_meta.name)
        return layout_nodes

    def get_propagation_direction(self, node: str):
        """
        The logic is propagating all layout nodes away from the accumulator node.
        """
        self.visited = []
        self.get_influenced_users(node)
        nodes_influenced_dir_users = self.visited
        self.visited = []
        self.get_influenced_inputs(node)
        nodes_influenced_dir_inputs = self.visited

        if "accum" in nodes_influenced_dir_users and "accum" not in nodes_influenced_dir_inputs:
            return "inputs"
        elif "accum" not in nodes_influenced_dir_users and "accum" in nodes_influenced_dir_inputs:
            return "users"
        else:
            raise RuntimeError("Unsolved propagation direction")

    # Get all influenced nodes if we propagate along the user direction
    def get_influenced_users(self, node: str):
        if node in self.visited:
            return
        self.visited.append(node)

        users = self.dag_ir.get_users(node)
        for user in users:
            self.get_influenced_users(user)
        user_inputs = []
        for user in users:
            user_inputs.append(set(self.dag_ir.get_all_inputs(user)))
        if len(user_inputs) > 0:
            user_inputs = set.union(*user_inputs)
            user_inputs.remove(node)
            for input in user_inputs:
                self.get_influenced_inputs(input)

    # Get all influenced nodes if we propagate along the input direction
    def get_influenced_inputs(self, node: str):
        if node in self.visited:
            return
        self.visited.append(node)

        inputs = self.dag_ir.get_all_inputs(node)
        for input in inputs:
            self.get_influenced_inputs(input)
        input_users = []
        for input in inputs:
            input_users.append(set(self.dag_ir.get_users(input)))
        if len(input_users) > 0:
            input_users = set.union(*input_users)
            input_users.remove(node)
            for user in input_users:
                self.get_influenced_users(user)

    def add_copy_before(self, layout_node_meta: LayoutNode, target: str):
        copied_node_meta = deepcopy(layout_node_meta)
        copied_node = f"{copied_node_meta.name}_copy{self.copy_cnt}"
        self.copy_cnt += 1
        copied_node_meta.name = copied_node
        self.dag_ir.add_node(copied_node_meta)
        # Add edges
        target_inputs = self.dag_ir.get_all_inputs(target)
        for src in target_inputs:
            self.dag_ir.remove_edge(src, target)
            self.dag_ir.add_edge(src, copied_node)
        self.dag_ir.add_edge(copied_node, target)
        self.layout_nodes_worklist.append(copied_node)

    def add_copy_after(self, layout_node_meta: LayoutNode, target: str):
        copied_node_meta = deepcopy(layout_node_meta)
        copied_node = f"{copied_node_meta.name}_copy{self.copy_cnt}"
        self.copy_cnt += 1
        copied_node_meta.name = copied_node
        self.dag_ir.add_node(copied_node_meta)
        # Add edges
        users = self.dag_ir.get_users(target)
        for user in users:
            self.dag_ir.remove_edge(target, user)
            self.dag_ir.add_edge(copied_node, user)
        self.dag_ir.add_edge(target, copied_node)
        self.layout_nodes_worklist.append(copied_node)

    # Propagate the layout `node` along the user direction
    def propagate_to_users(self, layout_node_meta: LayoutNode, node: str):
        """
        Propagate layout node to users
        """
        if node in self.visited:
            # Avoid applying twice
            return
        self.visited.append(node)

        node_meta = self.dag_ir.get_node_meta(node)
        if layout_node_meta.name != node:
            if isinstance(node_meta, LayoutNode):
                # Layout node is not transparent with layout node
                self.add_copy_before(layout_node_meta, node)
                return
            else:
                layout_node_meta.apply_to_user(node_meta)

        users = self.dag_ir.get_users(node)
        user_inputs = []
        for user in users:
            user_inputs.append(set(self.dag_ir.get_all_inputs(user)))
        for user in users:
            self.propagate_to_users(layout_node_meta, user)
        if len(user_inputs) > 0:
            user_inputs = set.union(*user_inputs)
            user_inputs.remove(node)
            for input in user_inputs:
                self.propagate_to_inputs(layout_node_meta.get_inverse_node(), input)

    # Propagate the layout `node` along the input direction
    def propagate_to_inputs(self, layout_node_meta: LayoutNode, node: str):
        """
        Propagate layout node to inputs
        """
        if node in self.visited:
            # Avoid applying twice
            return
        self.visited.append(node)

        node_meta = self.dag_ir.get_node_meta(node)
        if layout_node_meta.name != node:
            if isinstance(node_meta, LayoutNode):
                # Layout node is not transparent with layout node
                self.add_copy_after(layout_node_meta, node)
                return
            else:
                layout_node_meta.apply_to_input(node_meta)
        inputs = self.dag_ir.get_all_inputs(node)
        input_users = []
        for input in inputs:
            input_users.append(set(self.dag_ir.get_users(input)))
        for input in inputs:
            self.propagate_to_inputs(layout_node_meta, input)
        if len(input_users) > 0:
            input_users = set.union(*input_users)
            input_users.remove(node)
            for user in input_users:
                self.propagate_to_users(layout_node_meta.get_inverse_node(), user)
