"""Independent float64 pure-Python reference for the GS1 TF-IDF baseline.

b_i = sum_{j!=i} w_ij r_j / sum_{j!=i} w_ij, w_ij = exp(cos(x_i, x_j)/T), word TF-IDF
(trajectory.ai 'You Only Rollout Once', Appendix B). TF-IDF: lowercase tokens r"(?u)\\b\\w\\w+\\b",
raw counts, idf = ln((1+n)/(1+df)) + 1, L2-normalized.
"""
import math
import re
from collections import Counter

TOKEN = re.compile(r"(?u)\b\w\w+\b")


def tokens(text):
    return TOKEN.findall(text.lower())


def fit_idf(corpus):
    df = Counter()
    for doc in corpus:
        df.update(set(tokens(doc)))
    n = len(corpus)
    return {t: math.log((1 + n) / (1 + d)) + 1.0 for t, d in df.items()}


def vector(text, idf):
    v = {t: c * idf[t] for t, c in Counter(tokens(text)).items() if t in idf}
    norm = math.sqrt(math.fsum(x * x for x in v.values()))
    return {t: x / norm for t, x in v.items()} if norm else v


def cosine(a, b):
    if len(a) > len(b):
        a, b = b, a
    return math.fsum(x * b[t] for t, x in a.items() if t in b)


def gs1_advantages(rewards, prompts, corpus, tau=0.1):
    idf = fit_idf(corpus)
    vecs = [vector(p, idf) for p in prompts]
    out = []
    for i in range(len(prompts)):
        sims = [(j, cosine(vecs[i], vecs[j])) for j in range(len(prompts)) if j != i]
        m = max(s for _, s in sims)
        w = [(j, math.exp((s - m) / tau)) for j, s in sims]
        z = math.fsum(x for _, x in w)
        b = math.fsum(x * rewards[j] for j, x in w) / z
        out.append(float(rewards[i]) - b)
    return out