import json
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from transformers import CLIPTextModelWithProjection, CLIPTokenizer

EVALS = Path("/home/s56fasla/MOTIONX/evals")
CLIP_NAME = "openai/clip-vit-base-patch32"
ARMS = ["base", "baseline", "ski"]


def lcs(a, b):
    m = [[0] * (len(b) + 1) for _ in range(len(a) + 1)]
    for i in range(len(a)):
        for j in range(len(b)):
            m[i + 1][j + 1] = m[i][j] + 1 if a[i] == b[j] else max(m[i][j + 1], m[i + 1][j])
    return m[-1][-1]


def rouge_l(pred, ref):
    p, r = pred.lower().split(), ref.lower().split()
    if not p or not r:
        return 0.0
    l = lcs(p, r)
    prec, rec = l / len(p), l / len(r)
    return 0.0 if prec + rec == 0 else 2 * prec * rec / (prec + rec)


dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = CLIPTokenizer.from_pretrained(CLIP_NAME)
enc = CLIPTextModelWithProjection.from_pretrained(CLIP_NAME).to(dev).eval()


def embed(texts):
    out = []
    with torch.no_grad():
        for i in range(0, len(texts), 256):
            b = tok(texts[i:i + 256], padding=True, truncation=True,
                    max_length=77, return_tensors="pt").to(dev)
            out.append(F.normalize(enc(**b).text_embeds, dim=-1).cpu())
    return torch.cat(out).float()


results = {}
for arm in ARMS:
    f = EVALS / f"preds_{arm}.json"
    if not f.exists():
        print("missing", f, flush=True)
        continue
    preds = json.load(open(f))
    P = embed([p["pred"] for p in preds])
    R = embed([p["ref"] for p in preds])
    cos = float((P * R).sum(-1).mean())
    sim = P @ R.T
    tgt = torch.arange(len(preds))
    rank = (sim > sim[tgt, tgt].unsqueeze(1)).sum(1)
    rl = float(np.mean([rouge_l(p["pred"], p["ref"]) for p in preds]))
    results[arm] = {
        "n": len(preds),
        "clip_cos": round(cos, 4),
        "rouge_l": round(rl, 4),
        "retr_r1": round(float((rank < 1).float().mean()), 4),
        "retr_r5": round(float((rank < 5).float().mean()), 4),
        "medr": float(rank.float().median()) + 1,
        "mean_words": round(float(np.mean([len(p["pred"].split()) for p in preds])), 1),
        "chance_r1": round(1.0 / len(preds), 5),
    }
    print(arm, json.dumps(results[arm]), flush=True)

json.dump(results, open(EVALS / "caption_scores.json", "w"), indent=1)
print("saved", EVALS / "caption_scores.json", flush=True)