from __future__ import annotations

from typing import Any, Union

from torch import nn
from torchview import FunctionNode, ModuleNode, TensorNode, draw_graph

from rebasin.modules import (  # type: ignore[attr-defined]
    MODULE_TYPES,
    initialize_module,
)
from rebasin.paths import LinearPath, ModelGraph, ParallelPaths

NODE_TYPES = Union[ModuleNode, TensorNode, FunctionNode]


class PermutationInitializer:
    def __init__(
            self,
            model_a: nn.Module,
            model_b: nn.Module,
            input_data_b: Any,
            input_data_a: Any | None = None,
    ) -> None:
        self.model_a = model_a
        self.model_b = model_b
        self.input_data_b = input_data_b
        self.input_data_a = self.input_data_b if input_data_a is None else input_data_a

        self.id_to_module_a = {id(m): m for m in self.model_a.modules()}
        self.id_to_module_b = {id(m): m for m in self.model_b.modules()}
        self._err_msg = "A and B must have the same model architecture."

        self.model_graph = self.initialize_permutations()

    def initialize_permutations(self) -> ModelGraph:
        """Initialize the permutations of the model.
        """

        nextnodes_a: list[NODE_TYPES] = list(
            draw_graph(self.model_a, input_data=self.input_data_a, depth=1e12)
            .root_container
        )
        nextnodes_b: list[NODE_TYPES] = list(
            draw_graph(self.model_b, input_data=self.input_data_b, depth=1e12)
            .root_container
        )

        paths: list[LinearPath | ParallelPaths] = []

        while len(nextnodes_a) > 0 or len(nextnodes_b) > 0:
            assert len(nextnodes_a) == len(nextnodes_b), self._err_msg

            path: LinearPath | ParallelPaths
            if len(nextnodes_b) == 1:
                path, nextnodes_a, nextnodes_b = self.initialize_linear_path(
                    nextnodes_a, nextnodes_b
                )
            else:
                path, nextnodes_a, nextnodes_b = self.initialize_parallel_paths(
                    nextnodes_a, nextnodes_b
                )

            if bool(path):
                # This might be False (path empty) if we, for example,
                #   start with a nn.MultiheadAttention module.
                # Then, the three input tensors would yield a ParallelPaths,
                #   but it would be empty because those tensors combine into
                #   a single Module, the nn.MultiheadAttention.
                # In this case, we aren't interested in the empty ParallelPaths,
                #   only in the subsequent LinearPath (and what follows it).
                paths.append(path)

        return ModelGraph(*paths)

    def initialize_linear_path(
            self, nextnodes_a: list[NODE_TYPES], nextnodes_b: list[NODE_TYPES]
    ) -> tuple[LinearPath, list[NODE_TYPES], list[NODE_TYPES]]:

        modules: list[MODULE_TYPES] = []
        while 0 < len(nextnodes_a) < 2:
            assert len(nextnodes_a) == len(nextnodes_b) == 1, self._err_msg

            node_a, node_b = nextnodes_a.pop(), nextnodes_b.pop()
            assert type(node_a) == type(node_b), self._err_msg

            if isinstance(node_a, ModuleNode) and isinstance(node_b, ModuleNode):
                mod_a = self.id_to_module_a[node_a.compute_unit_id]
                mod_b = self.id_to_module_b[node_b.compute_unit_id]
                module = initialize_module(mod_a, mod_b)
                if module is not None:
                    modules.append(module)

            nextnodes_a = list(node_a.children)  # type: ignore[arg-type]
            nextnodes_b = list(node_b.children)  # type: ignore[arg-type]

            # If there are multiple parents in the next nodes,
            #   then we have reached a fork in the graph.
            # This is likely because initialize_linear_path was called by
            #   initialize_parallel_paths, and we reached the end of the parallel paths.
            if nextnodes_a and len(nextnodes_a[0].parents) > 1:
                assert nextnodes_b and len(nextnodes_b[0].parents) > 1, self._err_msg
                break

        return LinearPath(*modules), nextnodes_a, nextnodes_b

    def initialize_parallel_paths(
            self, nextnodes_a: list[NODE_TYPES], nextnodes_b: list[NODE_TYPES]
    ) -> tuple[ParallelPaths, list[NODE_TYPES], list[NODE_TYPES]]:
        paths: list[LinearPath] = []

        assert len(nextnodes_a) == len(nextnodes_b), self._err_msg

        finalnodes_a: list[NODE_TYPES] = []
        finalnodes_b: list[NODE_TYPES] = []
        prev_final_nodes_a: list[NODE_TYPES] = []
        prev_final_nodes_b: list[NODE_TYPES] = []

        for node_a, node_b in zip(nextnodes_a, nextnodes_b):  # noqa: B905
            path, finalnodes_a, finalnodes_b = self.initialize_linear_path(
                [node_a], [node_b]
            )
            assert len(finalnodes_a) == len(finalnodes_b), self._err_msg

            # All LinearPaths in this ParallelPaths must end in the same modules.
            if not prev_final_nodes_a:
                prev_final_nodes_a = finalnodes_a
                prev_final_nodes_b = finalnodes_b

            # assert prev_final_nodes_a == finalnodes_a, (
            #     "This network architecture is not supported! "
            #     "Please look at README.md for more information."
            # )
            # assert prev_final_nodes_b == finalnodes_b, (
            #     "This network architecture is not supported! "
            #     "Please look at README.md for more information."
            # )

            paths.append(path)

        return ParallelPaths(*paths), finalnodes_a, finalnodes_b
