"""
Dual Steering method — information-geometric steering via softmax m-geodesics.

Training computes concept directions (mean_pos - mean_neg) per layer,
identical to CAA. The geometry-aware Newton iteration happens at inference
time inside ``DualSteeringSteeringObject.apply_steering()``, which requires
access to the model's lm_head weight matrix.

Reference: Park et al. 2025 — "The Information Geometry of Softmax:
Probing and Steering".
"""

from __future__ import annotations

from typing import Any, Dict, List, Tuple
from dataclasses import dataclass, field
from collections import defaultdict

import torch
import torch.nn.functional as F

from wisent.core.steering_methods.core.atoms import BaseSteeringMethod
from wisent.core.activations.core.atoms import (
    LayerActivations,
    RawActivationMap,
    LayerName,
)
from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
from wisent.core.errors import InsufficientDataError

__all__ = [
    "DualSteeringMethod",
    "DualSteeringConfig",
    "DualSteeringResult",
]


@dataclass
class DualSteeringConfig:
    """Configuration for Dual Steering method."""

    num_steps: int = 3
    """Number of Newton iterations per forward pass."""
    alpha: float = 5e-3
    """Tikhonov regularisation for covariance inversion."""
    eta: float = 0.5
    """Step-size multiplier per Newton iteration."""
    top_k: int = 256
    """Number of top-probability tokens retained for covariance."""


@dataclass
class DualSteeringResult:
    """Result from Dual Steering training with per-layer concept directions."""

    concept_directions: Dict[LayerName, torch.Tensor]
    """Per-layer concept direction [D] (mean_pos - mean_neg)."""
    config: DualSteeringConfig
    layer_order: List[LayerName]
    metadata: Dict[str, Any] = field(default_factory=dict)


class DualSteeringMethod(BaseSteeringMethod):
    """
    Dual Steering method.

    Training is identical to CAA: compute mean(pos) - mean(neg) per layer.
    The geometry-aware Newton iteration is deferred to inference time.
    """

    name = "dual_steering"
    description = (
        "Information-geometric steering via softmax m-geodesics — "
        "Newton steps that minimise off-target KL divergence"
    )

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.config = DualSteeringConfig(
            num_steps=kwargs.get("num_steps", 3),
            alpha=kwargs.get("alpha", 5e-3),
            eta=kwargs.get("eta", 0.5),
            top_k=kwargs.get("top_k", 256),
        )

    def train(self, pair_set: ContrastivePairSet) -> LayerActivations:
        """Compatibility interface: returns normalised mean-diff vectors."""
        result = self.train_dual(pair_set)
        primary_map: RawActivationMap = {}
        for layer in result.layer_order:
            vec = result.concept_directions[layer]
            vec = F.normalize(vec.float(), p=2, dim=-1)
            primary_map[layer] = vec
        dtype = self.kwargs.get("dtype", None)
        return LayerActivations(primary_map, dtype=dtype)

    def train_dual(self, pair_set: ContrastivePairSet) -> DualSteeringResult:
        """Full Dual Steering training: concept direction per layer."""
        buckets = self._collect_from_set(pair_set)
        if not buckets:
            raise InsufficientDataError(reason="No valid activation pairs found")

        concept_dirs: Dict[LayerName, torch.Tensor] = {}
        layer_order: List[LayerName] = []
        layer_meta: Dict[str, Any] = {}

        for layer_name in sorted(buckets.keys()):
            pos_list, neg_list = buckets[layer_name]
            if not pos_list or not neg_list:
                continue
            pos = torch.stack(
                [t.detach().float().reshape(-1) for t in pos_list], dim=0
            )
            neg = torch.stack(
                [t.detach().float().reshape(-1) for t in neg_list], dim=0
            )

            direction = pos.mean(dim=0) - neg.mean(dim=0)
            concept_dirs[layer_name] = direction.detach()
            layer_order.append(layer_name)

            layer_meta[str(layer_name)] = {
                "n_pos": pos.shape[0],
                "n_neg": neg.shape[0],
                "direction_norm": direction.norm().item(),
            }

        if not layer_order:
            raise InsufficientDataError(
                reason="No layers produced valid concept directions"
            )

        return DualSteeringResult(
            concept_directions=concept_dirs,
            config=self.config,
            layer_order=layer_order,
            metadata={"per_layer": layer_meta},
        )

    def _collect_from_set(
        self, pair_set: ContrastivePairSet
    ) -> Dict[LayerName, Tuple[List[torch.Tensor], List[torch.Tensor]]]:
        """Build {layer_name: ([pos tensors], [neg tensors])} from pairs."""
        buckets: Dict[
            LayerName, Tuple[List[torch.Tensor], List[torch.Tensor]]
        ] = defaultdict(lambda: ([], []))
        for pair in pair_set.pairs:
            pos_la = getattr(pair.positive_response, "layers_activations", None)
            neg_la = getattr(pair.negative_response, "layers_activations", None)
            if pos_la is None or neg_la is None:
                continue
            layer_names = set(pos_la.to_dict().keys()) | set(
                neg_la.to_dict().keys()
            )
            for layer in layer_names:
                p = pos_la.to_dict().get(layer) if pos_la else None
                n = neg_la.to_dict().get(layer) if neg_la else None
                if isinstance(p, torch.Tensor) and isinstance(n, torch.Tensor):
                    buckets[layer][0].append(p)
                    buckets[layer][1].append(n)
        return buckets
