"""Evaluate embedding models on held-out Kazakh retrieval + cross-lingual matching.

  CUDA_VISIBLE_DEVICES=4 venv-emb/bin/python scripts/eval_embed.py

Models: v2 (checkpoints/embed-v2/final), base (BAAI/bge-m3), v1 (Darmm/darmm-embedding-multilingual).
Writes results/embed_v2.json.
"""
import json
from pathlib import Path

import torch
from datasets import load_from_disk
from sentence_transformers import SentenceTransformer

MODELS = {
    "v2": "checkpoints/embed-v2/final",
    "base-bge-m3": "BAAI/bge-m3",
    "v1-darmm": "Darmm/darmm-embedding-multilingual",
}


def encode(model, texts):
    return model.encode(texts, batch_size=128, normalize_embeddings=True, convert_to_tensor=True, show_progress_bar=False)


def retrieval_metrics(model, queries, docs):
    q = encode(model, queries)
    d = encode(model, docs)
    r1 = r5 = r10 = mrr = 0.0
    for i in range(0, len(queries), 512):
        sims = q[i : i + 512] @ d.T
        top = sims.topk(10, dim=1).indices.cpu()
        for row, gold in enumerate(range(i, min(i + 512, len(queries)))):
            hits = (top[row] == gold).nonzero()
            if len(hits):
                rank = int(hits[0][0]) + 1
                mrr += 1.0 / rank
                r10 += 1
                r5 += rank <= 5
                r1 += rank == 1
    n = len(queries)
    return {"recall@1": round(r1 / n, 4), "recall@5": round(r5 / n, 4), "recall@10": round(r10 / n, 4), "mrr@10": round(mrr / n, 4), "n": n}


def xl_accuracy(model, kk, other):
    a = encode(model, kk)
    b = encode(model, other)
    pred = (a @ b.T).argmax(dim=1).cpu()
    acc = (pred == torch.arange(len(kk))).float().mean().item()
    return {"acc@1": round(acc, 4), "n": len(kk)}


def main():
    ds = load_from_disk("data/kk_pairs")
    ret = ds["eval_retrieval"]
    queries, docs = list(ret["query"]), list(ret["doc"])
    kkru = ds["eval_xl_kkru"]
    kken = ds["eval_xl_kken"]

    results = {}
    for name, path in MODELS.items():
        model = SentenceTransformer(path, model_kwargs={"torch_dtype": torch.float16} if torch.cuda.is_available() else {})
        model.max_seq_length = 192
        results[name] = {
            "kk_retrieval": retrieval_metrics(model, queries, docs),
            "xl_kk_ru": xl_accuracy(model, list(kkru["kk"]), list(kkru["other"])),
            "xl_kk_en": xl_accuracy(model, list(kken["kk"]), list(kken["other"])),
        }
        print(name, json.dumps(results[name]), flush=True)
        del model
        torch.cuda.empty_cache()

    Path("results").mkdir(exist_ok=True)
    Path("results/embed_v2.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()