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

from src.datasets.lsh_index import LSHIndex

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,
        knn_device: str = "cpu",
        use_lsh: bool = False,
        lsh_num_tables: int = 8,
        lsh_num_bits: int = 14,
    ) -> 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.knn_device = knn_device
        self.use_lsh = use_lsh
        self.lsh_num_tables = lsh_num_tables
        self.lsh_num_bits = lsh_num_bits
        self.epoch = 0
        self._embeddings: torch.Tensor | None = None
        self._lsh: LSHIndex | 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. Embeddings are L2-normalized and moved to
        `self.knn_device` for the similarity step. If `use_lsh=True`, builds
        an LSH index over the embeddings for approximate candidate lookup.
        """
        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",
            )
        emb = F.normalize(embeddings.detach().float(), dim=-1)
        self._embeddings = emb.to(self.knn_device)
        self._cached_len = None

        if self.use_lsh:
            # LSH builds on CPU (bucket dicts are Python-native).
            self._lsh = LSHIndex(
                dim=emb.size(1),
                num_tables=self.lsh_num_tables,
                num_bits=self.lsh_num_bits,
                seed=self.seed,
            )
            self._lsh.build(emb.cpu())
        else:
            self._lsh = None

        LOGGER.info(
            "DSS embeddings updated: %d × %d on %s%s",
            self._embeddings.shape[0], self._embeddings.shape[1],
            self.knn_device,
            f" + LSH (tables={self.lsh_num_tables}, bits={self.lsh_num_bits})" if self.use_lsh else "",
        )

    def clear_embeddings(self) -> None:
        """Revert to mutex-only sampling for the next epoch."""
        self._embeddings = None
        self._lsh = 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.

        For each seed we compute (or retrieve from LSH) an order of
        candidates ranked by cosine similarity, then greedy-pack respecting
        mutex constraints. Runs on `knn_device` — GPU is strongly preferred
        for large N.
        """
        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 knn_device
        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))
            )

            if self._lsh is not None:
                # LSH path: restrict ranking to the candidate pool for this seed
                # (union of buckets across tables), then exact-rank the candidates.
                cand = self._lsh.candidates(emb[seed_idx].cpu())
                cand &= remaining
                if len(cand) < self.batch_size:
                    # Candidate pool too small — fall back to full ranking to
                    # preserve batch-fill guarantees.
                    sims = emb @ emb[seed_idx]
                    order = sims.argsort(descending=True).cpu().tolist()
                else:
                    cand_list = list(cand)
                    sims = emb[cand_list] @ emb[seed_idx]  # [|cand|]
                    inner_order = sims.argsort(descending=True).cpu().tolist()
                    order = [cand_list[i] for i in inner_order]
            else:
                sims = emb @ emb[seed_idx]  # [N]
                order = sims.argsort(descending=True).cpu().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. Drop the seed
                # and retry with the remaining pool.
                remaining.discard(seed_idx)
                continue

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

        if self.allow_partial and remaining:
            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
