"""
Subspace-projected Newton steering in the softmax information geometry.

WICHER projects the Newton step into a learned low-rank SVD concept
subspace (k << D), achieving O(k^3) complexity instead of O(D^3).

Key differences from wicher/newton.py:
- Operates in k-dim subspace instead of full D-dim space
- Adaptive per-component regularization (alpha / sigma_i^2)
- EMA momentum instead of heavy-ball
- Orthogonal complement preserved by construction
"""

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 subspace_newton_direction(
    hidden: torch.Tensor,
    lm_head_weight: torch.Tensor,
    concept_dir: torch.Tensor,
    concept_basis: torch.Tensor,
    component_variances: torch.Tensor,
    top_k: int = 256,
    alpha: float = 5e-3,
    prob_mass_threshold: float = 0.0,
) -> torch.Tensor:
    """
    Compute one Newton direction projected through the concept subspace.

    1. Project concept_dir into k-dim subspace: w_proj = V_k @ w
    2. Compute logits/softmax, select top-K tokens
    3. Project top-K unembeddings into subspace: W_proj = W_topk @ V_k^T
    4. p-weighted covariance in subspace: Cov_k [k,k]
    5. Adaptive regularization: Cov_reg = Cov_k + diag(alpha / (sigma^2 + eps))
    6. Solve k*k system: v_k = Cov_reg^-1 @ w_proj
    7. Lift back: v = V_k^T @ v_k

    Args:
        hidden: [D] current hidden state.
        lm_head_weight: [V, D] unembedding matrix.
        concept_dir: [D] concept direction.
        concept_basis: [k, D] SVD basis vectors.
        component_variances: [k] singular values squared.
        top_k: tokens retained for covariance.
        alpha: base regularization.
        prob_mass_threshold: adaptive K via probability mass.

    Returns:
        v: [D] normalised Newton direction.
    """
    w_proj = concept_basis @ concept_dir

    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)
        n_tokens = int((cumsum < prob_mass_threshold).sum().item()) + 1
        n_tokens = max(n_tokens, 1)
        n_tokens = min(n_tokens, probs.shape[-1])
        topk_probs = sorted_probs[:n_tokens]
        topk_idx = sorted_idx[:n_tokens]
    else:
        n_tokens = min(top_k, probs.shape[-1])
        topk_probs, topk_idx = torch.topk(probs, n_tokens)

    w_topk = lm_head_weight[topk_idx]

    w_proj_k = w_topk @ concept_basis.T

    p = topk_probs.unsqueeze(-1)
    mu_k = (p * w_proj_k).sum(dim=0)
    diff_k = w_proj_k - mu_k.unsqueeze(0)
    cov_k = (p * diff_k).T @ diff_k

    eps = torch.median(component_variances).clamp(min=1e-8)
    reg_diag = alpha / (component_variances + eps)
    cov_reg = cov_k + torch.diag(reg_diag)

    v_k = torch.linalg.solve(cov_reg, w_proj)

    v = concept_basis.T @ v_k
    v = F.normalize(v, p=2, dim=-1)
    return v


def wicher_newton_step(
    hidden: torch.Tensor,
    lm_head_weight: torch.Tensor,
    concept_dir: torch.Tensor,
    concept_basis: torch.Tensor,
    component_variances: torch.Tensor,
    num_steps: int = 3,
    alpha: float = 5e-3,
    eta: float = 0.5,
    top_k: int = 256,
    beta: float = 0.0,
    alpha_decay: float = 1.0,
    armijo: bool = False,
    prob_mass_threshold: float = 0.0,
) -> torch.Tensor:
    """
    Iterative subspace-projected Newton steering.

    Args:
        hidden: [D] or [B, D] hidden state.
        lm_head_weight: [V, D] unembedding weight matrix.
        concept_dir: [D] concept direction.
        concept_basis: [k, D] SVD basis vectors.
        component_variances: [k] singular values squared.
        num_steps: Newton iterations.
        alpha: base Tikhonov regularisation.
        eta: step-size multiplier.
        top_k: tokens retained for covariance.
        beta: EMA momentum coefficient (0 = disabled).
        alpha_decay: per-step decay for alpha.
        armijo: enable Armijo backtracking via entropy.
        prob_mass_threshold: adaptive K via probability mass.

    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 = subspace_newton_direction(
                lam[i], lm_head_weight, concept_dir,
                concept_basis, component_variances,
                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 beta > 0:
            v_batch = beta * v_prev + (1.0 - beta) * v_batch

        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
