"""Contribution-space redundancy and diversity-aware keep selection.

Negative TENP scores are never diversity-discounted: a neuron that points
against the expert output stays a prune candidate even if it looks unique.
"""

from __future__ import annotations

import torch

from .surgery import round_to_multiple


def n_keep_width(d_int: int, p_mid: float, multiple: int = 4) -> int:
    n_keep = max(multiple, round_to_multiple(int((1.0 - p_mid) * d_int), multiple))
    return min(n_keep, d_int)


def gram_cosine(gram: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """Column cosine from a Gram matrix G ≈ AᵀA, shape [d, d]."""
    if gram.ndim != 2 or gram.shape[0] != gram.shape[1]:
        raise ValueError(f"gram must be square, got {tuple(gram.shape)}")
    g = gram.detach().float()
    g = (g + g.T) * 0.5
    nrm = torch.diag(g).clamp_min(0).sqrt().clamp_min(eps)
    return g / (nrm.unsqueeze(0) * nrm.unsqueeze(1))


def weight_cosine(down_weight: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """Column cosine of W_down, shape [hidden, d_int] -> [d_int, d_int]."""
    w = down_weight.detach().float()
    nrm = w.norm(dim=0).clamp_min(eps)
    hat = w / nrm.unsqueeze(0)
    return hat.T @ hat


def contribution_redundancy(
    gram: torch.Tensor | None,
    down_weight: torch.Tensor,
    eps: float = 1e-8,
) -> torch.Tensor:
    """R⁺ = max(0, cos(a) cos(W)). Missing Gram falls back to weight cosine only."""
    r_w = weight_cosine(down_weight, eps=eps)
    if gram is None:
        r = r_w
    else:
        r = gram_cosine(gram, eps=eps) * r_w
    r = torch.maximum(r, torch.zeros_like(r))
    d = int(r.shape[0])
    eye = torch.eye(d, device=r.device, dtype=r.dtype)
    return r * (1.0 - eye) + eye


def greedy_diversity_keep(
    score: torch.Tensor,
    r_plus: torch.Tensor,
    n_keep: int,
) -> torch.Tensor:
    """Greedy keep: U_i = S_i (1 - r_i^max) for S_i > 0, else U_i = S_i."""
    s = score.detach().float().reshape(-1)
    d = int(s.numel())
    n_keep = min(max(int(n_keep), 1), d)
    if r_plus.shape != (d, d):
        raise ValueError(f"r_plus must be [{d}, {d}], got {tuple(r_plus.shape)}")
    r = r_plus.detach().float()
    remaining = torch.ones(d, dtype=torch.bool, device=s.device)
    r_max = torch.zeros(d, device=s.device, dtype=s.dtype)
    chosen: list[int] = []
    neg_inf = torch.tensor(float("-inf"), device=s.device, dtype=s.dtype)
    for _ in range(n_keep):
        utility = s.clone()
        pos = s > 0
        utility = torch.where(pos, s * (1.0 - r_max), s)
        utility = torch.where(remaining, utility, neg_inf)
        j = int(torch.argmax(utility).item())
        chosen.append(j)
        remaining[j] = False
        r_max = torch.maximum(r_max, r[j])
    idx = torch.tensor(chosen, device=s.device, dtype=torch.long)
    idx, _ = torch.sort(idx)
    return idx