"""GS1 / GRPO advantage computation (implementation under test)."""
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer


class TfidfKernel:
    def __init__(self, corpus, tau=0.1):
        self.vec = TfidfVectorizer(dtype=np.float64).fit(corpus)
        self.tau = tau

    def weights(self, prompts):
        X = self.vec.transform(prompts)
        S = (X @ X.T).toarray()
        logits = S / self.tau
        np.fill_diagonal(logits, -np.inf)
        logits -= logits.max(axis=1, keepdims=True)
        W = np.exp(logits)
        return W / W.sum(axis=1, keepdims=True)


def grpo_advantages(rewards_per_group):
    return [np.asarray(r, np.float64) - np.mean(np.asarray(r, np.float64)) for r in rewards_per_group]


def gs1_advantages(rewards, prompts, kernel):
    r = np.asarray(rewards, np.float64)
    W = kernel.weights(prompts)
    b = W @ r
    n = len(r)
    uniform = (r.sum() - r) / (n - 1)
    diag = {"n_eff": 1.0 / (W ** 2).sum(axis=1), "baseline": b, "uniform_loo": uniform}
    return r - b, diag