# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""The visitor and mutator infra for nn.Module."""

from typing import Any

from . import core as nn


class Mutator:
    """The mutator for nn.Module transform. Users can override the `visit_*` methods
    to apply transform in different structures, or even override the `visit` method
    to change the logic of traversal."""

    def visit_module(self, name: str, node: nn.Module) -> Any:
        """The base visiting method for mutation of nn.Module nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.Module
            The current node of nn.Module to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_effect(self, name: str, node: nn.Parameter) -> Any:
        """The base visiting method for mutation of nn.Parameter nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.Parameter
            The current node of nn.Parameter to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_param(self, name: str, node: nn.Effect) -> Any:
        """The base visiting method for mutation of nn.Effect nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.Effect
            The current node of nn.Effect to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_moduledict(self, name: str, node: nn.ModuleDict) -> Any:
        """The base visiting method for mutation of nn.ModuleDict nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.ModuleDict
            The current node of nn.ModuleDict to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_modulelist(self, name: str, node: nn.ModuleList) -> Any:
        """The base visiting method for mutation of nn.ModuleList nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.ModuleList
            The current node of nn.ModuleList to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_parameterdict(self, name: str, node: nn.ParameterDict) -> Any:
        """The base visiting method for mutation of nn.ParameterDict nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.ParameterDict
            The current node of nn.ParameterDict to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit_parameterlist(self, name: str, node: nn.ParameterList) -> Any:
        """The base visiting method for mutation of nn.ParameterList nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : nn.ParameterList
            The current node of nn.ParameterList to mutate.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """
        return self.visit(name, node)

    def visit(self, name: str, node: Any) -> Any:
        """The base dispatching method for visiting of all nodes.

        Parameters
        ----------
        name : str
            The name of the current node in parent's attribute.

        node : Any
            The current node to visit.

        Returns
        ------
        ret_node: Any
            The new node to replace current node.
        """

        def _get_child_name(parent: str, child: str) -> str:
            """Get the name of the child node/key given the parent's name."""
            if parent == "":
                # in the top level of the module
                return child
            else:
                return f"{parent}.{child}"

        if isinstance(node, nn.ParameterList):
            for i in range(len(node)):
                node[i] = self.visit_param(_get_child_name(name, str(i)), node[i])
        elif isinstance(node, nn.ParameterDict):
            for k, v in node.items():
                node[k] = self.visit_param(_get_child_name(name, k), v)
        elif isinstance(node, nn.ModuleList):
            for i in range(len(node)):
                if isinstance(node[i], nn.ParameterDict):
                    node[i] = self.visit_parameterdict(_get_child_name(name, str(i)), node[i])
                elif isinstance(node[i], nn.ParameterList):
                    node[i] = self.visit_parameterlist(_get_child_name(name, str(i)), node[i])
                elif isinstance(node[i], nn.ModuleDict):
                    node[i] = self.visit_moduledict(f"{name}.{i}", node[i])
                elif isinstance(node[i], nn.ModuleList):
                    node[i] = self.visit_modulelist(f"{name}.{i}", node[i])
                elif isinstance(node[i], nn.Module):
                    node[i] = self.visit_module(f"{name}.{i}", node[i])
                elif isinstance(node[i], nn.Effect):
                    node[i] = self.visit_effect(f"{name}.{i}", node[i])
                elif isinstance(node[i], nn.Parameter):
                    node[i] = self.visit_param(f"{name}.{i}", node[i])
        elif isinstance(node, nn.ModuleDict):
            for k, v in node.items():
                if isinstance(v, nn.ParameterDict):
                    node[k] = self.visit_parameterdict(_get_child_name(name, k), v)
                elif isinstance(v, nn.ParameterList):
                    node[k] = self.visit_parameterlist(_get_child_name(name, k), v)
                elif isinstance(v, nn.ModuleDict):
                    node[k] = self.visit_moduledict(_get_child_name(name, k), v)
                elif isinstance(v, nn.ModuleList):
                    node[k] = self.visit_modulelist(_get_child_name(name, k), v)
                elif isinstance(v, nn.Module):
                    node[k] = self.visit_module(_get_child_name(name, k), v)
                elif isinstance(v, nn.Effect):
                    node[k] = self.visit_effect(_get_child_name(name, k), v)
                elif isinstance(v, nn.Parameter):
                    node[k] = self.visit_param(_get_child_name(name, k), v)
        else:
            for key, value in node.__dict__.items():
                if isinstance(value, nn.ParameterDict):
                    setattr(node, key, self.visit_parameterdict(_get_child_name(name, key), value))
                elif isinstance(value, nn.ParameterList):
                    setattr(node, key, self.visit_parameterlist(_get_child_name(name, key), value))
                elif isinstance(value, nn.ModuleDict):
                    setattr(node, key, self.visit_moduledict(_get_child_name(name, key), value))
                elif isinstance(value, nn.ModuleList):
                    setattr(node, key, self.visit_modulelist(_get_child_name(name, key), value))
                elif isinstance(value, nn.Module):
                    setattr(node, key, self.visit_module(_get_child_name(name, key), value))
                elif isinstance(value, nn.Effect):
                    setattr(node, key, self.visit_effect(_get_child_name(name, key), value))
                elif isinstance(value, nn.Parameter):
                    setattr(node, key, self.visit_param(_get_child_name(name, key), value))
        return node
