"""NanoVDR / ColNanoVDR interactive retrieval demo.

Both students are text-only query encoders that run on CPU against a page index
their vision-language teacher built offline; no vision model runs at query time.
  - NanoVDR (single-vector): 69M DistilBERT student of Qwen3-VL-Embedding-2B,
    one vector scored by cosine similarity against the teacher's page vectors.
  - ColNanoVDR (multi-vector): 149M Ettin student of ColQwen3.5-4.5B, weighted
    token vectors scored by MaxSim against the teacher's page tokens.
Corpus: ViDoRe v3 Computer Science (1,360 pages).
"""

import json
import os
import time

import gradio as gr
import numpy as np
import torch
from sentence_transformers import MultiVectorEncoder, SentenceTransformer

DATA = "data"
torch.set_num_threads(max(1, os.cpu_count() or 1))

with open(os.path.join(DATA, "queries.json")) as f:
    EXAMPLES = json.load(f)
BY_TEXT = {q["text"]: q for q in EXAMPLES}

# ── ColNanoVDR: multi-vector index by ColQwen3.5-4.5B ─────────────────────────
mv_docs = torch.from_numpy(np.load(os.path.join(DATA, "colnanovdr", "doc_emb.npy")).astype(np.float32))
mv_len = torch.from_numpy(np.load(os.path.join(DATA, "colnanovdr", "doc_lengths.npy")))
mv_mask = torch.arange(mv_docs.shape[1])[None, :] < mv_len[:, None]
mv_model = MultiVectorEncoder("nanovdr/ColNanoVDR-Q-Ettin150M-ColQwen35-320-ML", trust_remote_code=True, device="cpu")

# ── NanoVDR: single-vector index by Qwen3-VL-Embedding-2B ─────────────────────
sv_docs = np.load(os.path.join(DATA, "corpus_embeddings.npy")).astype(np.float32)
sv_docs /= np.maximum(np.linalg.norm(sv_docs, axis=1, keepdims=True), 1e-12)
sv_model = SentenceTransformer("nanovdr/NanoVDR-Q-DistilBERT-Qwen3VL2B-2048-ML", device="cpu")

N_PAGES = sv_docs.shape[0]
NANOVDR, COLNANOVDR = "NanoVDR (single-vector, 69M)", "ColNanoVDR (multi-vector, 149M)"
ABOUT = {
    COLNANOVDR: (
        "**ColNanoVDR** ([model](https://huggingface.co/nanovdr/ColNanoVDR-Q-Ettin150M-ColQwen35-320-ML), "
        "[paper](https://arxiv.org/abs/2609.34899)): a 149M text-only encoder turns the query into one vector per token, "
        "each scaled by a learned weight, and **MaxSim** scores them against page tokens that "
        "[ColQwen3.5-4.5B](https://huggingface.co/athrael-soju/colqwen3.5-4.5B-v3) encoded offline."
    ),
    NANOVDR: (
        "**NanoVDR** ([model](https://huggingface.co/nanovdr/NanoVDR-Q-DistilBERT-Qwen3VL2B-2048-ML), "
        "[paper](https://arxiv.org/abs/2603.12824)): a 69M text-only encoder turns the query into a single vector, "
        "scored by **cosine similarity** against page vectors that "
        "[Qwen3-VL-Embedding-2B](https://huggingface.co/Qwen/Qwen3-VL-Embedding-2B) encoded offline."
    ),
}
TEACHER = {COLNANOVDR: ("colnanovdr", "ColQwen3.5-4.5B"), NANOVDR: ("nanovdr", "Qwen3-VL-Embedding-2B")}
for m in (mv_model, sv_model):
    m.encode_query(["warm up"]) if m is mv_model else m.encode(["warm up"])
print(f"Ready: {N_PAGES} pages; multi-vector index {tuple(mv_docs.shape)}, single-vector index {sv_docs.shape}")


def page(i):
    return os.path.join(DATA, "images", f"{i:04d}.jpg")


def token_weights(query, vecs):
    """Pair each non-special token with its learned weight (the norm of its vector)."""
    special = set(mv_model.tokenizer.all_special_ids)
    toks = [mv_model.tokenizer.decode([t]) for t in mv_model.tokenizer(query)["input_ids"] if t not in special]
    w = vecs.norm(dim=-1).tolist()
    if len(toks) != len(w):
        return [(query, None)]
    top = max(w) or 1.0
    return [(t, round(x / top, 2)) for t, x in zip(toks, w)]


def search(query, top_k, model_name):
    query, top_k = (query or "").strip(), int(top_k)
    if not query:
        return [], "Please enter a query.", [], [], ""
    t0 = time.perf_counter()
    if model_name == COLNANOVDR:
        q = mv_model.encode_query([query], convert_to_numpy=False)[0].float()
        t_enc = (time.perf_counter() - t0) * 1000
        t0 = time.perf_counter()
        with torch.no_grad():
            sim = torch.einsum("nld,kd->nlk", mv_docs, q).masked_fill(~mv_mask[:, :, None], float("-inf"))
            scores = sim.amax(1).sum(-1).numpy()
        weights, what = token_weights(query, q), f"{q.shape[0]} weighted token vectors · MaxSim"
    else:
        q = sv_model.encode([query], normalize_embeddings=True)[0]
        t_enc = (time.perf_counter() - t0) * 1000
        t0 = time.perf_counter()
        scores = sv_docs @ q
        weights, what = [], "1 vector · cosine similarity"
    ranked = np.argsort(-scores)[:top_k].tolist()
    t_score = (time.perf_counter() - t0) * 1000

    ex = BY_TEXT.get(query)
    rel = set(ex["relevant"]) if ex else set()
    cap = lambda r, i: f"#{r + 1}  page {i}" + ("  ✓ relevant" if i in rel else "")
    student = [(page(i), cap(r, i)) for r, i in enumerate(ranked)]
    info = f"Query encoded in **{t_enc:.0f} ms** on CPU ({what} over {N_PAGES:,} pages in {t_score:.0f} ms)"
    teacher, note = [], ""
    if ex:
        key, tname = TEACHER[model_name]
        tt = ex["teacher_top10"][key][:top_k]
        teacher = [(page(i), cap(r, i)) for r, i in enumerate(tt)]
        note = (f"Example query from ViDoRe v3: the {tname} teacher's own ranking is shown below for comparison "
                f"(precomputed). Overlap with the student's top-{top_k}: **{len(set(tt) & set(ranked))}/{top_k}**. "
                f"Pages marked ✓ are annotated as relevant.")
    return student, info, weights, teacher, note


def switch(model_name):
    return ABOUT[model_name], gr.update(visible=model_name == COLNANOVDR)


with gr.Blocks(title="NanoVDR Demo") as demo:
    gr.Markdown(
        f"""
        # NanoVDR: visual document retrieval with text-only query encoders

        Type a query. A small **text-only** student encodes it on **CPU** and is scored against a page index that its
        vision-language teacher built offline; the index is used unchanged and no vision model runs at query time.
        Choose the single-vector **NanoVDR** or the multi-vector **ColNanoVDR** student below.

        **Corpus**: {N_PAGES:,} pages from [ViDoRe v3 Computer Science](https://huggingface.co/datasets/vidore/vidore_v3_computer_science)
        (academic papers, slides, diagrams).
        """
    )
    model_in = gr.Radio([NANOVDR, COLNANOVDR], value=NANOVDR, label="Model")
    about = gr.Markdown(ABOUT[NANOVDR])
    with gr.Row():
        query_in = gr.Textbox(label="Query", lines=2, scale=4,
                              placeholder="e.g., How does the V-model integrate testing compared with Waterfall?")
        top_k = gr.Slider(1, 10, value=5, step=1, label="Top-K", scale=1)
    btn = gr.Button("Search", variant="primary")
    info = gr.Markdown()
    weights = gr.HighlightedText(label="Learned token weights (relative to the largest)", combine_adjacent=False,
                                 show_legend=False, visible=False)
    gallery = gr.Gallery(label="Student: retrieved pages", columns=5, height="auto", object_fit="contain")
    note = gr.Markdown()
    teacher_gallery = gr.Gallery(label="Teacher (example queries only, precomputed)", columns=5, height="auto",
                                 object_fit="contain")
    gr.Examples(examples=[[q["text"]] for q in EXAMPLES[:12]], inputs=[query_in], label="Example queries (ViDoRe v3)")

    outs = [gallery, info, weights, teacher_gallery, note]
    btn.click(search, [query_in, top_k, model_in], outs)
    query_in.submit(search, [query_in, top_k, model_in], outs)
    model_in.change(switch, [model_in], [about, weights]).then(search, [query_in, top_k, model_in], outs)

demo.launch(theme=gr.themes.Soft())