"""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.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,
    ).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(2):
        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"],
        )
        out = loss_fn(emb, epoch=0, total_epochs=10)
        out["total"].backward()
        opt.step()

        # 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"n_drone_caps={sum(1 for t in batch['caption_l1'] if t)}  "
            f"n_sat_caps={sum(1 for t in batch['sat_caption_l1'] if t)}"
        )
        assert torch.isfinite(out["total"]).all(), "Loss not finite!"

    print("OK: 2 train steps completed with finite loss")


if __name__ == "__main__":
    main()