"""Offline export of the SigLIP2 text embeddings for RoofSight's fixed prompts.

Run this ONCE on a machine with RAM headroom (~14 GB peak). It loads the full
`siglip2-g` adaptor — which instantiates the entire google/siglip2-giant model
just to keep its text tower — encodes the fixed prompt set, and writes the
resulting (num_prompts, D) tensor to disk.

The Space then loads with the `siglip2` adaptor factory removed from RADIO's
registry, so the text tower is never built there. See radio_backbone.load_model.

Also captures a reference feature map so the stripped runtime path can be
verified as numerically identical:

    python export_text_embeddings.py            # export + reference
    python verify_stripped_adaptor.py           # compare
"""

import os
import sys

import numpy as np
import torch
from PIL import Image

from radio_backbone import ROOF_PROMPTS, NON_ROOF_PROMPTS, PATCH_SIZE

OUT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "text_embeddings.pt")
REF_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_ref_features.pt")
REF_IMAGE = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "_run_output", "1310_susan_circle", "original.png",
)

# CPU + fp32 throughout: the reference must be deterministic so the parity
# check measures the adaptor swap, not autocast/GPU nondeterminism.
DEVICE = "cpu"


def main():
    all_labels = ROOF_PROMPTS + NON_ROOF_PROMPTS
    print(f"Loading FULL c-radio_v4-h + siglip2-g adaptor on {DEVICE} (this is the heavy path)...")

    model = torch.hub.load(
        "NVlabs/RADIO",
        "radio_model",
        version="c-radio_v4-h",
        adaptor_names=["siglip2-g"],
        vitdet_window_size=8,
        progress=True,
        skip_validation=True,
        trust_repo=True,
    )
    model.eval()
    model.to(DEVICE)

    adaptor = model.adaptors["siglip2-g"]
    assert hasattr(adaptor, "encode_text"), (
        "adaptor has no encode_text — the siglip2 factory was not used. "
        "This script must run WITHOUT the registry patch."
    )

    print(f"Encoding {len(all_labels)} prompts...")
    text_input = adaptor.tokenizer(all_labels)
    text_input = {k: v.to(DEVICE) for k, v in text_input.items()}
    with torch.no_grad():
        embeddings = adaptor.encode_text(text_input, normalize=True).float().cpu().clone()

    print(f"  embeddings: {tuple(embeddings.shape)}  dtype={embeddings.dtype}")
    torch.save({"prompts": all_labels, "embeddings": embeddings}, OUT_PATH)
    print(f"  wrote {OUT_PATH} ({os.path.getsize(OUT_PATH) / 1024:.1f} KB)")

    # --- Reference feature map, for the parity check ---
    if not os.path.exists(REF_IMAGE):
        print(f"\nNo reference image at {REF_IMAGE}; skipping reference capture.")
        return

    print(f"\nCapturing reference features from {REF_IMAGE}...")
    image = np.array(Image.open(REF_IMAGE).convert("RGB"))
    x = torch.from_numpy(image).permute(2, 0, 1).float().div_(255.0).unsqueeze(0)
    snapped = model.get_nearest_supported_resolution(*x.shape[-2:])
    if snapped != x.shape[-2:]:
        x = torch.nn.functional.interpolate(x, snapped, mode="bilinear", align_corners=False)

    with torch.no_grad():
        out = model(x.to(DEVICE))
    _, features = out["siglip2-g"]

    torch.save(
        {
            "features": features.float().cpu().clone(),
            "snapped": tuple(snapped),
            "image_shape": image.shape,
        },
        REF_PATH,
    )
    print(f"  features: {tuple(features.shape)}")
    print(f"  wrote {REF_PATH}")
    print("\nDone. Now run: python verify_stripped_adaptor.py")


if __name__ == "__main__":
    sys.exit(main())