from __future__ import annotations

"""Hard negative memory bank for contrastive learning.

MoCo-style FIFO queue of recent gallery embeddings. Each batch gets
B in-batch negatives + Q queue negatives, significantly increasing
the effective number of negatives without extra VRAM for forward pass.

Usage:
    bank = NegativeMemoryBank(size=4096, dim=1024)
    # In training loop:
    sim = bank.compute_similarity(query, gallery)  # [B, B + Q]
    bank.enqueue(gallery.detach())
"""

import torch
import torch.nn as nn


class NegativeMemoryBank(nn.Module):
    """FIFO queue of detached gallery embeddings for hard negatives.

    Args:
        size: Queue capacity (number of stored embeddings).
        dim: Embedding dimension.
    """

    def __init__(self, size: int = 4096, dim: int = 1024) -> None:
        super().__init__()
        self.size = size
        self.dim = dim
        # Queue stored as buffer (not a parameter, moves with .to(device)).
        self.register_buffer("queue", torch.randn(size, dim))
        self.queue = nn.functional.normalize(self.queue, dim=-1)
        self.register_buffer("ptr", torch.zeros(1, dtype=torch.long))
        self.register_buffer("full", torch.zeros(1, dtype=torch.bool))

    @torch.no_grad()
    def enqueue(self, embeddings: torch.Tensor) -> None:
        """Add embeddings to the queue (FIFO). Oldest are overwritten."""
        batch_size = embeddings.shape[0]
        ptr = int(self.ptr.item())

        if ptr + batch_size <= self.size:
            self.queue[ptr:ptr + batch_size] = embeddings.detach()
        else:
            # Wrap around.
            overflow = (ptr + batch_size) - self.size
            self.queue[ptr:] = embeddings[:batch_size - overflow].detach()
            self.queue[:overflow] = embeddings[batch_size - overflow:].detach()

        new_ptr = (ptr + batch_size) % self.size
        self.ptr[0] = new_ptr
        if not self.full.item() and (new_ptr < ptr or new_ptr == 0):
            self.full[0] = True

    def get_queue(self) -> torch.Tensor:
        """Return valid queue entries [Q, dim]."""
        if self.full.item():
            return self.queue
        ptr = int(self.ptr.item())
        if ptr == 0:
            return self.queue[:0]  # empty
        return self.queue[:ptr]

    @property
    def current_size(self) -> int:
        if self.full.item():
            return self.size
        return int(self.ptr.item())
