"""
Original from: https://github.com/jonhanke/nam_causal-head-gating
Adapted and extended for Causal Summer project by:
- Bowen
- 2026-01-27

Model adapters for accessing transformer internals across different architectures.

This module provides a unified interface for accessing transformer stages across
different HuggingFace causal language model architectures.
"""

from __future__ import annotations

from abc import ABC, abstractmethod

import torch.nn as nn


class ModelAdapter(ABC):
    """
    Abstract base class for accessing transformer model internals.

    Different model architectures organize their layers differently. This adapter
    provides a unified interface for CHG to work with any supported architecture.

    Attributes:
        model: The HuggingFace transformer model being adapted.
    """

    def __init__(self, model: nn.Module):
        """
        Initialize the adapter with a model.

        Args:
            model: A HuggingFace transformer model (e.g., from AutoModelForCausalLM).
        """
        self.model = model

    @abstractmethod
    def get_layers(self) -> nn.ModuleList:
        """
        Get the list of transformer layers.

        Returns:
            ModuleList of transformer layers.
        """
        pass

    @abstractmethod
    def get_attention_output_proj(self, layer: nn.Module) -> nn.Module:
        """
        Get the attention output projection (W_O) from a layer.

        This is the projection applied after attention heads are computed,
        which is where CHG applies its masks.

        Args:
            layer: A single transformer layer.

        Returns:
            The output projection module (typically a Linear layer).
        """
        pass

    def _resolve_layer(self, layer: int | nn.Module) -> nn.Module:
        """Resolve an integer layer index or module to an nn.Module."""
        if isinstance(layer, int):
            return self.get_layers()[layer]
        return layer

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        """Get the attention submodule from a transformer layer (or layer index)."""
        raise NotImplementedError(
            f"{self.__class__.__name__} does not expose an attention module."
        )

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        """Get the MLP/feed-forward submodule from a transformer layer (or layer index)."""
        raise NotImplementedError(
            f"{self.__class__.__name__} does not expose an MLP module."
        )

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        """Get the normalization module that feeds attention, if present."""
        return None

    def get_pre_mlp_norm(self, layer: nn.Module) -> nn.Module | None:
        """Get the normalization module that feeds the MLP, if present."""
        return None

    def get_final_norm(self) -> nn.Module | None:
        """Get the final normalization module before the LM head, if present."""
        return None

    def get_lm_head(self) -> nn.Module:
        """Get the projection module that produces token logits."""
        raise NotImplementedError(
            f"{self.__class__.__name__} does not expose an LM head module."
        )

    @property
    def num_layers(self) -> int:
        """Get the number of transformer layers."""
        return self.model.config.num_hidden_layers

    @property
    def num_heads(self) -> int:
        """Get the number of attention heads per layer."""
        return self.model.config.num_attention_heads

    @property
    def head_dim(self) -> int:
        """Get the dimension of each attention head.

        Note: For some models (e.g., Gemma-2), the attention input dimension
        (num_heads * head_dim) differs from hidden_size. In these cases,
        we use the explicit head_dim from config if available.
        """
        # Use explicit head_dim from config if available (e.g., Gemma-2)
        if hasattr(self.model.config, 'head_dim'):
            return self.model.config.head_dim
        # Fallback to computing from hidden_size
        hidden_size = self.model.config.hidden_size
        return hidden_size // self.num_heads

    @property
    def device(self):
        """Get the device the model is on."""
        return next(self.model.parameters()).device


class LlamaAdapter(ModelAdapter):
    """
    Adapter for Llama-style models.

    Supports:
        - meta-llama/Llama-2-*
        - meta-llama/Llama-3-*
        - meta-llama/Llama-3.1-*
        - meta-llama/Llama-3.2-*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.model.layers

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn.o_proj

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.input_layernorm

    def get_pre_mlp_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.post_attention_layernorm

    def get_final_norm(self) -> nn.Module | None:
        return self.model.model.norm

    def get_lm_head(self) -> nn.Module:
        return self.model.lm_head


class MistralAdapter(ModelAdapter):
    """
    Adapter for Mistral-style models.

    Supports:
        - mistralai/Mistral-*
        - mistralai/Mixtral-*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.model.layers

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn.o_proj

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.input_layernorm

    def get_pre_mlp_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.post_attention_layernorm

    def get_final_norm(self) -> nn.Module | None:
        return self.model.model.norm

    def get_lm_head(self) -> nn.Module:
        return self.model.lm_head


class GPT2Adapter(ModelAdapter):
    """
    Adapter for GPT-2 style models.

    Supports:
        - gpt2, gpt2-medium, gpt2-large, gpt2-xl
        - openai-community/gpt2*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.transformer.h

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).attn.c_proj

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).attn

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.ln_1

    def get_pre_mlp_norm(self, layer: nn.Module | None) -> nn.Module | None:
        return layer.ln_2

    def get_final_norm(self) -> nn.Module | None:
        return self.model.transformer.ln_f

    def get_lm_head(self) -> nn.Module:
        return self.model.lm_head

    @property
    def num_heads(self) -> int:
        return self.model.config.n_head


class GPTNeoXAdapter(ModelAdapter):
    """
    Adapter for GPT-NeoX style models.

    Supports:
        - EleutherAI/pythia-*
        - EleutherAI/gpt-neox-*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.gpt_neox.layers

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).attention.dense

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).attention

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.input_layernorm

    def get_pre_mlp_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.post_attention_layernorm

    def get_final_norm(self) -> nn.Module | None:
        return self.model.gpt_neox.final_layer_norm

    def get_lm_head(self) -> nn.Module:
        return self.model.embed_out


class GemmaAdapter(ModelAdapter):
    """
    Adapter for Gemma-style models.

    Supports:
        - google/gemma-*
        - google/gemma-2-*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.model.layers

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn.o_proj

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.input_layernorm

    def get_pre_mlp_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.post_attention_layernorm

    def get_final_norm(self) -> nn.Module | None:
        return self.model.model.norm

    def get_lm_head(self) -> nn.Module:
        return self.model.lm_head


class PhiAdapter(ModelAdapter):
    """
    Adapter for Phi-style models.

    Supports:
        - microsoft/phi-*
    """

    def get_layers(self) -> nn.ModuleList:
        return self.model.model.layers

    def get_attention_output_proj(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn.dense

    def get_attention_module(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).self_attn

    def get_mlp(self, layer: int | nn.Module) -> nn.Module:
        return self._resolve_layer(layer).mlp

    def get_pre_attention_norm(self, layer: nn.Module) -> nn.Module | None:
        return layer.input_layernorm

    def get_final_norm(self) -> nn.Module | None:
        return self.model.model.final_layernorm

    def get_lm_head(self) -> nn.Module:
        return self.model.lm_head


# Mapping from model architecture names to adapter classes
ADAPTER_REGISTRY: dict[str, type[ModelAdapter]] = {
    "LlamaForCausalLM": LlamaAdapter,
    "MistralForCausalLM": MistralAdapter,
    "MixtralForCausalLM": MistralAdapter,
    "GPT2LMHeadModel": GPT2Adapter,
    "GPTNeoXForCausalLM": GPTNeoXAdapter,
    "GemmaForCausalLM": GemmaAdapter,
    "Gemma2ForCausalLM": GemmaAdapter,
    "PhiForCausalLM": PhiAdapter,
    "Phi3ForCausalLM": PhiAdapter,
}

# List of supported architectures for documentation
SUPPORTED_ARCHITECTURES = list(ADAPTER_REGISTRY.keys())


def get_adapter(model: nn.Module) -> ModelAdapter:
    """
    Automatically detect the model architecture and return the appropriate adapter.

    Args:
        model: A HuggingFace transformer model.

    Returns:
        An appropriate ModelAdapter instance for the model.

    Raises:
        ValueError: If the model architecture is not supported.

    Example:
        >>> from transformers import AutoModelForCausalLM
        >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
        >>> adapter = get_adapter(model)
        >>> print(adapter.num_heads)  # 32
    """
    model_class_name = model.__class__.__name__

    if model_class_name in ADAPTER_REGISTRY:
        adapter_class = ADAPTER_REGISTRY[model_class_name]
        return adapter_class(model)

    # Try to find a matching adapter by checking the model's config
    if hasattr(model, 'config'):
        config = model.config
        model_type = getattr(config, 'model_type', None)

        # Map model_type to adapter
        model_type_mapping = {
            'llama': LlamaAdapter,
            'mistral': MistralAdapter,
            'mixtral': MistralAdapter,
            'gpt2': GPT2Adapter,
            'gpt_neox': GPTNeoXAdapter,
            'gemma': GemmaAdapter,
            'gemma2': GemmaAdapter,
            'phi': PhiAdapter,
            'phi3': PhiAdapter,
        }

        if model_type in model_type_mapping:
            adapter_class = model_type_mapping[model_type]
            return adapter_class(model)

    supported = ", ".join(SUPPORTED_ARCHITECTURES)
    raise ValueError(
        f"Unsupported model architecture: {model_class_name}. "
        f"Supported architectures: {supported}. "
        f"You can create a custom adapter by subclassing ModelAdapter."
    )
