"""Smoke test for DynamicSimilaritySampler: re-embed + batch composition.

Loads checkpoint, embeds a subset of train drones through model.encode_query,
feeds embeddings to the sampler, and verifies:
  - Sampler produces batches
  - Mutex constraint preserved
  - Batches differ from plain mutex (sorted by similarity to seed)
"""
import time
import torch
from torch.utils.data import DataLoader

from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
from src.datasets.gtauav_dataset import (
    GTAUAVDataset,
    GTAUAVDroneQuery,
    collate_drone_query,
)
from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform

CKPT = "out/gtauav/with_text/ckpt_epoch005.pt"
N_SUBSET = 1024  # small subset for smoke speed


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

    tf = get_dino_transform(image_size=256)
    ds = GTAUAVDataset(
        pair_json="meta/train_80.json",
        filter_meta="meta/seg_filter.json",
        image_transform=tf,
    )
    # Subset the entries for smoke speed
    ds.entries = ds.entries[:N_SUBSET]
    print(f"Using {len(ds)} train entries")

    query_ds = GTAUAVDroneQuery(ds)
    loader = DataLoader(
        query_ds, batch_size=32, shuffle=False, num_workers=2,
        collate_fn=collate_drone_query, pin_memory=True,
    )

    # Embed
    t0 = time.time()
    embs = []
    with torch.no_grad():
        for batch in loader:
            drone_img = batch["drone_img"].to("cuda", non_blocking=True)
            q = model.encode_query(
                drone_img,
                batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
            )
            embs.append(q.cpu())
    emb = torch.cat(embs, dim=0)
    print(f"Embedded {emb.shape[0]} queries in {time.time()-t0:.1f}s, dim={emb.shape[1]}")

    # Build sampler
    sat_cand = [e["sat_candidates"] for e in ds.entries]
    sampler = DynamicSimilaritySampler(sat_cand, batch_size=8, seed=42)
    sampler.set_epoch(0)

    # Mutex fallback
    mutex_batches = list(iter(sampler))
    print(f"Mutex-only: {len(mutex_batches)} batches")

    # DSS mode
    sampler.update_embeddings(emb)
    t0 = time.time()
    dss_batches = list(iter(sampler))
    print(f"DSS: {len(dss_batches)} batches in {time.time()-t0:.2f}s")

    # Verify mutex invariant in DSS output
    n_ok = 0
    for b in dss_batches[:20]:
        all_sats = set()
        overlap = False
        for idx in b:
            s = set(ds.entries[idx]["sat_candidates"])
            if s & all_sats:
                overlap = True; break
            all_sats |= s
        if not overlap: n_ok += 1
    print(f"Mutex preserved: {n_ok}/20 DSS batches clean")

    # Similarity check: within a DSS batch, drones should be more similar
    # to each other than drones in a random mutex batch.
    def batch_mean_pairwise_sim(batch):
        e = emb[batch]
        e = torch.nn.functional.normalize(e, dim=-1)
        sim = e @ e.t()
        # Exclude diagonal
        mask = ~torch.eye(len(batch), dtype=torch.bool)
        return sim[mask].mean().item()

    mutex_sim = sum(batch_mean_pairwise_sim(b) for b in mutex_batches[:20]) / 20
    dss_sim = sum(batch_mean_pairwise_sim(b) for b in dss_batches[:20]) / 20
    print(f"Mean in-batch cosine sim — mutex: {mutex_sim:.4f}  DSS: {dss_sim:.4f}")
    print(f"  → DSS batches {'MORE' if dss_sim > mutex_sim else 'LESS'} visually similar (expected: MORE)")


if __name__ == "__main__":
    main()