"""Minimal training smoke test: 2 batches forward+backward.

Verifies end-to-end that MutuallyExclusiveSampler + InfoNCELoss +
per-sample caption masking compose correctly for training.
"""
import torch
from torch.utils.data import DataLoader

from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
from src.losses.hard_negatives import NegativeMemoryBank
from src.losses.multi_infonce import InfoNCELoss
from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform

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


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

    tf = get_dino_transform(image_size=256)
    ds = GTAUAVDataset(
        pair_json="meta/train_80.json",
        filter_meta="meta/seg_filter.json",
        drone_transform=tf,
        sat_transform=tf,
    )

    sampler = MutuallyExclusiveSampler(
        [e["sat_candidates"] for e in ds.entries],
        batch_size=8, shuffle=True, seed=42,
    )
    sampler.set_epoch(0)
    loader = DataLoader(
        ds, batch_sampler=sampler, num_workers=2,
        collate_fn=collate_gtauav_batch, pin_memory=True,
    )

    loss_fn = InfoNCELoss(
        temperature_init=0.07, learnable_temperature=True,
        label_smoothing=0.1, weight_q2g=0.6, weight_g2q=0.4,
        tau_min=0.01, tau_max=0.1,
        hard_mining_k=512,
    ).to("cuda")
    neg_bank = NegativeMemoryBank(size=4096, dim=model.embed_dim).to("cuda")

    trainable = [p for p in model.trainable_parameters()] + list(loss_fn.parameters())
    opt = torch.optim.AdamW(trainable, lr=1e-4)

    it = iter(loader)
    for step in range(3):
        batch = next(it)
        opt.zero_grad()
        emb = model(
            drone_img=batch["drone_img"].to("cuda", non_blocking=True),
            sat_img=batch["sat_img"].to("cuda", non_blocking=True),
            caption_l1=batch["caption_l1"], caption_l2=batch["caption_l2"], caption_l3=batch["caption_l3"],
            sat_caption_l1=batch["sat_caption_l1"], sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"],
        )
        queue = neg_bank.get_queue()
        out = loss_fn(emb, epoch=0, total_epochs=10, queue_negatives=queue)
        out["total"].backward()
        opt.step()
        neg_bank.enqueue(emb["gallery"].detach())

        # Verify mutual exclusion in batch
        batch_sats = [set(ds.entries[i]["sat_candidates"]) for i in batch.get("__indices__", range(8))]
        # We can also check via sat_names (one sat per drone sampled)
        sat_names = batch["sat_names"]
        print(
            f"  step {step}: loss={out['total'].item():.4f}  "
            f"tau={out['temperature'].item():.4f}  "
            f"gate_q={out['gate_q'].item():.3f}  gate_g={out['gate_g'].item():.3f}  "
            f"queue_size={queue.shape[0] if queue is not None else 0}  "
            f"mining_k={loss_fn.hard_mining_k if queue is not None and queue.shape[0] > loss_fn.hard_mining_k else 'full'}"
        )
        assert torch.isfinite(out["total"]).all(), "Loss not finite!"

    print("OK: 3 train steps completed with finite loss (hard mining K=512)")


if __name__ == "__main__":
    main()