"""
Broyden quasi-Newton steering in the SVD concept subspace.

WICHER projects the Broyden step into a learned low-rank SVD concept
subspace (k << D), achieving O(k^2) complexity per step with NO
dependency on the lm_head weight matrix.

Key properties:
- Operates in k-dim subspace instead of full D-dim space
- Variance-weighted diagonal preconditioner from SVD training
- Sherman-Morrison rank-1 inverse Jacobian update per step
- EMA momentum via beta parameter
- No lm_head required — purely function-evaluation based
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def broyden_direction(
    z: torch.Tensor,
    w: torch.Tensor,
    H: torch.Tensor,
    alpha: float,
) -> torch.Tensor:
    """
    Compute one Broyden direction in the k-dim subspace.

    The residual is the gap between the current subspace projection
    and the desired shift along the concept direction w.

    Args:
        z: [k] current subspace coordinates.
        w: [k] concept direction projected into subspace.
        H: [k, k] current approximate inverse Jacobian.
        alpha: steering strength scalar.

    Returns:
        direction: [k] Broyden step direction.
    """
    residual = -alpha * w
    direction = H @ residual
    return direction


def _sherman_morrison_update(
    H: torch.Tensor,
    delta_z: torch.Tensor,
    delta_r: torch.Tensor,
) -> torch.Tensor:
    """
    Sherman-Morrison rank-1 update of approximate inverse Jacobian.

    H_new = H + ((delta_z - H @ delta_r) @ delta_z^T @ H)
                / (delta_z^T @ H @ delta_r)

    Args:
        H: [k, k] current inverse Jacobian approximation.
        delta_z: [k] change in position.
        delta_r: [k] change in residual.

    Returns:
        H_new: [k, k] updated inverse Jacobian.
    """
    H_dr = H @ delta_r
    denom = delta_z @ H_dr
    if denom.abs() < 1e-12:
        return H
    numerator = (delta_z - H_dr).unsqueeze(1) @ (delta_z @ H).unsqueeze(0)
    return H + numerator / denom


def wicher_broyden_step(
    hidden: 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,
    beta: float = 0.0,
    alpha_decay: float = 1.0,
) -> torch.Tensor:
    """
    Iterative subspace-projected Broyden steering.

    Args:
        hidden: [D] or [B, D] hidden state.
        concept_dir: [D] concept direction.
        concept_basis: [k, D] SVD basis vectors.
        component_variances: [k] singular values squared.
        num_steps: Broyden iterations.
        alpha: base steering strength.
        eta: step-size multiplier.
        beta: EMA momentum coefficient (0 = disabled).
        alpha_decay: per-step decay for alpha.

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

    B = hidden.shape[0]
    k = concept_basis.shape[0]
    device = hidden.device
    dtype = hidden.dtype

    w = concept_basis @ concept_dir

    eps = component_variances.sum().clamp(min=1e-8)
    H_diag = component_variances / eps
    H_init = torch.diag(H_diag)

    results = []
    for i in range(B):
        h = hidden[i].float()
        z = concept_basis @ h

        H = H_init.clone()
        alpha_t = alpha
        v_prev = torch.zeros(k, device=device, dtype=torch.float32)

        r = -alpha_t * w

        for _step in range(num_steps):
            direction = H @ r
            if beta > 0:
                direction = beta * v_prev + (1.0 - beta) * direction

            z_new = z + eta * direction
            r_new = -alpha_t * w

            delta_z = z_new - z
            delta_r = r_new - r
            if delta_r.norm() > 1e-12:
                H = _sherman_morrison_update(H, delta_z, delta_r)

            v_prev = direction
            z = z_new
            r = r_new
            alpha_t = alpha_t * alpha_decay

        h_new = h + concept_basis.T @ (z - concept_basis @ h)
        results.append(h_new.to(dtype))

    out = torch.stack(results, dim=0)
    return out.squeeze(0) if squeeze else out
