"""V-SPLADE Quality — Visual Document Retrieval Demo.

Upload a document page image and enter text queries to see:
  1. The top activated vocabulary tokens (sparse representation)
  2. Similarity scores between each query and the document image
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST be first (before torch)
import torch
import gradio as gr
from PIL import Image
from transformers import AutoProcessor

# Make the v-splade train/models package importable
import sys
from pathlib import Path

_TRAIN_DIR = Path(__file__).resolve().parent / "train"
if str(_TRAIN_DIR) not in sys.path:
    sys.path.insert(0, str(_TRAIN_DIR))

from models import build_model

MODEL_ID = "naver/v-splade-quality"


# ── Module-scope model load (eager, as required by ZeroGPU) ──────────────────
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = build_model(MODEL_ID, mode="inference_only", dtype=torch.bfloat16).to("cuda")
tokenizer = processor.tokenizer


@spaces.GPU(duration=60)
def retrieve(image: Image.Image, queries: str, topk: int = 15) -> tuple:
    """Encode a document image and score it against text queries.

    Args:
        image: A document page image (PNG/JPEG).
        queries: Newline-separated text queries to score against the image.
        topk: Number of top activated vocabulary tokens to display.

    Returns:
        A tuple of (top tokens HTML, query scores HTML, sparse stats text).
    """
    if image is None:
        return "<p style='color:red'>Please upload a document image.</p>", "", ""
    if not queries.strip():
        return "", "<p style='color:red'>Please enter at least one query.</p>", ""

    # Build lookup table on GPU (lazy — first call builds it).
    model.query_encoder.build_lookup_table()

    # ── Encode image → sparse embedding ──────────────────────────────────
    chat = [{"role": "user",
             "content": [{"type": "image"}, {"type": "text", "text": ""}]}]
    prompt = processor.apply_chat_template(chat, add_generation_prompt=True)
    inputs = processor(images=[image.convert("RGB")], text=prompt,
                       return_tensors="pt")
    inputs = {k: v.to("cuda") if torch.is_tensor(v) else v
              for k, v in inputs.items()}
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
    doc_vec = model.encode_passage(**inputs)[0].cpu()

    nnz = int((doc_vec > 0).sum())
    max_val = float(doc_vec.max())

    # ── Top-k activated vocabulary tokens ────────────────────────────────
    top_w, top_ids = torch.topk(doc_vec.float(), k=min(topk, doc_vec.shape[0]))
    tokens_html = "<table style='width:100%; border-collapse:collapse;'>"
    tokens_html += "<tr style='background:#f0f0f0;'><th style='padding:6px; text-align:left;'>Rank</th><th style='padding:6px; text-align:left;'>Token</th><th style='padding:6px; text-align:right;'>Weight</th></tr>"
    for rank, (idx, w) in enumerate(zip(top_ids, top_w), 1):
        tok_str = tokenizer.decode([int(idx)]).strip() or f"<id:{int(idx)}>"
        bar_width = max(2, int(float(w) / max_val * 200))
        tokens_html += (
            f"<tr><td style='padding:4px;'>{rank}</td>"
            f"<td style='padding:4px; font-family:monospace;'>{tok_str}</td>"
            f"<td style='padding:4px; text-align:right;'>"
            f"<div style='display:flex; align-items:center; gap:8px;'>"
            f"<div style='background:linear-gradient(90deg, #ff6b9d, #8b5cf6); width:{bar_width}px; height:16px; border-radius:3px;'></div>"
            f"<span>{float(w):.4f}</span></div></td></tr>"
        )
    tokens_html += "</table>"

    # ── Query similarity scores ──────────────────────────────────────────
    query_list = [q.strip() for q in queries.strip().split("\n") if q.strip()]
    scores_html = "<table style='width:100%; border-collapse:collapse;'>"
    scores_html += "<tr style='background:#f0f0f0;'><th style='padding:6px; text-align:left;'>Query</th><th style='padding:6px; text-align:right;'>Score</th><th style='padding:6px;'>Top matching tokens</th></tr>"

    max_score = 0.0
    results = []
    for q in query_list:
        tok = tokenizer(q, return_tensors="pt", add_special_tokens=False)
        q_vec = model.encode_query(
            tok["input_ids"].to("cuda"),
            tok["attention_mask"].to("cuda"),
        )[0].cpu()
        score = float((q_vec.float() * doc_vec.float()).sum())
        max_score = max(max_score, abs(score))
        # Top contributing tokens
        contrib = (q_vec.float() * doc_vec.float()).cpu()
        top_cw, top_cids = torch.topk(contrib, k=5)
        contribs = ", ".join(
            f"{tokenizer.decode([int(i)]).strip()}({float(w):.3f})"
            for i, w in zip(top_cids, top_cw) if float(w) > 0
        )
        results.append((q, score, contribs))

    for q, score, contribs in results:
        bar_width = max(2, int(abs(score) / max(max_score, 1e-6) * 200))
        color = "#22c55e" if score > 0 else "#ef4444"
        scores_html += (
            f"<tr><td style='padding:4px;'>{q}</td>"
            f"<td style='padding:4px; text-align:right;'><b>{score:.4f}</b></td>"
            f"<td style='padding:4px; font-size:0.9em; font-family:monospace;'>{contribs or '—'}</td></tr>"
        )
    scores_html += "</table>"

    stats = (
        f"Sparse vector: vocab_size={doc_vec.shape[0]}, nnz={nnz} "
        f"({nnz/doc_vec.shape[0]*100:.1f}% active), max={max_val:.4f}"
    )

    return tokens_html, scores_html, stats


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# 🔍 V-SPLADE Quality — Visual Document Retrieval\n"
        "Upload a document page image and enter text queries to see the "
        "sparse lexical representation and similarity scores in real-time. "
        "[Model card](https://huggingface.co/naver/v-splade-quality) · "
        "[Paper](https://arxiv.org/abs/2605.30917) · "
        "[Code](https://github.com/naver/v-splade)"
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Document page", type="pil", height=400
                )
                queries_input = gr.Textbox(
                    label="Queries (one per line)",
                    value="dog\nRecords\nBennett",
                    lines=5,
                    placeholder="Enter one query per line…",
                )
                run_btn = gr.Button("Retrieve", variant="primary")
            with gr.Column(scale=1):
                stats_output = gr.Textbox(label="Sparse vector stats", interactive=False)
                tokens_output = gr.HTML(label="Top activated vocabulary tokens")
                scores_output = gr.HTML(label="Query similarity scores")

        with gr.Accordion("Advanced settings", open=False):
            topk_slider = gr.Slider(
                minimum=5, maximum=50, value=15, step=1,
                label="Top-k tokens to display",
            )

        gr.Examples(
            examples=[
                ["sample_page.png", "dog\nRecords\nBennett", 15],
                ["sample_page.png", "dog breed\nveterinary\ntraining", 15],
                ["sample_page.png", "puppy\nchampionship\npedigree", 20],
            ],
            inputs=[image_input, queries_input, topk_slider],
            outputs=[tokens_output, scores_output, stats_output],
            fn=retrieve,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=retrieve,
        inputs=[image_input, queries_input, topk_slider],
        outputs=[tokens_output, scores_output, stats_output],
        api_name="retrieve",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)