"""Usage + a tiny built-in sanity check for TrynMini v2.

    python example.py

Encodes a few sentences, prints pairwise cosine similarities, and runs a small
labeled STS check (Spearman via numpy). Needs only numpy + tokenizers (+ safetensors
if the model ships as model.safetensors).
"""
import numpy as np
from modeling_trynmini import TrynMiniV2

# (sentence_a, sentence_b, gold 0..5)
PAIRS = [
    ("A man is playing a guitar.", "A person plays the guitar.", 4.6),
    ("A man is playing a guitar.", "A woman is slicing an onion.", 0.4),
    ("The cat sits on the mat.", "A cat is sitting on a mat.", 4.8),
    ("Stocks fell sharply today.", "Markets dropped a lot today.", 4.2),
    ("Stocks fell sharply today.", "The weather is sunny and warm.", 0.2),
    ("The plane landed safely.", "The aircraft touched down.", 4.5),
    ("The plane landed safely.", "The chef cooked a meal.", 0.3),
]


def _rank(a):
    a = np.asarray(a, float)
    order = a.argsort(kind="mergesort")
    inv = np.empty(len(a), int); inv[order] = np.arange(len(a))
    s = a[order]; obs = np.r_[True, s[1:] != s[:-1]]
    dense = obs.cumsum()[inv]; cnt = np.r_[np.nonzero(obs)[0], len(a)]
    return 0.5 * (cnt[dense] + cnt[dense - 1] + 1)


def spearman(a, b):
    ra, rb = _rank(a), _rank(b)
    if ra.std() == 0 or rb.std() == 0:
        return 0.0
    return float(np.corrcoef(ra, rb)[0, 1])


def main():
    m = TrynMiniV2.from_pretrained(".")
    print("Encoding demo:")
    sents = ["a man plays guitar", "someone plays a guitar", "the stock market fell"]
    e = m.encode(sents, dim=256)
    print(f"  shape={e.shape}  (L2-normalized)")
    print(f"  sim(0,1 paraphrase) = {float(e[0] @ e[1]):.3f}")
    print(f"  sim(0,2 unrelated)  = {float(e[0] @ e[2]):.3f}")

    s1 = [p[0] for p in PAIRS]; s2 = [p[1] for p in PAIRS]; gold = [p[2] for p in PAIRS]
    for d in (64, 128, 256):
        c = (m.encode(s1, dim=d) * m.encode(s2, dim=d)).sum(1)
        print(f"  built-in STS Spearman @dim={d:<3} = {spearman(c, gold):.3f}")


if __name__ == "__main__":
    main()