"""Text -> image search with a trained Native-Bird x SigLIP alignment checkpoint.

python search.py --ckpt runs/native_bird/best.pt "Nkịta na-agba ọsọ n'ahịhịa" "A dog running on grass" -k 5
python search.py --ckpt runs/native_bird/best.pt --split test --json out.json "query 1" "query 2"

The gallery is every image of --split (default test, 809 images). Image features come from the same frozen
image encoder the checkpoint was trained with (cached in data/features/, computed on first use; needs the
Flickr8k images at data.image_dir). Works on CPU or GPU.
"""
import argparse, json, os, torch
from nbmm.common import (seed_all, pick_device, load_split, image_features, text_features, load_sentence_model,
                         text_arm, TrainableText, AlignModel, rel, _set)


class Searcher:
    def __init__(self, ckpt, split="test", device="auto", overrides=()):
        ck = torch.load(rel(ckpt), map_location="cpu", weights_only=False)
        self.cfg = cfg = ck["cfg"]
        for o in overrides:
            k, v = o.split("=", 1); _set(cfg, k, v)
        seed_all(cfg["seed"])
        self.dev = pick_device(device)
        self.df = load_split(cfg, split)
        self.ids = self.df.image_id.tolist()
        I = image_features(cfg, self.df, split, self.dev)
        self.model = AlignModel(ck["d_img"], ck["d_txt"], cfg).to(self.dev).eval()
        self.model.load_state_dict(ck["model"])
        self.unfreeze = cfg["text_encoder"]["unfreeze_last_n_layers"] > 0
        if self.unfreeze:
            self.txt = TrainableText(cfg, self.dev).to(self.dev).eval()
            self.txt.load_state_dict(ck["text_trainable"], strict=False)
        else:
            self.arm = text_arm(cfg)
            self.st = load_sentence_model(self.arm, self.dev)
        with torch.no_grad():
            self.zi = self.model.img(I.to(self.dev))

    @torch.no_grad()
    def embed_text(self, texts):
        if self.unfreeze:
            e = self.txt(texts)
        else:
            e = self.st.encode([self.arm.get("prompt", "") + t for t in texts], convert_to_tensor=True,
                               normalize_embeddings=True).float().to(self.dev)
        return self.model.txt(e)

    @torch.no_grad()
    def search(self, texts, k=5):
        sims = self.embed_text(texts) @ self.zi.t()
        out = []
        for t, row in zip(texts, sims):
            order = row.argsort(descending=True)
            top = order[:k].tolist()
            out.append({"query": t, "top": [{"image_id": self.ids[j], "score": round(row[j].item(), 4)} for j in top],
                        "_order": order})
        return out

    def rank_of(self, res, image_id):
        j = self.ids.index(image_id)
        return int((res["_order"] == j).nonzero()[0].item()) + 1  # 1-based


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--split", default="test")
    ap.add_argument("-k", type=int, default=5)
    ap.add_argument("--json", default=None)
    ap.add_argument("queries", nargs="+")
    a = ap.parse_args()
    s = Searcher(a.ckpt, a.split)
    res = s.search(a.queries, a.k)
    for r in res:
        r.pop("_order")
        print(r["query"])
        for i, h in enumerate(r["top"], 1):
            print(f"  {i}. {h['image_id']}  {h['score']}")
    if a.json:
        json.dump(res, open(a.json, "w"), indent=2, ensure_ascii=False)


if __name__ == "__main__":
    main()