"""
Newton-step steering along m-geodesics in the softmax information geometry.

Implements Algorithm 1 from Park et al. 2025: "The Information Geometry
of Softmax: Probing and Steering". Each step solves a regularised
linear system using the softmax-weighted covariance of top-K unembedding
vectors, then updates the hidden state along the resulting direction.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def topk_softmax_covariance(
    hidden: torch.Tensor,
    lm_head_weight: torch.Tensor,
    concept_dir: torch.Tensor,
    top_k: int = 256,
    alpha: float = 5e-3,
) -> torch.Tensor:
    """
    Compute one Newton direction from the softmax Fisher geometry.

    1. Compute logits and softmax probabilities.
    2. Select top-K tokens by probability.
    3. Gather their unembedding rows W_k.
    4. Compute p-weighted covariance of W_k.
    5. Solve (Cov + alpha * I) v = concept_dir.
    6. Return normalised v.

    Args:
        hidden: [D] current hidden state.
        lm_head_weight: [V, D] unembedding matrix.
        concept_dir: [D] concept direction (mean_pos - mean_neg).
        top_k: number of tokens to retain.
        alpha: Tikhonov regularisation.

    Returns:
        v: [D] normalised Newton direction.
    """
    logits = hidden @ lm_head_weight.T
    probs = F.softmax(logits, dim=-1)

    k = min(top_k, probs.shape[-1])
    topk_probs, topk_idx = torch.topk(probs, k)

    w_k = lm_head_weight[topk_idx]

    p = topk_probs.unsqueeze(-1)
    mu = (p * w_k).sum(dim=0)
    diff = w_k - mu.unsqueeze(0)
    cov = (p * diff).T @ diff

    d = hidden.shape[-1]
    cov_reg = cov + alpha * torch.eye(d, device=cov.device, dtype=cov.dtype)

    v = torch.linalg.solve(cov_reg, concept_dir)
    v = F.normalize(v, p=2, dim=-1)
    return v


def dual_newton_step(
    hidden: torch.Tensor,
    lm_head_weight: torch.Tensor,
    concept_dir: torch.Tensor,
    num_steps: int = 3,
    alpha: float = 5e-3,
    eta: float = 0.5,
    top_k: int = 256,
) -> torch.Tensor:
    """
    Iterative Newton steering along the m-geodesic.

    Runs ``num_steps`` iterations of the Newton update, each time
    recomputing the softmax covariance at the current hidden state.

    Args:
        hidden: [D] or [B, D] hidden state.
        lm_head_weight: [V, D] unembedding weight matrix.
        concept_dir: [D] concept direction.
        num_steps: number of Newton iterations.
        alpha: Tikhonov regularisation for covariance.
        eta: step-size multiplier per iteration.
        top_k: tokens retained for covariance.

    Returns:
        Updated hidden state, same shape as input.
    """
    squeeze = hidden.dim() == 1
    if squeeze:
        hidden = hidden.unsqueeze(0)

    lam = hidden.clone()
    for _ in range(num_steps):
        batch_dirs = []
        for i in range(lam.shape[0]):
            v = topk_softmax_covariance(
                lam[i], lm_head_weight, concept_dir,
                top_k=top_k, alpha=alpha,
            )
            batch_dirs.append(v)
        v_batch = torch.stack(batch_dirs, dim=0)
        lam = lam + eta * v_batch

    return lam.squeeze(0) if squeeze else lam
