"""
DualSteeringSteeringObject — inference-time Newton steering via softmax geometry.

When ``set_model_weights()`` has been called (by WisentModel.attach()),
``apply_steering()`` runs full Newton iterations using the lm_head weight
matrix. Otherwise, it uses additive steering (h + strength * dir).

The lm_head weight is **not** serialised — it is injected at runtime.
"""

from __future__ import annotations

from typing import Any, Dict, Optional
import torch
import torch.nn.functional as F

from wisent.core.steering_methods.steering_object import (
    BaseSteeringObject,
    SteeringObjectMetadata,
)
from .newton import dual_newton_step

__all__ = ["DualSteeringSteeringObject"]


class DualSteeringSteeringObject(BaseSteeringObject):
    """Steering object that applies Newton steps in softmax geometry."""

    method_name = "dual_steering"

    def __init__(
        self,
        metadata: SteeringObjectMetadata,
        concept_directions: Dict[int, torch.Tensor],
        num_steps: int = 3,
        alpha: float = 5e-3,
        eta: float = 0.5,
        top_k: int = 256,
    ):
        super().__init__(metadata)
        self.concept_directions = concept_directions
        self.num_steps = num_steps
        self.alpha = alpha
        self.eta = eta
        self.top_k = top_k
        self._lm_head_weight: Optional[torch.Tensor] = None

    def set_model_weights(self, model: Any) -> None:
        """
        Extract and cache the lm_head weight matrix from the HF model.

        Called by ``WisentModel.apply_steering_object()`` right after
        ``detach()``. The weight is stored as a detached float tensor
        and is NOT serialised.
        """
        lm_head = getattr(model, "lm_head", None)
        if lm_head is None:
            lm_head = getattr(model, "embed_out", None)
        if lm_head is None:
            for name, module in model.named_modules():
                if "lm_head" in name or "embed_out" in name:
                    lm_head = module
                    break
        if lm_head is not None:
            weight = getattr(lm_head, "weight", None)
            if weight is not None:
                self._lm_head_weight = weight.detach().float()

    def get_steering_vector(
        self, layer: int, hidden_state: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Compatibility: return normalised concept direction for layer."""
        if layer not in self.concept_directions:
            raise KeyError(f"No concept direction for layer {layer}")
        return F.normalize(self.concept_directions[layer].float(), p=2, dim=-1)

    def compute_gate(self, hidden_state: torch.Tensor) -> torch.Tensor:
        batch_size = hidden_state.shape[0] if hidden_state.dim() > 1 else 1
        return torch.ones(
            batch_size, device=hidden_state.device, dtype=hidden_state.dtype
        )

    def compute_intensity(
        self, hidden_state: torch.Tensor, layer: int
    ) -> torch.Tensor:
        batch_size = hidden_state.shape[0] if hidden_state.dim() > 1 else 1
        return torch.ones(
            batch_size, device=hidden_state.device, dtype=hidden_state.dtype
        )

    def apply_steering(
        self,
        hidden_state: torch.Tensor,
        layer: int,
        base_strength: float = 1.0,
    ) -> torch.Tensor:
        """
        Apply dual steering.

        If lm_head weights are available: full Newton iteration.
        Otherwise: additive steering (h + strength * concept_dir).
        """
        if layer not in self.concept_directions:
            return hidden_state

        concept_dir = self.concept_directions[layer].float()
        original_shape = hidden_state.shape
        original_dtype = hidden_state.dtype

        if self._lm_head_weight is not None:
            return self._apply_newton(
                hidden_state, concept_dir, base_strength,
                original_shape, original_dtype,
            )
        return self._apply_additive(
            hidden_state, concept_dir, base_strength,
            original_shape, original_dtype,
        )

    def _apply_newton(
        self,
        hidden_state: torch.Tensor,
        concept_dir: torch.Tensor,
        base_strength: float,
        original_shape: torch.Size,
        original_dtype: torch.dtype,
    ) -> torch.Tensor:
        """Full Newton iteration using lm_head weights."""
        lm_w = self._lm_head_weight.to(hidden_state.device)
        cd = concept_dir.to(hidden_state.device)

        if hidden_state.dim() == 1:
            h = hidden_state.unsqueeze(0).float()
        elif hidden_state.dim() == 2:
            h = hidden_state.float()
        else:
            b, s, hd = hidden_state.shape
            h = hidden_state.reshape(b * s, hd).float()

        h_new = dual_newton_step(
            h, lm_w, cd * base_strength,
            num_steps=self.num_steps,
            alpha=self.alpha,
            eta=self.eta,
            top_k=self.top_k,
        )
        h_new = h_new.to(original_dtype)

        if hidden_state.dim() == 1:
            return h_new.squeeze(0)
        if hidden_state.dim() == 3:
            return h_new.reshape(original_shape)
        return h_new

    def _apply_additive(
        self,
        hidden_state: torch.Tensor,
        concept_dir: torch.Tensor,
        base_strength: float,
        original_shape: torch.Size,
        original_dtype: torch.dtype,
    ) -> torch.Tensor:
        """Additive steering when lm_head is not available."""
        cd = F.normalize(concept_dir, p=2, dim=-1).to(
            device=hidden_state.device, dtype=original_dtype
        )
        return hidden_state + base_strength * cd

    def to_dict(self) -> Dict[str, Any]:
        """Serialise to dictionary (lm_head NOT included)."""
        meta = self.metadata
        return {
            "method": self.method_name,
            "metadata": {
                "method": meta.method,
                "model_name": meta.model_name,
                "benchmark": meta.benchmark,
                "category": meta.category,
                "extraction_strategy": meta.extraction_strategy,
                "num_pairs": meta.num_pairs,
                "layers": meta.layers,
                "hidden_dim": meta.hidden_dim,
                "created_at": meta.created_at,
                "extra": meta.extra,
                "calibration_norms": {
                    str(k): v for k, v in meta.calibration_norms.items()
                },
                "extraction_component": meta.extraction_component,
            },
            "concept_directions": {
                str(k): v for k, v in self.concept_directions.items()
            },
            "num_steps": self.num_steps,
            "alpha": self.alpha,
            "eta": self.eta,
            "top_k": self.top_k,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DualSteeringSteeringObject":
        """Deserialise from dictionary."""
        meta_data = data["metadata"]
        cal_raw = meta_data.get("calibration_norms", {})
        calibration_norms = {int(k): float(v) for k, v in cal_raw.items()}
        metadata = SteeringObjectMetadata(
            method=meta_data["method"],
            model_name=meta_data["model_name"],
            benchmark=meta_data["benchmark"],
            category=meta_data["category"],
            extraction_strategy=meta_data["extraction_strategy"],
            num_pairs=meta_data["num_pairs"],
            layers=meta_data["layers"],
            hidden_dim=meta_data["hidden_dim"],
            created_at=meta_data.get("created_at", ""),
            extra=meta_data.get("extra", {}),
            calibration_norms=calibration_norms,
            extraction_component=meta_data.get(
                "extraction_component", "residual_stream"
            ),
        )

        def to_tensor(v):
            return torch.tensor(v) if isinstance(v, list) else v

        concept_dirs = {
            int(k): to_tensor(v)
            for k, v in data["concept_directions"].items()
        }

        return cls(
            metadata=metadata,
            concept_directions=concept_dirs,
            num_steps=data.get("num_steps", 3),
            alpha=data.get("alpha", 5e-3),
            eta=data.get("eta", 0.5),
            top_k=data.get("top_k", 256),
        )
