import json
import torch
import numpy as np

from pathlib import Path
from tqdm import tqdm
from search import model, tokenizer

# 1. Load filenames.json and image_embeddings.npy
# build a lookup: filename → row number (0..999)

EMB_DIR = Path("embeddings")

image_embeddings = np.load(EMB_DIR / "image_embeddings.npy")              # (1000, 512)
filenames = json.loads((EMB_DIR / "filenames.json").read_text())         # 1000 names

row_of = {name: i for i, name in enumerate(filenames)}

# 2. Read Flickr8k.token.txt
#  for each line: split on "\t", strip the "#n" from the filename
#  keep it only if the filename is in the lookup
#  → captions:   list of 5,000 strings
#  → true_rows:  list of 5,000 row numbers (which image each caption belongs to)

CAPTIONS_FILE = Path("data/Flickr8k.token.txt")

captions = []     # the caption texts
true_rows = []    # for each caption, the row of the image it describes

for line in CAPTIONS_FILE.read_text().splitlines():
    if not line.strip():
        continue
    image_ref, caption = line.split("\t", 1)      # "abc.jpg#0", "A dog runs ..."
    filename = image_ref.split("#")[0]            # "abc.jpg"
    if filename in row_of:                        # only the 1,000 indexed images
        captions.append(caption.strip())
        true_rows.append(row_of[filename])

true_rows = np.array(true_rows)

# 3. Embed the 5,000 captions with CLIP's text encoder
#  in batches (e.g. 256), normalize, like build_index.py did for images
#  → text_embeddings: (5000, 512)

BATCH_SIZE = 256

text_vectors = []
for start in tqdm(range(0, len(captions), BATCH_SIZE)):
    batch = captions[start:start + BATCH_SIZE]
    tokens = tokenizer(batch)
    with torch.no_grad():
        vectors = model.encode_text(tokens)
        vectors = vectors / vectors.norm(dim=-1, keepdim=True)
    text_vectors.append(vectors.cpu().numpy())

text_embeddings = np.concatenate(text_vectors)

# 4. Similarity matrix
#  sims = text_embeddings @ image_embeddings.T      → (5000, 1000)

sims = text_embeddings @ image_embeddings.T          # (5000, 1000)

# 5. Rank of the true image for each caption
#  true_scores = sims[i, true_rows[i]] for each i
#  rank = how many images scored HIGHER than the true one
#        (0 = it came first, 4 = it came fifth, …)

true_scores = sims[np.arange(len(captions)), true_rows]          # (5000,)
ranks = (sims > true_scores[:, None]).sum(axis=1)                # (5000,)  0 = first place

# 6. Recall@K = fraction of captions with rank < K, for K = 1, 5, 10
#  print a small table

print()
print("Text → image retrieval on Flickr8k test (1,000 images, 5,000 captions)")
print("Model: CLIP ViT-B/32 (laion2b_s34b_b79k), zero-shot")
print("-" * 40)
for k in (1, 5, 10):
    recall = (ranks < k).mean()
    print(f"Recall@{k:<2}  {recall:6.1%}")
print("-" * 40)
print(f"Median rank: {int(np.median(ranks)) + 1}")