"""Diversity-weighted leave-one-out baseline for grouped policy samples."""
import numpy as np


def diversity_baseline(costs, embeddings, p=1.0, eps=1e-8):
    """Return baseline, minimization advantage, normalized weights, distances."""
    costs = np.asarray(costs, dtype=float)
    z = np.asarray(embeddings, dtype=float)
    if costs.ndim != 1 or z.ndim != 2 or len(costs) != len(z):
        raise ValueError("costs must be [B], embeddings must be [B,d]")
    if len(costs) < 2 or p < 0 or eps <= 0:
        raise ValueError("need B>=2, p>=0, eps>0")
    z = z - z.mean(axis=0, keepdims=True)
    d = ((z[:, None, :] - z[None, :, :]) ** 2).sum(axis=-1)
    w = (d + eps) ** p
    np.fill_diagonal(w, 0.0)
    denom = w.sum(axis=1)
    baseline = (w @ costs) / denom
    advantage = baseline - costs
    return baseline, advantage, w / denom[:, None], d


def uniform_loo(costs):
    costs = np.asarray(costs, dtype=float)
    return (costs.sum() - costs) / (len(costs) - 1)


def effective_count(normalized_weights):
    q = np.asarray(normalized_weights)
    return 1.0 / (q * q).sum(axis=-1)


def torch_diversity_loss(log_probs, costs, embeddings, p=1.0, eps=1e-8,
                         normalize_advantage=False):
    """SSPO loss; log_probs are [B] summed trajectory log probabilities.

    Baseline quantities are detached as required. Costs may be tensors or arrays;
    embeddings may retain gradients, but no gradient is allowed through them here.
    """
    import torch
    if log_probs.ndim != 1 or embeddings.ndim != 2 or costs.ndim != 1:
        raise ValueError("log_probs/costs must be [B], embeddings must be [B,d]")
    if log_probs.shape[0] != costs.shape[0] or costs.shape[0] != embeddings.shape[0]:
        raise ValueError("batch dimensions must agree")
    z = embeddings.detach() - embeddings.detach().mean(dim=0, keepdim=True)
    d = ((z[:, None, :] - z[None, :, :]) ** 2).sum(dim=-1)
    w = (d + eps).pow(p)
    eye = torch.eye(w.shape[0], dtype=torch.bool, device=w.device)
    w = w.masked_fill(eye, 0.0)
    denom = w.sum(dim=1)
    b = (w @ costs.detach()) / denom
    adv = b - costs.detach()
    if normalize_advantage:
        adv = adv / costs.detach().std().clamp_min(eps)
    return -(adv.detach() * log_probs).mean()