"""Steering Vectors for inference-time refusal intervention.

The existing OBLITERATUS pipeline only supports permanent weight modification.
Steering vectors provide a complementary approach: modifying activations at
inference time without changing any weights.

This is based on:
  - Turner et al. (2023): "Activation Addition: Steering Language Models
    Without Optimization"
  - Rimsky et al. (2024): "Steering Llama 2 via Contrastive Activation
    Addition" (CAA)
  - Li et al. (2024): "Inference-Time Intervention: Eliciting Truthful
    Answers from a Language Model"

The approach:
  1. Compute a steering vector from the refusal direction (or any
     concept direction)
  2. At inference time, add or subtract scaled multiples of the vector
     to the residual stream at specified layers
  3. This steers the model toward or away from refusal without modifying
     any weights

Advantages over weight projection:
  - **Reversible**: Steering can be turned on/off per-request
  - **Tunable**: The steering strength (alpha) can be adjusted continuously
  - **Composable**: Multiple steering vectors can be combined
  - **Non-destructive**: Model weights are never modified

Limitations (vs. weight projection):
  - Requires wrapping the model's forward pass (hooks)
  - Slight inference-time overhead per token
  - Effect is per-token, not permanent

This module provides:
  1. SteeringVector construction from refusal directions or contrastive pairs
  2. Hook-based application to any HuggingFace model
  3. Multi-layer steering with per-layer alpha scaling
  4. Evaluation utilities for measuring steering effectiveness

References:
    - Turner et al. (2023): Activation Addition (arXiv:2308.10248)
    - Rimsky et al. (2024): Contrastive Activation Addition for Llama 2
    - Li et al. (2024): Inference-Time Intervention (arXiv:2306.03341)
"""

from __future__ import annotations

from dataclasses import dataclass, field

import torch
import torch.nn as nn


@dataclass
class SteeringVector:
    """A steering vector that can be applied at inference time."""

    direction: torch.Tensor          # (hidden_dim,) unit vector
    source_layer: int | None         # layer it was extracted from (None if synthetic)
    label: str                       # human-readable label (e.g. "refusal", "truthfulness")
    default_alpha: float             # recommended steering strength
    metadata: dict = field(default_factory=dict)


@dataclass
class SteeringConfig:
    """Configuration for inference-time steering."""

    vectors: list[SteeringVector]
    target_layers: list[int]         # which layers to steer at
    alpha: float = 1.0               # global scaling factor
    per_layer_alpha: dict[int, float] = field(default_factory=dict)  # per-layer overrides
    position: str = "all"            # "all", "last", or "first" — which positions to steer
    normalize: bool = True           # normalize vector to unit norm before scaling


@dataclass
class SteeringResult:
    """Result of applying steering vectors."""

    config: SteeringConfig
    hooks_installed: int
    total_steered_layers: int


class SteeringVectorFactory:
    """Create steering vectors from various sources."""

    @staticmethod
    def from_refusal_direction(
        refusal_direction: torch.Tensor,
        source_layer: int | None = None,
        alpha: float = -1.0,
    ) -> SteeringVector:
        """Create a steering vector from a pre-computed refusal direction.

        By default, alpha=-1.0 steers AWAY from refusal (removes it).
        Use alpha=+1.0 to steer TOWARD refusal (reinforces it).

        Args:
            refusal_direction: (hidden_dim,) refusal direction vector.
            source_layer: Layer the direction was extracted from.
            alpha: Steering strength. Negative = away from refusal.

        Returns:
            SteeringVector ready for application.
        """
        d = refusal_direction.float().squeeze()
        d = d / d.norm().clamp(min=1e-10)
        return SteeringVector(
            direction=d,
            source_layer=source_layer,
            label="refusal",
            default_alpha=alpha,
        )

    @staticmethod
    def from_contrastive_pairs(
        positive_activations: list[torch.Tensor],
        negative_activations: list[torch.Tensor],
        label: str = "contrastive",
        alpha: float = 1.0,
    ) -> SteeringVector:
        """Create a steering vector from contrastive activation pairs.

        The vector is the difference in mean activations:
          vector = mean(positive) - mean(negative)

        Args:
            positive_activations: Activations from "positive" examples
                (e.g., harmful prompts that trigger refusal).
            negative_activations: Activations from "negative" examples
                (e.g., harmless prompts without refusal).
            label: Human-readable label.
            alpha: Default steering strength.

        Returns:
            SteeringVector from contrastive difference.
        """
        pos_mean = torch.stack([a.float().squeeze() for a in positive_activations]).mean(dim=0)
        neg_mean = torch.stack([a.float().squeeze() for a in negative_activations]).mean(dim=0)
        diff = pos_mean - neg_mean
        d = diff / diff.norm().clamp(min=1e-10)
        return SteeringVector(
            direction=d,
            source_layer=None,
            label=label,
            default_alpha=alpha,
            metadata={"n_positive": len(positive_activations),
                       "n_negative": len(negative_activations),
                       "raw_magnitude": diff.norm().item()},
        )

    @staticmethod
    def combine(
        vectors: list[SteeringVector],
        weights: list[float] | None = None,
        label: str = "combined",
    ) -> SteeringVector:
        """Combine multiple steering vectors into one.

        Args:
            vectors: List of SteeringVector to combine.
            weights: Per-vector weights. If None, equal weights.
            label: Label for the combined vector.

        Returns:
            Combined SteeringVector.
        """
        if not vectors:
            raise ValueError("Need at least one vector to combine")

        if weights is None:
            weights = [1.0 / len(vectors)] * len(vectors)

        combined = sum(w * v.direction for w, v in zip(weights, vectors))
        combined = combined / combined.norm().clamp(min=1e-10)

        mean_alpha = sum(v.default_alpha for v in vectors) / len(vectors)

        return SteeringVector(
            direction=combined,
            source_layer=None,
            label=label,
            default_alpha=mean_alpha,
            metadata={"n_combined": len(vectors), "weights": weights},
        )


class SteeringHookManager:
    """Manages inference-time hooks for applying steering vectors.

    This class installs PyTorch forward hooks on specified layers
    to add/subtract steering vectors from the residual stream.
    """

    def __init__(self):
        self._hooks: list = []
        self._active = False

    def install(
        self,
        model: nn.Module,
        config: SteeringConfig,
        layer_modules: list[nn.Module] | None = None,
    ) -> SteeringResult:
        """Install steering hooks on the model.

        Args:
            model: The transformer model.
            config: SteeringConfig specifying vectors, layers, and alphas.
            layer_modules: If provided, use these as the layer modules.
                Otherwise, attempts to find them automatically.

        Returns:
            SteeringResult with installation details.
        """
        self.remove()  # Clean up any existing hooks

        if layer_modules is None:
            layer_modules = self._find_layer_modules(model)

        n_installed = 0
        for layer_idx in config.target_layers:
            if layer_idx >= len(layer_modules):
                continue

            module = layer_modules[layer_idx]
            alpha = config.per_layer_alpha.get(layer_idx, config.alpha)

            hook = self._make_hook(config.vectors, alpha, config.position, config.normalize)
            handle = module.register_forward_hook(hook)
            self._hooks.append(handle)
            n_installed += 1

        self._active = True

        return SteeringResult(
            config=config,
            hooks_installed=n_installed,
            total_steered_layers=n_installed,
        )

    def remove(self):
        """Remove all installed hooks."""
        for handle in self._hooks:
            handle.remove()
        self._hooks.clear()
        self._active = False

    @property
    def is_active(self) -> bool:
        return self._active

    def _make_hook(
        self,
        vectors: list[SteeringVector],
        alpha: float,
        position: str,
        normalize: bool,
    ):
        """Create a forward hook that applies steering vectors."""
        def hook(module, input, output):
            # output is typically (hidden_states, ...) or just hidden_states
            if isinstance(output, tuple):
                hidden = output[0]
                rest = output[1:]
            else:
                hidden = output
                rest = None

            for vec in vectors:
                d = vec.direction.to(hidden.device, hidden.dtype)
                if normalize:
                    d = d / d.norm().clamp(min=1e-10)

                scale = alpha * vec.default_alpha
                steering = scale * d

                if hidden.ndim == 3:
                    # (batch, seq_len, hidden_dim) — typical transformer output
                    if position == "last":
                        hidden = hidden.clone()
                        hidden[:, -1, :] = hidden[:, -1, :] + steering
                    elif position == "first":
                        hidden = hidden.clone()
                        hidden[:, 0, :] = hidden[:, 0, :] + steering
                    else:  # "all"
                        hidden = hidden + steering.unsqueeze(0).unsqueeze(0)
                elif hidden.ndim == 2:
                    # (batch, hidden_dim) — e.g., linear layer output
                    hidden = hidden + steering.unsqueeze(0)
                else:
                    # Unsupported shape — add along last dim as best effort
                    hidden = hidden + steering

            if rest is not None:
                return (hidden,) + rest
            return hidden

        return hook

    @staticmethod
    def _find_layer_modules(model: nn.Module) -> list[nn.Module]:
        """Auto-detect transformer layer modules."""
        # Common attribute paths for transformer layers
        for attr_path in [
            "model.layers", "transformer.h", "gpt_neox.layers",
            "model.decoder.layers", "encoder.layer",
        ]:
            obj = model
            try:
                for part in attr_path.split("."):
                    obj = getattr(obj, part)
                return list(obj)
            except AttributeError:
                continue
        return []


def compute_steering_effectiveness(
    clean_projection: float,
    steered_projection: float,
    direction: str = "remove",
) -> float:
    """Compute how effective steering was.

    Args:
        clean_projection: Refusal projection without steering.
        steered_projection: Refusal projection with steering.
        direction: "remove" (want to reduce) or "add" (want to increase).

    Returns:
        Effectiveness score (0-1). 1.0 = perfectly effective.
    """
    if direction == "remove":
        if abs(clean_projection) < 1e-10:
            return 1.0  # Already no refusal
        return max(0.0, 1.0 - abs(steered_projection) / abs(clean_projection))
    else:
        if abs(steered_projection) < 1e-10:
            return 0.0
        return min(1.0, abs(steered_projection) / max(abs(clean_projection), 1e-10))


def format_steering_report(result: SteeringResult) -> str:
    """Format steering application report."""
    lines = []
    lines.append("Steering Vector Application")
    lines.append("=" * 35)
    lines.append("")
    lines.append(f"Hooks installed: {result.hooks_installed}")
    lines.append(f"Layers steered: {result.total_steered_layers}")
    lines.append(f"Global alpha: {result.config.alpha}")
    lines.append(f"Position mode: {result.config.position}")
    lines.append(f"Vectors applied: {len(result.config.vectors)}")
    for v in result.config.vectors:
        lines.append(f"  - {v.label} (alpha={v.default_alpha:+.2f}, dim={v.direction.shape[0]})")
    return "\n".join(lines)
