"""
RBSP Vector Quantization -- Binary and Product Quantization.

Binary quantization: compress float vectors to bit arrays.  Distance = Hamming
(XOR + popcount).

Product quantization: split vector into subspaces, quantize each independently.

Key insight: Python arbitrary-precision ``int`` acts as a SIMD register.
``XOR + bin().count('1')`` gives Hamming distance in two operations regardless
of dimension.

Zero external dependencies.
"""

from __future__ import annotations

import heapq
import math
import random
from typing import Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def popcount(n: int) -> int:
    """Count set bits.  Uses ``bin(n).count('1')`` which CPython optimises."""
    return bin(n).count("1")


def hamming_distance(a: int, b: int) -> int:
    """Hamming distance via XOR + popcount.  O(1) regardless of dimension."""
    return popcount(a ^ b)


# ---------------------------------------------------------------------------
# Internal: lightweight k-means for ProductQuantizer.fit()
# ---------------------------------------------------------------------------

def _euclidean_sq(a: List[float], b: List[float]) -> float:
    """Squared Euclidean distance between two same-length vectors."""
    total = 0.0
    for x, y in zip(a, b):
        d = x - y
        total += d * d
    return total


def _vec_add(dst: List[float], src: List[float]) -> None:
    """In-place element-wise addition: ``dst += src``."""
    for i in range(len(dst)):
        dst[i] += src[i]


def _kmeans(
    vectors: List[List[float]],
    k: int,
    max_iter: int = 20,
    seed: int = 42,
) -> List[List[float]]:
    """Simple k-means clustering.

    Returns *k* centroids as ``List[List[float]]``.
    Uses ``random.sample`` for initialisation.

    Args:
        vectors:  Training vectors (must all share the same dimensionality).
        k:        Number of clusters.
        max_iter: Maximum iterations (default 20).
        seed:     RNG seed for reproducibility.
    """
    if not vectors:
        return []
    dim = len(vectors[0])
    n = len(vectors)

    rng = random.Random(seed)

    # Clamp k to the number of distinct vectors available
    k = min(k, n)

    # Initialise centroids via random sample
    indices = rng.sample(range(n), k)
    centroids: List[List[float]] = [list(vectors[i]) for i in indices]

    assignments: List[int] = [0] * n

    for _ in range(max_iter):
        changed = False

        # Assign each vector to the nearest centroid
        for vi in range(n):
            best_c = 0
            best_dist = _euclidean_sq(vectors[vi], centroids[0])
            for ci in range(1, k):
                d = _euclidean_sq(vectors[vi], centroids[ci])
                if d < best_dist:
                    best_dist = d
                    best_c = ci
            if assignments[vi] != best_c:
                assignments[vi] = best_c
                changed = True

        if not changed:
            break

        # Recompute centroids
        new_centroids: List[List[float]] = [[0.0] * dim for _ in range(k)]
        counts: List[int] = [0] * k
        for vi in range(n):
            c = assignments[vi]
            _vec_add(new_centroids[c], vectors[vi])
            counts[c] += 1

        for ci in range(k):
            if counts[ci] > 0:
                inv = 1.0 / counts[ci]
                for d in range(dim):
                    new_centroids[ci][d] *= inv
            else:
                # Keep previous centroid for empty clusters
                new_centroids[ci] = list(centroids[ci])

        centroids = new_centroids

    return centroids


# ---------------------------------------------------------------------------
# Binary Quantizer
# ---------------------------------------------------------------------------

class BinaryQuantizer:
    """Quantise float vectors to binary codes.

    Each dimension is quantised to 1 bit: ``value >= threshold -> 1, else 0``.
    Stored as a Python ``int`` (arbitrary precision = free SIMD).

    A ``fit()`` call computes per-dimension thresholds (median of training
    data).  After fitting, ``quantize()`` converts any same-dimensional vector
    to a compact integer code, and ``hamming_distance()`` compares two codes in
    O(1).

    Args:
        dim: Vector dimensionality.  If not provided it is inferred from the
             first call to ``fit()`` or ``insert()``.
    """

    def __init__(self, dim: Optional[int] = None) -> None:
        self._dim = dim
        self._thresholds: Optional[List[float]] = None
        # Indexed store: id -> binary code
        self._codes: Dict[str, int] = {}

    # -- Properties -----------------------------------------------------------

    @property
    def dim(self) -> Optional[int]:
        """Dimensionality (``None`` until first ``fit`` or ``insert``)."""
        return self._dim

    @property
    def thresholds(self) -> Optional[List[float]]:
        """Per-dimension quantisation thresholds (``None`` before ``fit``)."""
        return self._thresholds

    @property
    def size(self) -> int:
        """Number of indexed vectors."""
        return len(self._codes)

    # -- Core API -------------------------------------------------------------

    def fit(self, vectors: List[List[float]]) -> None:
        """Compute per-dimension thresholds (median) from training vectors.

        Args:
            vectors: Training data.  All vectors must share the same length.

        Raises:
            ValueError: If *vectors* is empty.
        """
        if not vectors:
            raise ValueError("Cannot fit on empty vector set")

        dim = len(vectors[0])
        self._dim = dim

        # Transpose to get per-dimension lists, then compute median
        thresholds: List[float] = []
        for d in range(dim):
            col = sorted(v[d] for v in vectors)
            mid = len(col) // 2
            if len(col) % 2 == 1:
                thresholds.append(col[mid])
            else:
                thresholds.append((col[mid - 1] + col[mid]) / 2.0)

        self._thresholds = thresholds

    def quantize(self, vector: List[float]) -> int:
        """Convert a float vector to a binary code (Python ``int``).

        Uses per-dimension thresholds from ``fit()`` if available, otherwise
        falls back to zero-threshold (positive -> 1).

        Args:
            vector: Input float vector.

        Returns:
            An integer whose *i*-th bit is 1 iff ``vector[i] >= threshold[i]``.
        """
        if self._dim is None:
            self._dim = len(vector)
        code = 0
        thresholds = self._thresholds
        for i, v in enumerate(vector):
            if thresholds is not None:
                if v >= thresholds[i]:
                    code |= 1 << i
            else:
                if v >= 0.0:
                    code |= 1 << i
        return code

    @staticmethod
    def hamming_distance(a: int, b: int) -> int:
        """Hamming distance between two binary codes.

        Equivalent to the module-level ``hamming_distance`` but available as a
        method for convenience.
        """
        return popcount(a ^ b)

    def insert(self, id: str, vector: List[float]) -> None:
        """Quantise and store a vector under *id*.

        Args:
            id:     Unique string identifier.
            vector: Float vector to index.
        """
        self._codes[id] = self.quantize(vector)

    def batch_insert(self, ids: List[str], vectors: List[List[float]]) -> None:
        """Bulk-insert multiple vectors.

        Args:
            ids:     List of unique string identifiers (same length as *vectors*).
            vectors: List of float vectors.

        Raises:
            ValueError: If *ids* and *vectors* differ in length.
        """
        if len(ids) != len(vectors):
            raise ValueError(
                f"ids ({len(ids)}) and vectors ({len(vectors)}) must have "
                f"the same length"
            )
        for id_, vec in zip(ids, vectors):
            self._codes[id_] = self.quantize(vec)

    def search(
        self,
        query: List[float],
        k: int = 10,
    ) -> List[Tuple[str, int]]:
        """Find the *k* nearest indexed vectors by Hamming distance.

        Args:
            query: Query float vector (will be quantised internally).
            k:     Number of results to return.

        Returns:
            List of ``(id, hamming_distance)`` tuples, ordered by ascending
            distance.
        """
        q_code = self.quantize(query)
        # heapq.nsmallest is O(n log k) which beats full sort for small k
        scored = [(id_, popcount(code ^ q_code)) for id_, code in self._codes.items()]
        return heapq.nsmallest(k, scored, key=lambda x: x[1])


# ---------------------------------------------------------------------------
# Product Quantizer
# ---------------------------------------------------------------------------

class ProductQuantizer:
    """Product quantisation: split vector into *M* subspaces, cluster each to
    *n_clusters* centroids.

    Reduces memory from ``D * 4`` bytes to ``M`` bytes per vector (when
    *n_clusters* = 256 so each code is one byte).  Distance is computed via
    lookup tables (ADC -- Asymmetric Distance Computation).

    Args:
        m:    Number of subspaces.  The vector dimensionality must be
              divisible by *m*.
        seed: RNG seed forwarded to k-means initialisation.
    """

    def __init__(self, m: int = 8, seed: int = 42) -> None:
        self._m = m
        self._seed = seed
        self._dim: Optional[int] = None
        self._dsub: Optional[int] = None
        self._n_clusters: int = 256

        # codebooks[j] is a list of n_clusters centroids, each of length dsub
        self._codebooks: Optional[List[List[List[float]]]] = None

        # Indexed store: id -> list of M code bytes
        self._codes: Dict[str, List[int]] = {}

    # -- Properties -----------------------------------------------------------

    @property
    def m(self) -> int:
        """Number of subspaces."""
        return self._m

    @property
    def dim(self) -> Optional[int]:
        """Full vector dimensionality (``None`` before ``fit``)."""
        return self._dim

    @property
    def dsub(self) -> Optional[int]:
        """Per-subspace dimensionality (``None`` before ``fit``)."""
        return self._dsub

    @property
    def n_clusters(self) -> int:
        """Number of centroids per subspace."""
        return self._n_clusters

    @property
    def size(self) -> int:
        """Number of indexed vectors."""
        return len(self._codes)

    # -- Core API -------------------------------------------------------------

    def fit(
        self,
        vectors: List[List[float]],
        n_clusters: int = 256,
    ) -> None:
        """Train codebooks by running k-means on each subspace.

        Args:
            vectors:    Training data.
            n_clusters: Centroids per subspace (default 256, i.e. 1-byte codes).

        Raises:
            ValueError: If *vectors* is empty or dimensionality is not
                        divisible by *m*.
        """
        if not vectors:
            raise ValueError("Cannot fit on empty vector set")

        dim = len(vectors[0])
        if dim % self._m != 0:
            raise ValueError(
                f"Dimension {dim} is not divisible by m={self._m}"
            )

        self._dim = dim
        self._dsub = dim // self._m
        self._n_clusters = n_clusters

        dsub = self._dsub
        codebooks: List[List[List[float]]] = []

        for j in range(self._m):
            # Extract subvectors for subspace j
            start = j * dsub
            sub_vectors = [v[start : start + dsub] for v in vectors]
            centroids = _kmeans(
                sub_vectors, n_clusters, max_iter=20, seed=self._seed + j
            )
            codebooks.append(centroids)

        self._codebooks = codebooks

    def encode(self, vector: List[float]) -> List[int]:
        """Encode a vector into *M* code bytes.

        Each byte is the index of the nearest centroid in the corresponding
        subspace's codebook.

        Args:
            vector: Full-dimensional float vector.

        Returns:
            List of *M* integers in ``[0, n_clusters)``.

        Raises:
            RuntimeError: If the quantizer has not been fitted.
        """
        if self._codebooks is None:
            raise RuntimeError("ProductQuantizer has not been fitted")

        dsub = self._dsub
        codes: List[int] = []

        for j in range(self._m):
            start = j * dsub
            sub = vector[start : start + dsub]
            book = self._codebooks[j]

            best_idx = 0
            best_dist = _euclidean_sq(sub, book[0])
            for ci in range(1, len(book)):
                d = _euclidean_sq(sub, book[ci])
                if d < best_dist:
                    best_dist = d
                    best_idx = ci
            codes.append(best_idx)

        return codes

    def decode(self, codes: List[int]) -> List[float]:
        """Approximate reconstruction: look up centroids and concatenate.

        Args:
            codes: *M* code bytes as returned by ``encode()``.

        Returns:
            Reconstructed float vector of length ``dim``.

        Raises:
            RuntimeError: If the quantizer has not been fitted.
        """
        if self._codebooks is None:
            raise RuntimeError("ProductQuantizer has not been fitted")

        result: List[float] = []
        for j in range(self._m):
            centroid = self._codebooks[j][codes[j]]
            result.extend(centroid)
        return result

    def distance_table(self, query: List[float]) -> List[List[float]]:
        """Precompute a distance lookup table for *query* (ADC).

        Returns a table ``T`` such that ``T[j][c]`` is the squared Euclidean
        distance between the *j*-th subvector of *query* and centroid *c* in
        subspace *j*.  This allows computing approximate distances to any
        encoded vector in O(M) time.

        Args:
            query: Full-dimensional query vector.

        Returns:
            A list of *M* lists, each of length *n_clusters*.

        Raises:
            RuntimeError: If the quantizer has not been fitted.
        """
        if self._codebooks is None:
            raise RuntimeError("ProductQuantizer has not been fitted")

        dsub = self._dsub
        table: List[List[float]] = []

        for j in range(self._m):
            start = j * dsub
            q_sub = query[start : start + dsub]
            book = self._codebooks[j]
            dists = [_euclidean_sq(q_sub, book[c]) for c in range(len(book))]
            table.append(dists)

        return table

    def _adc_distance(
        self,
        table: List[List[float]],
        codes: List[int],
    ) -> float:
        """Compute ADC distance using a precomputed distance table."""
        total = 0.0
        for j in range(self._m):
            total += table[j][codes[j]]
        return total

    def insert(self, id: str, vector: List[float]) -> None:
        """Encode and store a vector under *id*.

        Args:
            id:     Unique string identifier.
            vector: Float vector to index.
        """
        self._codes[id] = self.encode(vector)

    def search(
        self,
        query: List[float],
        k: int = 10,
    ) -> List[Tuple[str, float]]:
        """Find the *k* nearest indexed vectors by ADC distance.

        Args:
            query: Query float vector.
            k:     Number of results to return.

        Returns:
            List of ``(id, squared_distance)`` tuples, ordered by ascending
            distance.
        """
        table = self.distance_table(query)
        scored = [
            (id_, self._adc_distance(table, codes))
            for id_, codes in self._codes.items()
        ]
        return heapq.nsmallest(k, scored, key=lambda x: x[1])
