"""Smoke-test for the rewritten `_evaluate` function.

Loads checkpoint ckpt_epoch005.pt and runs the new full-gallery eval with
max_batches=5 to verify end-to-end without waiting for a full epoch.
"""
from torch.utils.data import DataLoader

from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
from src.training.train_gtauav import _evaluate

CKPT = "out/gtauav/with_text/ckpt_epoch005.pt"

def main() -> None:
    model, _ = AsymmetricEncoder.load_checkpoint(CKPT, device="cuda")

    eval_tf = get_dino_transform(image_size=256)
    ds = GTAUAVDataset(
        pair_json="meta/test_20.json",
        filter_meta="meta/seg_filter.json",
        image_transform=eval_tf,
    )
    loader = DataLoader(
        ds,
        batch_size=32,
        shuffle=False,
        num_workers=2,
        collate_fn=collate_gtauav_batch,
        pin_memory=True,
    )

    print("Running _evaluate (max_batches=5 on queries, full gallery)...")
    metrics = _evaluate(
        model=model,
        loader=loader,
        device="cuda",
        max_batches=5,
        desc="smoke",
    )
    print("--- metrics ---")
    for k, v in metrics.items():
        print(f"  {k}: {v:.4f}" if isinstance(v, float) else f"  {k}: {v}")


if __name__ == "__main__":
    main()