import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from scripts.omniret_embedding import OmniRetEmbedder

IDS = ("-1rZFviqTTQ_000003", "-25e5qcELvw_000011")
EXPECTED = torch.tensor(
    [
        [0.32468632, -0.03668483, 0.28203371, 0.04291208, 0.42771006, 0.06420071],
        [-0.03580582, 0.16362236, -0.03176199, 0.33771127, -0.00980220, 0.16253276],
        [0.41171762, 0.00378584, 0.22568333, 0.01310393, 0.43492085, 0.08068858],
    ]
)

queries = [
    {
        "instruction": "Retrieve the video that aligns with the audio.",
        "audio": ROOT / "examples" / "audios" / f"{IDS[0]}.wav",
    },
    {
        "instruction": "Retrieve the audio that matches the given image.",
        "image": ROOT / "examples" / "images" / f"{IDS[1]}.jpg",
    },
    {
        "instruction": "Retrieve the audio that matches the given video.",
        "video": ROOT / "examples" / "videos" / f"{IDS[0]}.mp4",
    },
]
documents = [
    {modality: ROOT / "examples" / f"{modality}s" / f"{media_id}.{extension}"}
    for modality, extension in (("audio", "wav"), ("image", "jpg"), ("video", "mp4"))
    for media_id in IDS
]

model = OmniRetEmbedder(ROOT, torch_dtype=torch.bfloat16)
embeddings = model.process(queries + documents)
scores = embeddings[: len(queries)] @ embeddings[len(queries) :].T
print(scores.tolist())
if not torch.allclose(scores.cpu(), EXPECTED, atol=1e-3, rtol=0):
    raise SystemExit("scores differ from checkpoint-2220 expectations")