import enum
import operator

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.quantized as nnq
import torch.nn.qat as nnqat
import torch.nn.intrinsic.qat as nniqat
toq = torch.ops.quantized

from torch.fx import GraphModule
from torch.fx.graph import Graph, Node

from .utils import getattr_from_fqn

from typing import Dict, Tuple, List, Optional, Set, Callable, Any

def _get_output_nodes(g: Graph) -> List[Node]:
    return [n for n in g.nodes if n.op == 'output']

def get_base_name_to_sets_of_related_ops() -> Dict[str, Set[Callable]]:
    base_name_to_sets_of_related_ops: Dict[str, Set[Callable]] = {
        # conv modules
        'torch.nn.Conv2d': set([
            nn.Conv2d,
            nnq.Conv2d,
            nnqat.Conv2d,
            # Note: matching weights may not work with nniqat.ConvBn2d directly
            # leaving that as a problem for a future PR to solve.
            nniqat.ConvBn2d,
        ]),
        # linear modules
        'torch.nn.Linear': set([
            nn.Linear,
            nnq.Linear,
            nnqat.Linear,
        ]),
        # linear functionals
        'torch.nn.functional.linear': set([
            F.linear,
            toq.linear,
            toq.linear_relu,
        ]),
        # add
        'torch.add': set([
            torch.add,
            toq.add,
            operator.add,  # x + y
        ]),
        # cat
        'torch.cat': set([
            torch.cat,
            toq.cat,
        ]),
        # mul
        'torch.mul': set([
            torch.mul,
            toq.mul,
        ]),
    }
    return base_name_to_sets_of_related_ops

def get_type_a_related_to_b(
    base_name_to_sets_of_related_ops: Dict[str, Set[Callable]],
) -> Set[Tuple[Callable, Callable]]:
    # TODO(future PR): allow customizations
    # TODO(future PR): reuse existing quantization mappings
    # TODO(future PR): add the rest of modules and ops here
    type_a_related_to_b: Set[Tuple[Callable, Callable]] = set()

    for base_name, s in base_name_to_sets_of_related_ops.items():
        s_list = list(s)
        # add every bidirectional pair
        for idx_0 in range(0, len(s_list) - 1):
            for idx_1 in range(idx_0 + 1, len(s_list)):
                type_a_related_to_b.add((s_list[idx_0], s_list[idx_1]))
                type_a_related_to_b.add((s_list[idx_1], s_list[idx_0]))

    return type_a_related_to_b

def get_non_matchable_functions() -> Set[Callable]:
    """
    `call_function` nodes pointing to these functions are non-matchable.
    """
    # TODO(future PR): allow customizations
    return set([
        torch.quantize_per_tensor,
    ])

def get_non_matchable_modules() -> Set[Callable]:
    """
    `call_module` nodes pointing to instances of these types are non-matchable.
    """
    # TODO(future PR): allow customizations
    return set([
        torch.quantization.ObserverBase,
        torch.quantization.FakeQuantizeBase,
    ])

def get_reversed_fusions() -> Set[Tuple[Callable, Callable]]:
    """
    Set of potential fusions, in reverse order.  The order is reversed
    to match how fusion patterns are defined in quantization code.
    """
    return set([
        (F.relu, F.linear),
    ])

# TODO(future PR): we should see if we can reuse quantization's fusion
# patterns here.
def end_node_matches_reversed_fusion(
    end_node: Node,
    reversed_fusion: Tuple[Callable, Callable],
) -> bool:
    """
    Returns true if a pattern ending with `end_node` matches
    the fusion pattern.
    """
    if end_node.op == 'call_function':
        cur_node = end_node
        for fusion_idx in range(len(reversed_fusion)):
            cur_fusion_op = reversed_fusion[fusion_idx]
            if cur_node.target != cur_fusion_op:
                return False
            if len(cur_node.args) > 0 and isinstance(cur_node.args[0], Node):
                cur_node = cur_node.args[0]
            else:
                return False
        return True
    # TODO(future PR): handle call_module
    return False


class _NSGraphMatchableSubgraphsIterator:
    """
    Iterates through the graph of gm, starting with the output nodes
    and continuing backwards.
    1. Returns matchable subgraphs, in order. A subgraph is defined by
       (start_node, end_node).
    2. Skips over non-matchable subgraphs
    """
    def __init__(
        self,
        gm: GraphModule,
        non_matchable_functions: Set[Callable],
        non_matchable_modules: Set[Callable],
    ):
        self.gm: GraphModule = gm
        self.non_matchable_functions: Set[Callable] = non_matchable_functions
        self.non_matchable_modules: Set[Callable] = non_matchable_modules
        self.seen_nodes: Set[Node] = set()
        self.stack: List[Node] = []
        for start_node in _get_output_nodes(self.gm.graph):
            self.stack.append(start_node)

    def __iter__(self):
        return self

    def __next__(self) -> Tuple[Node, Node]:
        """
        Returns the next matchable subgraph, defined by (start_node, end_node)
        """
        while len(self.stack) > 0:
            cur_end_node = self.stack.pop()
            if cur_end_node in self.seen_nodes:
                continue

            # for subgraphs which are single nodes, start_node == end_node
            # for subgraphs with more than one node, start node != end_node
            cur_start_node = cur_end_node

            # Check for potential fusions. For now, we are greedy
            # and always skip all non-base nodes of a fusion.  For example,
            # if we match linear-relu backwards, we will always skip the
            # relu node and attempt to match the linear node.  This can
            # be made configurable later if needed.
            for _reverse_fusion_ops in get_reversed_fusions():
                is_match = end_node_matches_reversed_fusion(
                    cur_end_node, _reverse_fusion_ops)
                if is_match:
                    # navigate to the base node
                    for fusion_idx in range(len(_reverse_fusion_ops) - 1):
                        self.seen_nodes.add(cur_start_node)
                        # for now, assume that there are no other nodes
                        # which need to be added to the stack
                        cur_start_node = cur_start_node.args[0]  # type: ignore
                    break

            self.seen_nodes.add(cur_start_node)
            # add args of previous nodes to stack
            # TODO(future PR): handle kwargs as needed
            for arg in cur_start_node.args:
                self._recursively_add_node_arg_to_stack(arg)

            # skip observers, etc
            # note: this check is done on the start_node, i.e.
            # if we are matching linear-relu in reverse, this would do the matchable
            # check on the linear
            if not self._is_matchable(cur_start_node):
                continue

            return cur_start_node, cur_end_node

        raise StopIteration

    def _recursively_add_node_arg_to_stack(self, arg: Any) -> None:
        """
        Adds all of the nodes in this arg to the stack, properly navigating
        through list, dicts and tuples.
        """
        if isinstance(arg, Node):
            self.stack.append(arg)
        elif isinstance(arg, torch.fx.immutable_collections.immutable_list) or type(arg) is tuple:
            for inner_arg in arg:
                self._recursively_add_node_arg_to_stack(inner_arg)
        elif isinstance(arg, torch.fx.immutable_collections.immutable_dict):
            for key, value in arg.items():
                self._recursively_add_node_arg_to_stack(value)

    def _is_matchable(self, node: Node) -> bool:
        if node.op == 'call_function':
            return not (node.target in self.non_matchable_functions)
        elif node.op == 'call_module':
            assert isinstance(node.target, str)
            # target_mod = getattr(self.gm, node.target)
            target_mod = getattr_from_fqn(self.gm, node.target)
            return not \
                any(isinstance(target_mod, t)  # type: ignore
                    for t in self.non_matchable_modules)
        else:
            return False

class GraphMatchingException(Exception):
    """
    Exception raised when two graphs cannot be matched.
    """
    pass

class NodeTypeRelationship(enum.Enum):
    # same type
    # example: F.linear and toq.linear, or nn.Conv2d and nn.Conv2d
    EQUAL = enum.auto()
    # same node_relationship set, but not the same type
    # example: F.linear and toq.linear
    RELATED_BUT_NOT_EQUAL = enum.auto()
    # not related
    NOT_RELATED = enum.auto()

def _get_node_relationship_type(
    node_a: Node,
    node_b: Node,
    gm_a: GraphModule,
    gm_b: GraphModule,
    type_a_related_to_b: Set[Tuple[Callable, Callable]],
) -> NodeTypeRelationship:
    if node_a.op != node_b.op:
        # for now, comparing call_module to call_function is not supported
        # this can be added later if needed
        return NodeTypeRelationship.NOT_RELATED

    if node_a.op == 'call_function':
        if node_a.target == node_b.target:
            # nodes with equivalent targets always match (i.e. F.linear and F.linear)
            return NodeTypeRelationship.EQUAL
        key = (node_a.target, node_b.target)
        if key in type_a_related_to_b:
            return NodeTypeRelationship.RELATED_BUT_NOT_EQUAL
        else:
            return NodeTypeRelationship.NOT_RELATED
    elif node_a.op == 'call_module':
        # for call_module, we need to look up the modules to do the type check
        assert isinstance(node_a.target, str)
        mod_a = getattr_from_fqn(gm_a, node_a.target)
        assert isinstance(node_b.target, str)
        mod_b = getattr_from_fqn(gm_b, node_b.target)
        # modules with equivalent types always match (i.e. nn.Conv2d and nn.Conv2d)
        if type(mod_a) == type(mod_b):
            return NodeTypeRelationship.EQUAL
        key = (type(mod_a), type(mod_b))
        if key in type_a_related_to_b:
            return NodeTypeRelationship.RELATED_BUT_NOT_EQUAL
        else:
            return NodeTypeRelationship.NOT_RELATED
    return NodeTypeRelationship.NOT_RELATED

def _get_name_for_subgraph(
    start_node_a: Node,
    end_node_a: Node,
    gm_a: GraphModule,
    base_name_to_sets_of_related_ops: Dict[str, Set[Callable]],
    existing_names: Set[str],
) -> str:
    """
    Returns a unique name for a subgraph. This name is based on two things:
    1. the name of the set containing the underlying type of the base op in the
       subgraph (i.e. 'torch.nn.functional.linear' if this is related to a linear op)
    2. the number of previous subgraphs with related underlying type of the base op

    For example, in the graph

    linear0 -> relu0 -> linear1 -> relu1

    The subgraphs are (linear0, relu0) and (linear1, relu1).  If we iterate
    from the output node backwards, the name given to (linear1, relu1) will be
    `base_op_torch.nn.functional.linear_0`, and the name given to (linear0, relu0)
    will be `base_op_torch.nn.functional.linear_1`.

    Why are we not just using the node name? Answer: because of two requirements:
    A. fusions must be supported
    B. some Numeric Suite APIs can be called without having all of the models in memory

    For example, let's say we need to match nodes of

    (1) ... -> linear0 -> relu0 -> ...

    And

    (2) ... -> linear_relu0 -> ...

    Without being able to inspect them together. With the current naming scheme, if
    we iterate through both of these graphs in the same order, and assuming the rest
    of the graphs match, both of these subgraphs will get the same name without
    (1) and (2) knowing anything about each other.
    """
    target_type = _get_node_target_type(start_node_a, gm_a)
    target_base_type = None
    for base_name, sets_of_related_ops in base_name_to_sets_of_related_ops.items():
        if target_type in sets_of_related_ops:
            target_base_type = base_name
    target_base_name = 'base_op_' + str(target_base_type)
    counter = 0
    proposed_name = target_base_name + '_' + str(counter)
    while proposed_name in existing_names:
        counter += 1
        proposed_name = target_base_name + '_' + str(counter)
    existing_names.add(proposed_name)
    return proposed_name

def _get_node_target_type(node: Node, gm: GraphModule) -> Optional[Callable]:
    if node.op == 'call_function':
        return node.target  # type: ignore
    elif node.op == 'call_module':
        assert isinstance(node.target, str)
        mod = getattr_from_fqn(gm, node.target)
        return type(mod)
    return None

def get_matching_subgraph_pairs(
    gm_a: GraphModule,
    gm_b: GraphModule,
) -> Dict[str, Tuple[Tuple[Node, Node], Tuple[Node, Node]]]:
    """
    Matches matchable subgraphs of graph_a to graph_b.

    For a node, "matchable" is defined as a node which is not an observer,
    fake_quants, quant or dequant.

    A subgraph can contain one or more nodes.  A subgraph is matchable if
    at least one node inside of it is matchable.  Currently, all nodes in
    a subgraph must be matchable (because we assume no observers will be
    inserted in the middle of a fusion).

    A subgraph is defined by (start_node, end_node).  We assume that only
    start_node and end_node are linked with the surrounding graph, all other
    nodes in a subgraph are self-contained.

    A pair of nodes is "related" if both nodes represent the same mathematical
    operation across different quantization flavors. For example,
    `F.linear` and `torch.ops.quantized.linear` are related, and
    `F.linear` and `torch.nn.Conv` are not related.

    For each matchable pair of nodes node_a and node_b, they will match
    if node_a and node_b are related.

    For graphs A and B, they will match iff:
    1. the number of matchable subgraphs in A and B is equivalent
    2. when iterating through the matchable subgraphs of A and B in the same order, each
       corresponding pair of base nodes is related.

    This enables us to find the corresponding subgraphs between
    graphs of related models.  For example, if we had two graphs such as:

    graph_a: x0 -> conv_0 (type: nn.Conv2d) -> obs_0 -> x1
             w  -/
             b  -/

    graph_b: x0 -> quant_0 -> qconv_0 (type: nnq.Conv2d) -> dequant_0 -> x1
           packed_params_0 -/

    This function will return the following result:
    {
        'conv_0': (  # the name of the node in graph_b
          (conv_0, conv_0),  # (start_node_a, end_node_a)
          (qconv_0, qconv_0),  # (start_node_b, end_node_b)
        ),
    }

    Or, if we have a fusion pattern,

    graph_a: x0 -> linear_0 -> relu_0 -> obs_0 -> x1
             w  -/
             b  -/

    graph_b: x0 -> quant_0 -> linear_relu_0 -> dequant_0 -> x1
           packed_params_0 -/

    This function will return the following result:
    {
        'linear_relu_0': (  # the name of the node in graph_b
          (linear_0, relu_0),  # (start_node_a, end_node_a)
          (linear_relu_0, linear_relu_0),  # (start_node_b, end_node_b)
        ),
    }
    """
    non_matchable_functions = get_non_matchable_functions()
    non_matchable_modules = get_non_matchable_modules()
    graph_a_iterator = _NSGraphMatchableSubgraphsIterator(
        gm_a, non_matchable_functions, non_matchable_modules)
    graph_b_iterator = _NSGraphMatchableSubgraphsIterator(
        gm_b, non_matchable_functions, non_matchable_modules)
    results = {}
    base_name_to_sets_of_related_ops = get_base_name_to_sets_of_related_ops()
    type_a_related_to_b = \
        get_type_a_related_to_b(base_name_to_sets_of_related_ops)

    existing_names_a: Set[str] = set()
    existing_names_b: Set[str] = set()

    while True:
        # fetch the next nodes from a and b
        cur_start_node_a, cur_start_node_b = None, None
        cur_end_node_a, cur_end_node_b = None, None
        try:
            cur_start_node_a, cur_end_node_a = next(graph_a_iterator)
        except StopIteration:
            pass
        try:
            cur_start_node_b, cur_end_node_b = next(graph_b_iterator)
        except StopIteration:
            pass

        # look up types of a and b for useful error messages
        type_a, type_b = None, None
        if cur_end_node_a is not None:
            type_a = _get_node_target_type(cur_end_node_a, gm_a)
        if cur_end_node_b is not None:
            type_b = _get_node_target_type(cur_end_node_b, gm_b)

        # check for results and determine what to do next
        if cur_end_node_a is not None and cur_end_node_b is not None:
            assert isinstance(cur_start_node_a, Node)
            assert isinstance(cur_start_node_b, Node)
            # both nodes were fetched, check for node_relationship
            # note: node_relationship is checked on the start node, i.e.
            # if a linear-relu pattern is checked, we would check for node_relationship
            # of the linear
            node_relationship = _get_node_relationship_type(
                cur_start_node_a, cur_start_node_b,
                gm_a, gm_b, type_a_related_to_b)
            if node_relationship == NodeTypeRelationship.NOT_RELATED:
                msg = f"({cur_start_node_a}, {type_a}) and ({cur_start_node_b}, {type_b}) are not related"
                raise GraphMatchingException(msg)
            elif node_relationship == NodeTypeRelationship.EQUAL:
                # For now, skip nodes with equal types. In the future, this can
                # be made configurable.
                continue
            key_name_a = _get_name_for_subgraph(
                cur_start_node_a, cur_end_node_a, gm_a, base_name_to_sets_of_related_ops,
                existing_names_a)
            key_name_b = _get_name_for_subgraph(
                cur_start_node_b, cur_end_node_b, gm_b, base_name_to_sets_of_related_ops,
                existing_names_b)
            assert key_name_a == key_name_b, \
                f"Subgraph names {key_name_a} and {key_name_b} do not match"
            results[key_name_a] = (
                (cur_start_node_a, cur_end_node_a),
                (cur_start_node_b, cur_end_node_b),
            )
            continue
        elif cur_end_node_a is None and cur_end_node_b is None:
            # we reached the end of both graphs
            break
        else:
            # only one node was fetched, no match possible, throw error
            msg = f"Matchable nodes count mismatch: ({cur_end_node_a}, {type_a}) and ({cur_end_node_b}, {type_b})"
            raise GraphMatchingException(msg)

    return results
