"""Original Wanda/SparseGPT importance adapted to SwiGLU mid-neuron groups.

The algorithms remain unstructured at the score level, but the selected
coordinates are removed as complete SwiGLU neurons so checkpoint shapes shrink.
"""

from __future__ import annotations

import torch


def wanda_mid_scores(energy: torch.Tensor, down_weight: torch.Tensor) -> torch.Tensor:
    """Return classic Wanda scores for V/A/T and each SwiGLU mid coordinate."""
    if energy.ndim != 2 or energy.shape[0] != 3:
        raise ValueError(f"energy must be [3, d_int], got {tuple(energy.shape)}")
    if down_weight.ndim != 2 or down_weight.shape[1] != energy.shape[1]:
        raise ValueError(
            f"down_weight must be [hidden, {energy.shape[1]}], got {tuple(down_weight.shape)}"
        )
    col_norm = down_weight.detach().float().norm(dim=0)
    return col_norm.unsqueeze(0) * energy.detach().float().clamp_min(0).sqrt()


def _damped_hinv(gram: torch.Tensor, percdamp: float) -> torch.Tensor:
    """Upper-triangular Cholesky factor of damped H^{-1} (miniSeedance SparseGPT)."""
    if gram.ndim != 2 or gram.shape[0] != gram.shape[1]:
        raise ValueError(f"gram must be square, got {tuple(gram.shape)}")
    h = gram.detach().float().clone()
    h = (h + h.T) * 0.5
    diag = torch.diag(h)
    dead = diag == 0
    if bool(dead.any()):
        ids = torch.arange(h.shape[0], device=h.device)
        h[ids[dead], ids[dead]] = 1.0
        diag = torch.diag(h)
    damping = float(percdamp) * diag.clamp_min(1e-8).mean()
    ids = torch.arange(h.shape[0], device=h.device)
    inv = None
    for _ in range(8):
        trial = h.clone()
        trial[ids, ids] += damping
        chol, info = torch.linalg.cholesky_ex(trial)
        if int(info.max()) == 0:
            inv = torch.cholesky_inverse(chol)
            break
        damping *= 10.0
    if inv is None:
        raise torch.linalg.LinAlgError("unable to stabilize Gram matrix for SparseGPT")
    return torch.linalg.cholesky(inv, upper=True)


def sparsegpt_mid_scores(
    gram: torch.Tensor,
    energy: torch.Tensor,
    down_weight: torch.Tensor,
    percdamp: float = 0.01,
) -> torch.Tensor:
    """Return OBS/SparseGPT salience for grouped SwiGLU mid coordinates."""
    if energy.ndim != 2 or energy.shape[0] != 3:
        raise ValueError(f"energy must be [3, d_int], got {tuple(energy.shape)}")
    if down_weight.ndim != 2 or down_weight.shape[1] != energy.shape[1]:
        raise ValueError(
            f"down_weight must be [hidden, {energy.shape[1]}], got {tuple(down_weight.shape)}"
        )
    hinv = _damped_hinv(gram, percdamp)
    d = torch.diag(hinv).clamp_min(1e-8)
    base = down_weight.detach().float().pow(2).sum(dim=0) / d.pow(2)
    return energy.detach().float().clamp_min(0) * base.unsqueeze(0)


def sparsegpt_compensate_down(
    down_weight: torch.Tensor,
    gram: torch.Tensor,
    keep_idx: torch.Tensor,
    percdamp: float = 0.01,
    eps: float = 1e-8,
) -> torch.Tensor:
    """OBS-style compensation on down_proj before physical slice (miniSeedance SparseGPT)."""
    keep_idx = keep_idx.long()
    w = down_weight.detach().clone().float()
    d_int = w.shape[1]
    keep_set = set(keep_idx.tolist())
    drop = [i for i in range(d_int) if i not in keep_set]
    if not drop:
        return w.to(down_weight.dtype)
    work = gram.device
    if work.type == "cpu" and torch.cuda.is_available():
        work = torch.device("cuda")
        w = w.to(work)
        gram = gram.to(work)
    hinv = _damped_hinv(gram, percdamp)
    d_diag = torch.diag(hinv).clamp_min(eps)
    importance = w.pow(2).sum(dim=0) / d_diag.pow(2)
    drop_t = torch.tensor(drop, device=w.device, dtype=torch.long)
    drop_sorted = drop_t[torch.argsort(importance[drop_t])]
    for i in drop_sorted.tolist():
        col = w[:, i]
        d = hinv[i, i]
        err = col / d
        w[:, i:] -= err.unsqueeze(1).matmul(hinv[i, i:].unsqueeze(0))
        w[:, i] = 0
    return w.to(device=down_weight.device, dtype=down_weight.dtype)