#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Regenerate `image-smoke-reference.json`: fp32 reference embeddings of the shipped test images (passed as file
paths, no prompt, exactly as sentence-transformers embeds images), of one search query per image and of two
fictional documents, computed with sentence-transformers on the ORIGINAL google/embeddinggemma-2 checkpoint
(CPU, float32).

    python3 make_image_smoke_reference.py --model google/embeddinggemma-2 --out image-smoke-reference.json

Needs: sentence-transformers >= 6.1, transformers >= 5.19, torch, torchvision, pillow. The shipped file was produced
with sentence-transformers 6.1.0 / transformers 5.19.0 / torch 2.14.1 / torchvision 0.29.1.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

HERE = Path(__file__).resolve().parent
TEXTS = [
    {"kind": "document", "text": "The Lumen lamp range sold best in Westmere this quarter; Lake Tamsin stayed below two hundred units."},
    {"kind": "document", "text": "Verbale della riunione: il lancio delle tende Nimbus slitta a marzo, i preordini aprono con sei rivenditori."},
]
SMOKE_IMAGES = ["chart_bar", "ui_login", "slide_launch", "doc_invoice", "diagram_flow", "scene_sunset", "shape_red_circle",
                "icon_rgba", "gray_letter", "palette_chart"]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--model", default="google/embeddinggemma-2")
    ap.add_argument("--images", default=str(HERE / "images"))
    ap.add_argument("--out", default=str(HERE / "image-smoke-reference.json"))
    ap.add_argument("--threshold", type=float, default=0.999)
    ap.add_argument("--threshold-image", type=float, default=0.998)
    args = ap.parse_args()
    import sentence_transformers
    import torch
    import torchvision
    import transformers
    from sentence_transformers import SentenceTransformer

    m = SentenceTransformer(args.model, device="cpu", model_kwargs={"torch_dtype": torch.float32}, config_kwargs={"audio_config": None})
    m.eval()
    items = {it["name"]: it for it in json.loads((Path(args.images) / "manifest.json").read_text())}
    out = {
        "reference": "sentence-transformers fp32 on CPU, original checkpoint; images passed as paths with no prompt",
        "versions": {"sentence_transformers": sentence_transformers.__version__, "transformers": transformers.__version__,
                     "torch": torch.__version__, "torchvision": torchvision.__version__},
        "threshold_cosine": args.threshold,
        "threshold_cosine_image": args.threshold_image,
        "images": [],
        "texts": [],
    }
    for name in SMOKE_IMAGES:
        it = items[name]
        path = str(Path(args.images) / it["file"])
        with torch.no_grad():
            v = m.encode([{"image": path}], batch_size=1, normalize_embeddings=True, convert_to_numpy=True)[0]
        out["images"].append({"file": f"{Path(args.images).name}/{it['file']}", "query": it["query"], "embedding": [round(float(x), 7) for x in v.tolist()]})
        print("image", it["file"], len(v))
    for name in SMOKE_IMAGES:
        q = items[name]["query"]
        v = m.encode(q, prompt_name="query", normalize_embeddings=True, convert_to_numpy=True)
        out["texts"].append({"kind": "search_query", "text": q, "embedding": [round(float(x), 7) for x in v.tolist()]})
    for t in TEXTS:
        v = m.encode(t["text"], prompt_name="document", normalize_embeddings=True, convert_to_numpy=True)
        out["texts"].append({**t, "embedding": [round(float(x), 7) for x in v.tolist()]})
    Path(args.out).write_text(json.dumps(out, ensure_ascii=False, indent=1), encoding="utf-8")
    print("wrote", args.out)


if __name__ == "__main__":
    main()