from __future__ import annotations

"""Batch sampler that prevents false negatives from GTA-UAV's semi-positive graph.

In GTA-UAV a single satellite tile can be a valid (semi-)positive for multiple
drone frames. When those frames land in the same InfoNCE batch, the shared
satellite becomes a negative for every drone except one — which trains the
model to push apart embeddings that should actually be close.

MutuallyExclusiveSampler resolves this by greedily building batches where no
two drone indices share ANY entry in their `sat_candidates` set. This keeps
the diagonal InfoNCE formulation valid: every off-diagonal satellite is a
genuine negative for every query in the row/column.

References: Zhu et al., Game4Loc (arXiv:2409.16925), §3.2 "Sample ID".
"""

import logging
import random
from typing import Iterator, Sequence

from torch.utils.data.sampler import Sampler

LOGGER = logging.getLogger("caption_test.mutex_sampler")


class MutuallyExclusiveSampler(Sampler[list[int]]):
    """Batch sampler yielding drone-index lists with disjoint sat_candidates.

    Args:
        sat_candidates_per_item: For each dataset index i, the list of
            satellite names considered valid matches (positive + semi-positive).
        batch_size: Target batch size. Partial batches are dropped.
        shuffle: Shuffle the index pool each epoch (use set_epoch for reproducibility).
        seed: Base RNG seed — the effective seed is `seed + epoch`.
        allow_partial: If True, yield the trailing partial batch. Default False.
    """

    def __init__(
        self,
        sat_candidates_per_item: Sequence[Sequence[str]],
        batch_size: int,
        shuffle: bool = True,
        seed: int = 0,
        allow_partial: bool = False,
    ) -> None:
        if batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {batch_size}")
        self._item_sats: list[frozenset[str]] = [
            frozenset(s) for s in sat_candidates_per_item
        ]
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.seed = seed
        self.allow_partial = allow_partial
        self.epoch = 0
        self._cached_len: int | None = None

    def set_epoch(self, epoch: int) -> None:
        """Advance epoch (invalidates cached length) — call from training loop."""
        self.epoch = epoch
        self._cached_len = None

    def _generate_batches(self) -> list[list[int]]:
        rng = random.Random(self.seed + self.epoch) if self.shuffle else None
        remaining = list(range(len(self._item_sats)))
        if rng is not None:
            rng.shuffle(remaining)

        batches: list[list[int]] = []

        # Each outer iteration produces at most one batch. Items that conflict
        # with the batch's claimed-sat set roll over to the next iteration.
        while remaining:
            batch: list[int] = []
            claimed: set[str] = set()
            next_remaining: list[int] = []

            for idx in remaining:
                sats = self._item_sats[idx]
                if len(batch) < self.batch_size and not (sats & claimed):
                    batch.append(idx)
                    claimed |= sats
                else:
                    next_remaining.append(idx)

            if len(batch) == self.batch_size:
                batches.append(batch)
            elif self.allow_partial and batch:
                batches.append(batch)
                break  # leftover rollover produces no more full batches
            else:
                break  # drop partial trailing batch

            remaining = next_remaining
            if rng is not None:
                rng.shuffle(remaining)

        return batches

    def __iter__(self) -> Iterator[list[int]]:
        batches = self._generate_batches()
        self._cached_len = len(batches)
        for batch in batches:
            yield batch

    def __len__(self) -> int:
        if self._cached_len is None:
            # Estimate by actually generating — correct count needed by
            # DataLoader/tqdm. Cached until next set_epoch.
            self._cached_len = len(self._generate_batches())
        return self._cached_len
