"""Mutually-exclusive batch sampler для diagonal InfoNCE (протокол §6.4).

Проблема, которую решает (см. EXPERIMENT_SETUP.md §8):
    Loss — диагональный InfoNCE: в батче B пар, правильные на диагонали, всё
    вне диагонали считается негативами. На GTA-UAV один спутниковый тайл часто
    накрывает несколько дронов (avg 4.58 semi-позитивов). При случайном shuffle
    два дрона с ОБЩИМ позитивным тайлом попадают в один батч → этот тайл для
    одного дрона стоит вне диагонали и ошибочно штрафуется как негатив
    (false-negative), зашумляя градиент.

Решение:
    Формируем батчи так, чтобы множества позитивных тайлов любых двух дронов
    в одном батче НЕ пересекались. Тогда, какой бы тайл ни был сэмплирован на
    диагональ, в батче нет второго дрона, для которого он позитив.

Гарантия консервативная (по ПОЛНЫМ спискам pos+semi-pos), поэтому от выбора
конкретного тайла в __getitem__ результат не зависит.

Порядок воспроизводим: перестановка зависит только от seed и номера эпохи
(инкрементируется на каждый __iter__), что согласуется с set_seed в train.py.
"""
from __future__ import annotations

import logging

import torch
from torch.utils.data import Sampler

LOGGER = logging.getLogger("cvgl.data.me_sampler")


class MutuallyExclusiveBatchSampler(Sampler[list[int]]):
    """Batch-sampler: в одном батче нет двух элементов с общим позитивным тайлом.

    Жадное bin-packing first-fit: элементы обходятся в случайном (по seed+epoch)
    порядке; каждый кладётся в первый открытый батч, чьи уже занятые тайлы не
    пересекаются с его позитивами и который ещё не заполнен, иначе открывается
    новый батч. Заполненные до batch_size батчи отдаются наружу.

    Args:
        positive_sets: Для каждого элемента датасета — множество id позитивных
            тайлов (int). Пустое множество допустимо (элемент совместим с любым
            батчом).
        batch_size: Размер батча.
        seed: База для генератора перестановок (совпадает с train seed).
        drop_last: Отбрасывать ли неполные батчи (True — как в обучении).
    """

    def __init__(
        self,
        positive_sets: list[set[int]],
        batch_size: int,
        seed: int = 42,
        drop_last: bool = True,
    ) -> None:
        self.positive_sets = [frozenset(s) for s in positive_sets]
        self.batch_size = batch_size
        self.seed = seed
        self.drop_last = drop_last
        self.n = len(self.positive_sets)
        self._epoch = 0

    def set_epoch(self, epoch: int) -> None:
        """Явно задать номер эпохи (опционально; иначе инкремент на __iter__)."""
        self._epoch = epoch

    def __iter__(self):
        g = torch.Generator()
        g.manual_seed(self.seed + self._epoch)
        order = torch.randperm(self.n, generator=g).tolist()
        self._epoch += 1

        open_batches: list[dict] = []  # [{"members": [...], "used": set()}]
        n_yielded = 0

        for idx in order:
            pos = self.positive_sets[idx]
            placed = False
            for b in open_batches:
                if len(b["members"]) < self.batch_size and b["used"].isdisjoint(pos):
                    b["members"].append(idx)
                    b["used"] |= pos
                    placed = True
                    if len(b["members"]) == self.batch_size:
                        open_batches.remove(b)
                        n_yielded += 1
                        yield b["members"]
                    break
            if not placed:
                open_batches.append({"members": [idx], "used": set(pos)})

        # Хвост: неполные батчи.
        for b in open_batches:
            if not self.drop_last:
                n_yielded += 1
                yield b["members"]

        LOGGER.debug(
            "MutuallyExclusiveBatchSampler: %d батчей выдано (n=%d, bs=%d, drop_last=%s)",
            n_yielded, self.n, self.batch_size, self.drop_last,
        )

    def __len__(self) -> int:
        # Приблизительно: точное число батчей зависит от коллизий и известно
        # только после прохода. Для drop_last=True это верхняя оценка.
        if self.drop_last:
            return self.n // self.batch_size
        return (self.n + self.batch_size - 1) // self.batch_size
