"""Cache frozen COCO-80 class text embeddings from a CLIP text encoder.

Runs once, saves (80, 512) L2-normalized fp32 tensor to disk. The detection
head loads this as a frozen buffer and computes cosine similarity against
learned feature projections.
"""
import os
import torch
from transformers import CLIPTextModel, CLIPTokenizer

COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
    "toothbrush",
]

# Prompt template matches CLIP's training distribution better than bare class names.
TEMPLATE = "a photo of a {}"


def main():
    out_path = os.environ.get("COCO_TEXT_EMBED_PATH",
                              "/home/zootest/datasets/coco/coco_text_embed_clip_vitb32.pt")
    model_name = "openai/clip-vit-base-patch32"
    print(f"Loading {model_name}...")
    tok = CLIPTokenizer.from_pretrained(model_name)
    enc = CLIPTextModel.from_pretrained(model_name).cuda().eval()

    prompts = [TEMPLATE.format(c) for c in COCO_CLASSES]
    toks = tok(prompts, padding=True, return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = enc(**toks)
    pooled = out.pooler_output  # (80, 512)
    pooled = torch.nn.functional.normalize(pooled.float(), p=2, dim=-1)

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    torch.save({
        "embeddings": pooled.cpu(),
        "classes": COCO_CLASSES,
        "model": model_name,
        "template": TEMPLATE,
    }, out_path)
    print(f"Saved {pooled.shape} to {out_path}")


if __name__ == "__main__":
    main()