from contextlib import contextmanager
from dataclasses import dataclass, replace
from typing import Any, Callable, Generator, Optional, overload

import torch
from torch import Tensor, nn
from torch.utils.hooks import RemovableHandle

from .layer_matching import (
    LayerType,
    ModelLayerConfig,
    collect_matching_layers,
    guess_and_enhance_layer_config,
)
from .torch_utils import get_module, untuple_tensor

PatchOperator = Callable[[Tensor, Tensor], Tensor]


@dataclass
class SteeringPatchHandle:
    """
    A handle that can be used to remove a steering patch from a model after
    running `steering_vector.patch_activations()`.
    """

    model_hooks: list[RemovableHandle]

    def remove(self) -> None:
        """Remove the steering patch from the model"""
        for hook in self.model_hooks:
            hook.remove()


@dataclass
class SteeringVector:
    """A steering vector that can be applied to a model."""

    layer_activations: dict[int, Tensor]
    layer_type: LayerType = "decoder_block"

    def patch_activations(
        self,
        model: nn.Module,
        layer_config: Optional[ModelLayerConfig] = None,
        operator: Optional[PatchOperator] = None,
        multiplier: float = 1.0,
        min_token_index: int | None = None,
        token_indices: list[int] | slice | Tensor | None = None,
    ) -> SteeringPatchHandle:
        """
        Patch the activations of the given model with this steering vector.
        This will modify the model in-place, and return a handle that can be used to undo the patching.
        This method does the same thing as `apply`, but requires manually undoing the patching to
        restore the model to its original state. For most cases, `apply` is easier to use. Tokens to patch
        can be selected using either `min_token_index` or `token_indices`, but not both. If neither is provided,
        all tokens will be patched.

        Args:
            model: The model to patch
            layer_config: A dictionary mapping layer types to layer matching functions.
                If not provided, this will be inferred automatically.
            operator: A function that takes the original activation and the steering vector
                and returns a modified vector that is added to the original activation.
            multiplier: A multiplier to scale the patch activations. Default is 1.0.
            min_token_index: The minimum token index to apply the patch to. Default is None.
            token_indices: Either a list of token indices to apply the patch to, a slice, or a mask tensor. Default is None.
        Example:
            >>> model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
            >>> steering_vector = SteeringVector(...)
            >>> handle = steering_vector.patch_activations(model)
            >>> model.forward(...)
            >>> handle.remove()
        """
        assert (min_token_index is None) or (
            token_indices is None
        ), "Can not pass both min_token_index and token_indices"
        if isinstance(token_indices, Tensor):
            assert torch.all(
                torch.logical_or(token_indices == 0, token_indices == 1)
            ), "token_indices tensor must be a mask (containing only 0s and 1s)"
        token_indices = (
            token_indices if token_indices is not None else slice(min_token_index, None)
        )
        layer_config = guess_and_enhance_layer_config(
            model, layer_config, self.layer_type
        )
        hooks: list[RemovableHandle] = []
        if self.layer_type not in layer_config:
            raise ValueError(
                f"layer_type {self.layer_type} not provided in layer config"
            )
        matcher = layer_config[self.layer_type]
        matching_layers = collect_matching_layers(model, matcher)
        for layer_num, target_activation in self.layer_activations.items():
            layer_name = matching_layers[layer_num]

            target_activation = multiplier * self.layer_activations[layer_num]

            module = get_module(model, layer_name)
            handle = module.register_forward_hook(
                # create the hook via function call since python only creates new scopes on functions
                _create_additive_hook(
                    target_activation.reshape(1, 1, -1), token_indices, operator
                )
            )
            hooks.append(handle)
        return SteeringPatchHandle(hooks)

    @contextmanager
    def apply(
        self,
        model: nn.Module,
        layer_config: Optional[ModelLayerConfig] = None,
        operator: Optional[PatchOperator] = None,
        multiplier: float = 1.0,
        min_token_index: int = 0,
        token_indices: list[int] | slice | Tensor | None = None,
    ) -> Generator[None, None, None]:
        """
        Apply this steering vector to the given model. Tokens to patch
        can be selected using either `min_token_index` or `token_indices`, but not both.
        If neither is provided, all tokens will be patched.

        Args:
            model: The model to patch
            layer_config: A dictionary mapping layer types to layer matching functions.
                If not provided, this will be inferred automatically.
            operator: A function that takes the original activation and the steering vector
                and returns a modified vector that is added to the original activation.
            multiplier: A multiplier to scale the patch activations. Default is 1.0.
            min_token_index: The minimum token index to apply the patch to. Default is None.
            token_indices: Either a list of token indices to apply the patch to, a slice, or a mask tensor. Default is None.
        Example:
            >>> model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
            >>> steering_vector = SteeringVector(...)
            >>> with steering_vector.apply(model):
            >>>     model.forward(...)
        """
        try:
            handle = self.patch_activations(
                model=model,
                layer_config=layer_config,
                operator=operator,
                multiplier=multiplier,
                min_token_index=min_token_index,
                token_indices=token_indices,
            )
            yield
        finally:
            handle.remove()

    # types copied from torch.Tensor

    @overload
    def to(
        self, dtype: torch.dtype, non_blocking: bool = False, copy: bool = False
    ) -> "SteeringVector":
        ...

    @overload
    def to(
        self,
        device: Optional[torch.device | str] = None,
        dtype: Optional[torch.dtype] = None,
        non_blocking: bool = False,
        copy: bool = False,
    ) -> "SteeringVector":
        ...

    @overload
    def to(
        self, other: Tensor, non_blocking: bool = False, copy: bool = False
    ) -> "SteeringVector":
        ...

    def to(self, *args: Any, **kwargs: Any) -> "SteeringVector":
        """
        Return a new steering vector moved to the given device/dtype.

        This method calls ``torch.Tensor.to`` on each of the layer activations.
        """
        layer_activations = {
            layer_num: act.to(*args, **kwargs)
            for layer_num, act in self.layer_activations.items()
        }
        return replace(self, layer_activations=layer_activations)


def _create_additive_hook(
    target_activation: Tensor,
    token_indices: list[int] | slice | Tensor,
    operator: PatchOperator | None = None,
) -> Any:
    """Create a hook function that adds the given target_activation to the model output"""

    def hook_fn(_m: Any, _inputs: Any, outputs: Any) -> Any:
        original_tensor = untuple_tensor(outputs)
        act = target_activation.to(original_tensor.device)
        delta = act
        if operator is not None:
            delta = operator(original_tensor, act)
        if isinstance(token_indices, Tensor):
            mask = token_indices
        else:
            mask = torch.zeros(original_tensor.shape[1])
            mask[token_indices] = 1
        mask = mask.reshape(1, -1, 1)
        mask = mask.to(original_tensor.device)
        original_tensor[None] = original_tensor + (mask * delta)
        return outputs

    return hook_fn
