"""
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.

v2 enhancements:
- Momentum-accelerated Newton (heavy-ball)
- Annealing Tikhonov regularisation
- Armijo backtracking via entropy
- Adaptive top-K via probability mass threshold
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def _batch_entropy(hidden: torch.Tensor, lm_head_weight: torch.Tensor) -> torch.Tensor:
    """Compute entropy of softmax distribution for each row in hidden.

    Args:
        hidden: [B, D] hidden states.
        lm_head_weight: [V, D] unembedding matrix.

    Returns:
        [B] entropy values.
    """
    logits = hidden @ lm_head_weight.T
    log_probs = F.log_softmax(logits, dim=-1)
    probs = log_probs.exp()
    return -(probs * log_probs).sum(dim=-1)


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

    1. Compute logits and softmax probabilities.
    2. Select top-K tokens by probability (adaptive if prob_mass_threshold > 0).
    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 (used when prob_mass_threshold <= 0).
        alpha: Tikhonov regularisation.
        prob_mass_threshold: if > 0, select K tokens to cover this fraction
            of probability mass instead of using fixed top_k.

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

    if prob_mass_threshold > 0:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cumsum = torch.cumsum(sorted_probs, dim=0)
        k = int((cumsum < prob_mass_threshold).sum().item()) + 1
        k = max(k, 1)
        k = min(k, probs.shape[-1])
        topk_probs = sorted_probs[:k]
        topk_idx = sorted_idx[:k]
    else:
        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,
    momentum: float = 0.0,
    alpha_decay: float = 1.0,
    armijo: bool = False,
    prob_mass_threshold: float = 0.0,
) -> 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.

    v2 enhancements:
    - momentum: heavy-ball momentum coefficient [0, 0.9]. 0 = disabled.
    - alpha_decay: multiply alpha by this each step [0.3, 1.0]. 1.0 = constant.
    - armijo: if True, halve eta when entropy increases (backtracking).
    - prob_mass_threshold: adaptive top-K via cumulative probability mass.

    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.
        momentum: heavy-ball momentum coefficient (0 disables).
        alpha_decay: per-step decay factor for alpha (1.0 = no decay).
        armijo: enable Armijo backtracking via entropy.
        prob_mass_threshold: adaptive K via probability mass (0 = fixed top_k).

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

    lam = hidden.clone()
    v_prev = torch.zeros_like(lam)
    alpha_t = alpha

    for step 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_t,
                prob_mass_threshold=prob_mass_threshold,
            )
            batch_dirs.append(v)
        v_batch = torch.stack(batch_dirs, dim=0)

        if momentum > 0:
            v_batch = v_batch + momentum * v_prev

        if armijo:
            entropy_before = _batch_entropy(lam, lm_head_weight)
            lam_candidate = lam + eta * v_batch
            entropy_after = _batch_entropy(lam_candidate, lm_head_weight)
            overshoot = entropy_after > entropy_before
            if overshoot.any():
                lam_half = lam + (eta / 2) * v_batch
                lam = torch.where(
                    overshoot.unsqueeze(-1), lam_half, lam_candidate
                )
            else:
                lam = lam_candidate
        else:
            lam = lam + eta * v_batch

        v_prev = v_batch
        alpha_t = alpha_t * alpha_decay

    return lam.squeeze(0) if squeeze else lam
