from __future__ import annotations

"""Dynamic Similarity Sampler — hard batch construction via embedding kNN.

Builds on MutuallyExclusiveSampler: keeps the no-false-negatives guarantee
(disjoint `sat_candidates` within a batch), but additionally pressures the
model by packing visually-similar drone queries together. The intuition is
that in-batch negatives which are easy to distinguish contribute little
gradient — if the batch contains lookalikes, the model has to learn finer
features.

Workflow:
  1. Training loop calls `update_embeddings(query_embeddings)` at the start
     of each epoch (or every N epochs) with current query embeddings.
  2. `_generate_batches` picks a random seed drone, ranks the rest by cosine
     similarity to the seed, then greedily packs highest-similarity drones
     whose `sat_candidates` don't conflict.
  3. If no embeddings are set yet, falls back to mutex-only behavior (first
     epoch, or warmup before the model has useful representations).
"""

import logging
import random
from typing import Iterator, Sequence

import torch
import torch.nn.functional as F
from torch.utils.data.sampler import Sampler

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


class DynamicSimilaritySampler(Sampler[list[int]]):
    """Batch sampler: visually-similar drones packed per batch, mutex-preserved.

    Args:
        sat_candidates_per_item: Valid satellite names per dataset index
            (positive + semi-positive).
        batch_size: Target batch size. Partial trailing batches are dropped.
        shuffle: Randomize seed ordering each epoch.
        seed: Base RNG seed (effective seed is `seed + epoch`).
        allow_partial: Yield trailing partial batch if non-empty.
    """

    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._embeddings: torch.Tensor | None = None
        self._cached_len: int | None = None

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch
        self._cached_len = None

    def update_embeddings(self, embeddings: torch.Tensor) -> None:
        """Store query embeddings [N, D] for the next epoch's sampling.

        Call from the training loop before starting each epoch where you want
        DSS to be active. Stored on CPU as float32 and L2-normalized.
        """
        if embeddings.size(0) != len(self._item_sats):
            raise ValueError(
                f"embeddings have {embeddings.size(0)} rows but sampler tracks "
                f"{len(self._item_sats)} items",
            )
        self._embeddings = F.normalize(embeddings.detach().cpu().float(), dim=-1)
        self._cached_len = None
        LOGGER.info(
            "DSS embeddings updated: %d × %d", *self._embeddings.shape,
        )

    def clear_embeddings(self) -> None:
        """Revert to mutex-only sampling for the next epoch."""
        self._embeddings = None
        self._cached_len = None

    def _generate_batches_mutex_only(self) -> list[list[int]]:
        """Fallback: greedy mutex packing without similarity ranking."""
        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]] = []
        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
            else:
                break
            remaining = next_remaining
            if rng is not None:
                rng.shuffle(remaining)
        return batches

    def _generate_batches_dss(self) -> list[list[int]]:
        """Similarity-guided batches using stored embeddings."""
        assert self._embeddings is not None
        rng = random.Random(self.seed + self.epoch) if self.shuffle else None
        emb = self._embeddings  # already L2-normalized, on CPU
        n = emb.size(0)

        remaining = set(range(n))
        batches: list[list[int]] = []

        while len(remaining) >= self.batch_size:
            seed_idx = (
                rng.choice(list(remaining))
                if rng is not None
                else next(iter(remaining))
            )

            # Cosine similarity from seed to all items (fp32 matmul, ~25K ops).
            sims = emb @ emb[seed_idx]  # [N]
            order = sims.argsort(descending=True).tolist()

            batch: list[int] = []
            claimed: set[str] = set()
            for idx in order:
                if idx not in remaining:
                    continue
                sats = self._item_sats[idx]
                if sats & claimed:
                    continue
                batch.append(idx)
                claimed |= sats
                if len(batch) == self.batch_size:
                    break

            if len(batch) < self.batch_size:
                # Couldn't fill from this seed's neighborhood (heavy mutex
                # conflict). Drop the seed and retry with remaining pool.
                remaining.discard(seed_idx)
                continue

            batches.append(batch)
            for i in batch:
                remaining.discard(i)

        if self.allow_partial and remaining:
            # Mutex-pack whatever's left.
            claimed = set()
            tail: list[int] = []
            for idx in remaining:
                sats = self._item_sats[idx]
                if not (sats & claimed):
                    tail.append(idx)
                    claimed |= sats
            if tail:
                batches.append(tail)

        return batches

    def _generate_batches(self) -> list[list[int]]:
        if self._embeddings is None:
            return self._generate_batches_mutex_only()
        return self._generate_batches_dss()

    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:
            self._cached_len = len(self._generate_batches())
        return self._cached_len