# /// script
# requires-python = ">=3.10"
# dependencies = ["transformers>=4.30", "torch", "onnx"]
# ///
"""Export the CLIP ViT-B/32 text encoder to ONNX and save vocab files.

Compatible with all YOLOv8-worldv2 sizes (s/m/l/x) — run once and reuse.

The C++ node uses these three files to compute text embeddings at runtime:
  clip_text_encoder.onnx  input:  input_ids  (1, 77) int64
                          output: text_embeds (1, 512) float32  (L2-normalize in C++)
  clip_vocab.txt          token → id mapping (tab-separated, one per line)
  clip_merges.txt         BPE merge rules (space-separated pairs, one per line)

Usage:
    uv run scripts/export_clip_text_encoder.py
    uv run scripts/export_clip_text_encoder.py --output-dir models/
"""

import argparse
import sys
from pathlib import Path

CLIP_MODEL = "openai/clip-vit-base-patch32"


def parse_args():
    p = argparse.ArgumentParser(description="Export CLIP ViT-B/32 text encoder to ONNX")
    p.add_argument(
        "--output-dir", "-o",
        default="models",
        help="Directory for output files. Default: models/",
    )
    return p.parse_args()


def main():
    args = parse_args()

    try:
        import torch
        from transformers import CLIPTextModelWithProjection, CLIPTokenizer
    except ImportError as e:
        print(f"Error: {e}\nRun: uv run scripts/export_clip_text_encoder.py", file=sys.stderr)
        sys.exit(1)

    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)

    print(f"Loading {CLIP_MODEL} ...")
    model = CLIPTextModelWithProjection.from_pretrained(CLIP_MODEL)
    tokenizer = CLIPTokenizer.from_pretrained(CLIP_MODEL)
    model.eval()

    # --- Export ONNX text encoder ---
    # Input:  input_ids   (1, 77) int64   — BPE token IDs, padded/truncated to 77
    # Output: text_embeds (1, 512) float32 — projected features (NOT normalized; C++ normalizes)
    onnx_path = out / "clip_text_encoder.onnx"
    print(f"Exporting ONNX → {onnx_path} ...")
    dummy = torch.zeros(1, 77, dtype=torch.long)
    torch.onnx.export(
        model,
        dummy,
        str(onnx_path),
        input_names=["input_ids"],
        output_names=["text_embeds"],
        dynamic_axes={
            "input_ids":   {0: "batch_size"},
            "text_embeds": {0: "batch_size"},
        },
        opset_version=14,
    )

    # --- Save vocab as tab-separated text (token\tid) ---
    # Avoids JSON parsing in C++. Tokens with special chars are escaped.
    vocab_path = out / "clip_vocab.txt"
    print(f"Saving vocab  → {vocab_path} ...")
    vocab = tokenizer.get_vocab()
    with open(vocab_path, "w", encoding="utf-8") as f:
        for token, idx in sorted(vocab.items(), key=lambda x: x[1]):
            token_esc = (token
                         .replace("\\", "\\\\")
                         .replace("\t", "\\t")
                         .replace("\n", "\\n"))
            f.write(f"{token_esc}\t{idx}\n")

    # --- Save BPE merge rules ---
    merges_path = out / "clip_merges.txt"
    print(f"Saving merges → {merges_path} ...")
    with open(merges_path, "w", encoding="utf-8") as f:
        for (a, b) in tokenizer.bpe_ranks.keys():
            f.write(f"{a} {b}\n")

    print(f"\nDone. Files written to {out.resolve()}/")
    print(f"  clip_text_encoder.onnx  — ONNX model")
    print(f"  clip_vocab.txt          — {len(vocab)} tokens")
    print(f"  clip_merges.txt         — {len(tokenizer.bpe_ranks)} merge rules")
    print(f"\nSet env vars:")
    print(f"  ZED_ONNX_TEXT_ENCODER_PATH={out.resolve()}/clip_text_encoder.onnx")
    print(f"  ZED_ONNX_CLIP_VOCAB={out.resolve()}/clip_vocab.txt")
    print(f"  ZED_ONNX_CLIP_MERGES={out.resolve()}/clip_merges.txt")


if __name__ == "__main__":
    main()