from __future__ import annotations

from collections import defaultdict
from contextlib import AbstractContextManager
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch.utils.hooks import RemovableHandle

from .. import util
from .graph import InterventionGraph
from .protocols import InterventionProtocol

class Interleaver(AbstractContextManager):
    
    """The Interleaver is responsible for executing a function involving a PyTorch model alongside a user's custom functionality
    (represented by an `InterventionGraph`). This is called interleaving.
    
    The `InterventionGraph` has information about which components (modules) of the model the user's custom logic will interact with.
    As the `Interleaver` is a context, entering it adds the appropriate hooks to these components which act as a bridge between the model's
    original computation graph and the `InterventionGraph`. Exiting the `Interleaver` removes these hooks.
    
    Attributes:
        graph (InterventionGraph): The computation graph representing the user's custom intervention logic.
        batch_groups (Optional[List[Tuple[int, int]]]): A batch group is a section of tensor values related to a given intervention. 
            They are a tuple of (batch_start, batch_length). So if batch group 0 was (0, 4) it means it starts at index 0 and goes until index 3.
            The batch index is assumed to be the first dimension of all Tensors.
            InterventionProtocol Nodes know which batch group they are a part of in their arguments. That value is the index into the batch_groups.
        input_hook (Optional[Callable]). Function to hook onto the inputs of modules for interleaving. Defaults to None and therefore `InterventionProtocol.intervene`.
        output_hook (Optional[Callable]). Function to hook onto the outputs of modules for interleaving. Defaults to None and therefore `InterventionProtocol.intervene`.
        batch_size (Optional[int]). Total batch size. Used to determine which Tensors need to be narrowed to their batch group.
            i.e If a Tensor's first dimension isn't batch_size, we dont need to narrow it to convert it for its batch_group.
            Defaults to None and therefore the sum of the last batch_group.
    """

    def __init__(
        self,
        graph: InterventionGraph,
        batch_groups: Optional[List[Tuple[int, int]]] = None,
        input_hook: Optional[Callable] = None,
        output_hook: Optional[Callable] = None,
        batch_size: Optional[int] = None,
    ) -> None:

        self.graph = graph

        self.batch_groups = [] if batch_groups is None else batch_groups

        if input_hook is None:
            input_hook = (
                lambda activations, module_path, module: InterventionProtocol.intervene(
                    activations, module_path, module, "input", self
                )
            )

        if output_hook is None:
            output_hook = (
                lambda activations, module_path, module: InterventionProtocol.intervene(
                    activations, module_path, module, "output", self
                )
            )

        self.input_hook = input_hook
        self.output_hook = output_hook

        self.handles: List[RemovableHandle] = []

        self.batch_size = (
            sum(self.batch_groups[-1]) if batch_size is None else batch_size
        )

    def __enter__(self) -> Interleaver:
        """Registers input and output hooks to modules involved in the `InterventionGraph`.

        Returns:
            Interleaver: Interleaver
        """

        # Keys of `InterventionGraph.interventions` are the module paths + if they are for input or output.
        # e.x 'transformer.h.0.mlp.output'
        for module_key in self.graph.interventions.keys():

            module_atoms = module_key.split(".")

            # Get just the hook type i.e input/output
            *module_atoms, hook_type = module_atoms

            # Get just the module path 
            module_path = ".".join(module_atoms)

            # Get the torch module using the module_path
            module: torch.nn.Module = util.fetch_attr(self.graph.model, module_path)

            if hook_type == "input":

                # Input hook activations are a tuple of (positional args, key-word arguments)
                # Include the module_path not the module
                def input_hook(module, input, kwargs, module_path=module_path):
                    return self.input_hook((input, kwargs), module_path, module)

                self.handles.append(
                    module.register_forward_pre_hook(
                        input_hook, with_kwargs=True, prepend=True
                    )
                )

            elif hook_type == "output":

                def output_hook(module, input, output, module_path=module_path):
                    return self.output_hook(output, module_path, module)

                self.handles.append(
                    module.register_forward_hook(output_hook, prepend=True)
                )

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:

        # Remove all hooks
        for handle in self.handles:
            handle.remove()

        if isinstance(exc_val, Exception):
            raise exc_val



