import json
import sys
from pathlib import Path

import numpy as np
import open_clip
import torch

# --- Settings (must match build_index.py) ---
EMB_DIR = Path("embeddings")
MODEL_NAME = "ViT-B-32"
PRETRAINED = "laion2b_s34b_b79k"

# --- Load everything once, when the module is imported ---
model, _, _ = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED)
model.eval()
tokenizer = open_clip.get_tokenizer(MODEL_NAME)

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


def embed_text(query: str) -> np.ndarray:
    """Turn a text query into a normalized 512-number vector."""
    tokens = tokenizer([query])                       # batch of 1 query
    with torch.no_grad():
        vec = model.encode_text(tokens)               # (1, 512)
        vec = vec / vec.norm(dim=-1, keepdim=True)
    return vec[0].cpu().numpy()                       # (512,)


def search(query: str, k: int = 5) -> list[tuple[str, float]]:
    """Return the top-k (filename, score) pairs, best first."""
    q = embed_text(query)
    scores = embeddings @ q                           # (1000, 512) @ (512,) -> (1000,)
    top = np.argsort(scores)[::-1][:k]                # positions of the k highest scores
    return [(filenames[i], float(scores[i])) for i in top]


if __name__ == "__main__":
    query = " ".join(sys.argv[1:]) or "a dog running on the beach"
    for rank, (name, score) in enumerate(search(query, k=5), start=1):
        print(f"{rank}. {name}  score={score:.3f}")