from __future__ import annotations

"""Random-projection LSH for approximate cosine-similarity kNN.

Builds H independent hash tables, each with B bits. Every embedding hashes to
a B-bit code per table based on its sign against B random hyperplanes.
Two vectors collide in a table with probability (1 − θ/π)^B where θ is the
angle between them, so cosine-similar vectors collide frequently across
tables.

For each query we union the candidate sets from all H tables, then refine by
exact cosine similarity. Typical regime: `num_tables=8, num_bits=14` gives
recall@10 ≥ 0.95 on CVGL-scale features while keeping candidate pools under
a few percent of N.

Not used by default at the current 25K-item scale (exact kNN is fast enough);
this module exists as a drop-in acceleration path when the dataset grows.
"""

import logging

import torch
import torch.nn.functional as F

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


class LSHIndex:
    """Random-projection cosine-LSH index.

    Args:
        dim: Embedding dimension.
        num_tables: Number of independent hash tables (more → higher recall,
            larger candidate pool).
        num_bits: Bits per code (more → lower collision rate per table, so
            you need more tables to compensate). Typical 12-16.
        seed: RNG seed for hyperplane draw (makes index reproducible).
    """

    def __init__(
        self,
        dim: int,
        num_tables: int = 8,
        num_bits: int = 14,
        seed: int = 0,
    ) -> None:
        if num_bits > 62:
            raise ValueError("num_bits must be <= 62 to fit a signed int64 hash")
        self.dim = dim
        self.num_tables = num_tables
        self.num_bits = num_bits

        gen = torch.Generator().manual_seed(seed)
        # Hyperplanes: [num_tables, num_bits, dim] — normalized for numerical
        # stability (sign is unaffected, but magnitudes stay bounded).
        planes = torch.randn(num_tables, num_bits, dim, generator=gen)
        self.planes = F.normalize(planes, dim=-1)

        # Bit weights for packing {0,1}^B → int: 2^0, 2^1, ..., 2^(B-1).
        self._bit_weights = (
            1 << torch.arange(num_bits, dtype=torch.int64)
        )  # [num_bits]

        # Populated by .build(). List of dicts (one per table): hash_int → [item_idx, ...]
        self._tables: list[dict[int, list[int]]] = [{} for _ in range(num_tables)]
        self._n_items = 0

    @torch.no_grad()
    def _hash(self, emb: torch.Tensor) -> torch.Tensor:
        """Hash [N, D] embeddings to [N, num_tables] int64 codes.

        Project each embedding onto every hyperplane and take the sign
        pattern as a bit vector; pack the bits into an int64 per table.
        """
        if emb.dim() != 2:
            raise ValueError(f"expected [N, D], got shape {tuple(emb.shape)}")
        flat_planes = self.planes.reshape(-1, self.dim)  # [T*B, D]
        projections = emb @ flat_planes.t()  # [N, T*B]
        bits = (projections > 0).view(-1, self.num_tables, self.num_bits)  # [N, T, B]
        codes = (bits.to(torch.int64) * self._bit_weights).sum(dim=-1)  # [N, T]
        return codes

    @torch.no_grad()
    def build(self, embeddings: torch.Tensor) -> None:
        """Insert all [N, D] embeddings into every table. Call once per epoch."""
        if embeddings.dim() != 2 or embeddings.size(1) != self.dim:
            raise ValueError(
                f"embeddings must be [N, {self.dim}], got {tuple(embeddings.shape)}",
            )
        for t in range(self.num_tables):
            self._tables[t].clear()
        codes = self._hash(embeddings.float())  # [N, num_tables]
        codes_cpu = codes.cpu().tolist()
        for item_idx, per_table in enumerate(codes_cpu):
            for t, code in enumerate(per_table):
                self._tables[t].setdefault(code, []).append(item_idx)
        self._n_items = embeddings.size(0)
        # Stats for logging.
        bucket_sizes = [
            len(bucket) for table in self._tables for bucket in table.values()
        ]
        if bucket_sizes:
            LOGGER.info(
                "LSH built: N=%d, tables=%d, bits=%d, "
                "avg bucket=%.1f (max=%d)",
                self._n_items, self.num_tables, self.num_bits,
                sum(bucket_sizes) / len(bucket_sizes), max(bucket_sizes),
            )

    @torch.no_grad()
    def candidates(self, query: torch.Tensor) -> set[int]:
        """Union of hash-bucket contents for a [D] query vector."""
        codes = self._hash(query.float().unsqueeze(0)).squeeze(0)  # [num_tables]
        out: set[int] = set()
        codes_list = codes.tolist()
        for t, code in enumerate(codes_list):
            bucket = self._tables[t].get(code)
            if bucket is not None:
                out.update(bucket)
        return out

    @property
    def n_items(self) -> int:
        return self._n_items
