"""
FastKNN — FAISS-backed K-Nearest Neighbours for ML data sourcing.

The class is purpose-built for machine learning workflows: the ``label``
column acts as an anchor that identifies which rows of ``query_df`` are
*positive examples* (``label == 1``).  Those vectors are used to retrieve
the most similar candidates from the ``search_df`` pool — a pattern commonly
used for hard-negative mining, data augmentation, and data sourcing.

Supported FAISS index types
---------------------------
* ``flat``     – Exact brute-force (default).  Perfect recall.
* ``ivf_flat`` – Approximate; partitions space into clusters.  Fast at scale.
* ``ivf_pq``   – Approximate + Product Quantisation.  Lowest memory footprint.
* ``hnsw``     – Graph-based approximate search.  Fast queries, high recall.

All FAISS parameters are optional and default to sensible values, so
existing call-sites that only pass the original five positional arguments
continue to work without any changes.
"""

import logging
import time
from typing import List, Optional, Tuple

import faiss
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_VALID_INDEX_TYPES: Tuple[str, ...] = ("flat", "ivf_flat", "ivf_pq", "hnsw")
_VALID_METRICS: Tuple[str, ...] = ("l2", "ip")
_VALID_FILLNA: Tuple[str, ...] = ("zero", "mean", "median", "mode", "max", "min")


# ---------------------------------------------------------------------------
# FastKNN
# ---------------------------------------------------------------------------

class FastKNN:
    """
    Fast K-Nearest Neighbours similarity search using FAISS.

    Designed for ML data-sourcing workflows where a binary ``label`` column
    anchors the query space: only ``query_df`` rows where ``label == 1`` are
    used as query vectors.  The index is built over ``search_df``, and the
    matched rows are returned for downstream use (e.g. hard-negative mining,
    data augmentation).

    Args:
        query_df: Query DataFrame containing ``label`` and all ``knn_keys``
            columns.
        search_df: Search (database) DataFrame to index and retrieve from.
        knn_keys: Feature column names used to build the similarity vectors.
        label: Binary label column name.  Rows where ``query_df[label] == 1``
            are treated as positive anchors.

        fillna_method: Missing-value imputation strategy.  One of
            ``'zero'`` (default), ``'mean'``, ``'median'``, ``'mode'``,
            ``'max'``, ``'min'``.
        k_neighbors: Maximum neighbours retrieved per query vector.  Capped
            automatically at ``len(search_df)``.  Defaults to ``1000``.
        remove_duplicate: Drop duplicate rows from both DataFrames before
            processing.  Defaults to ``False``.

        index_type: FAISS index family.  Defaults to ``'flat'`` (exact).
            Options: ``'flat'``, ``'ivf_flat'``, ``'ivf_pq'``, ``'hnsw'``.
        metric: Distance metric — ``'l2'`` (squared Euclidean, default) or
            ``'ip'`` (inner product).  Use ``metric='ip'`` together with
            ``normalize=True`` to get **cosine similarity**.
        normalize: If ``True``, L2-normalise all vectors before indexing and
            searching.  Combined with ``metric='ip'`` this yields cosine
            similarity.  Defaults to ``False``.

        nlist: Number of Voronoi clusters for IVF indexes (``ivf_flat``,
            ``ivf_pq``).  ``None`` auto-selects ``max(1, int(sqrt(n)))``.
            Ignored for ``'flat'`` and ``'hnsw'``.
        nprobe: Clusters visited per query for IVF indexes.  Higher → better
            recall, higher latency.  Capped at ``nlist`` automatically.
            Defaults to ``8``.

        pq_m: PQ sub-vector count (``ivf_pq`` only).  Must divide the feature
            dimension; auto-adjusted if not.  Defaults to ``8``.
        pq_nbits: Bits per PQ sub-vector codebook entry (``ivf_pq`` only).
            Typical value ``8`` (256 centroids/sub-vector).  Defaults to ``8``.

        hnsw_m: Bi-directional links per graph node (``hnsw`` only).
            Higher → better recall, more memory.  Range 16–64.  Defaults to
            ``32``.
        ef_construction: Build-time beam width (``hnsw`` only).  Higher →
            better graph, slower build.  Defaults to ``64``.
        ef_search: Query-time beam width (``hnsw`` only).  Higher → better
            recall, slower queries.  Auto-raised to ``k_neighbors`` if
            smaller.  Defaults to ``32``.

    Example::

        knn = FastKNN(
            query_df=labelled_df,
            search_df=pool_df,
            knn_keys=["emb_0", "emb_1", "emb_2"],
            label="is_positive",
        )
        sourced_df, sourced_df_with_label = knn.run()
        knn.validate()

        # Approximate search with cosine similarity
        knn_cos = FastKNN(
            query_df=labelled_df,
            search_df=pool_df,
            knn_keys=["emb_0", "emb_1", "emb_2"],
            label="is_positive",
            index_type="ivf_flat",
            metric="ip",
            normalize=True,
            nlist=200,
            nprobe=16,
        )
    """

    def __init__(
        self,
        query_df: pd.DataFrame,
        search_df: pd.DataFrame,
        knn_keys: List[str],
        label: str,
        # ── existing parameters (backward-compatible) ─────────────────
        fillna_method: str = "zero",
        k_neighbors: int = 1000,
        remove_duplicate: bool = False,
        # ── FAISS index selection ─────────────────────────────────────
        index_type: str = "flat",
        metric: str = "l2",
        normalize: bool = False,
        # ── IVF parameters (ivf_flat / ivf_pq) ───────────────────────
        nlist: Optional[int] = None,
        nprobe: int = 8,
        # ── IVF-PQ parameters (ivf_pq only) ──────────────────────────
        pq_m: int = 8,
        pq_nbits: int = 8,
        # ── HNSW parameters (hnsw only) ───────────────────────────────
        hnsw_m: int = 32,
        ef_construction: int = 64,
        ef_search: int = 32,
    ) -> None:
        # -- validate enum args ----------------------------------------
        if index_type not in _VALID_INDEX_TYPES:
            raise ValueError(
                f"index_type must be one of {_VALID_INDEX_TYPES}, got {index_type!r}"
            )
        if metric not in _VALID_METRICS:
            raise ValueError(
                f"metric must be one of {_VALID_METRICS}, got {metric!r}"
            )
        if fillna_method not in _VALID_FILLNA:
            raise ValueError(
                f"fillna_method must be one of {_VALID_FILLNA}, got {fillna_method!r}"
            )
        if k_neighbors < 1:
            raise ValueError(f"k_neighbors must be ≥ 1, got {k_neighbors}")
        if nprobe < 1:
            raise ValueError(f"nprobe must be ≥ 1, got {nprobe}")

        # -- store parameters ------------------------------------------
        self.query_df = query_df
        self.search_df = search_df
        self.knn_keys = knn_keys
        self.label = label

        self.fillna_method = fillna_method
        self.k_neighbors = k_neighbors
        self.remove_duplicate = remove_duplicate

        self.index_type = index_type
        self.metric = metric
        self.normalize = normalize

        self.nlist = nlist
        self.nprobe = nprobe
        self.pq_m = pq_m
        self.pq_nbits = pq_nbits
        self.hnsw_m = hnsw_m
        self.ef_construction = ef_construction
        self.ef_search = ef_search

        # -- keep raw copies for validation / retrieval ----------------
        self.query_df_raw = query_df.copy()
        self.search_df_raw = search_df.copy()

        # -- result state ----------------------------------------------
        self.indices: Optional[List[int]] = None
        self.sourced_search_df: pd.DataFrame = pd.DataFrame()
        self.sourced_search_df_with_label: Optional[pd.DataFrame] = None

        # -- pre-process immediately upon initialisation ---------------
        t = time.time()
        self._pre_process()
        logger.info("Preprocessing took %.4f s", time.time() - t)

    # ------------------------------------------------------------------
    # Pre-processing
    # ------------------------------------------------------------------

    def _pre_process(self) -> None:
        """Filter, column-select, impute, and optionally deduplicate."""
        logger.info("Pre-processing started")

        # ML anchor: keep only positive-label rows as query vectors
        if self.label in self.query_df.columns:
            self.query_df = self.query_df[self.query_df[self.label] == 1]

        # Restrict to feature columns
        self.query_df = self.query_df[self.knn_keys]
        self.search_df = self.search_df[self.knn_keys]
        logger.info("Feature columns selected: %s", self.knn_keys)

        # Impute missing values
        self.query_df = self._fillna(self.query_df)
        self.search_df = self._fillna(self.search_df)
        logger.info("NaN imputation: method='%s'", self.fillna_method)

        if self.remove_duplicate:
            self.query_df = self.query_df.drop_duplicates()
            self.search_df = self.search_df.drop_duplicates()
            logger.info("Duplicates removed")

        logger.info(
            "Pre-processing complete — query: %d rows, search: %d rows",
            len(self.query_df),
            len(self.search_df),
        )

    def _fillna(self, df: pd.DataFrame) -> pd.DataFrame:
        """Impute missing values according to *fillna_method*."""
        try:
            if self.fillna_method == "zero":
                return df.fillna(0)

            strategy_map = {
                "mean":   df.mean(),
                "median": df.median(),
                "mode":   df.mode().iloc[0],
                "max":    df.max(),
                "min":    df.min(),
            }
            fill_values = strategy_map[self.fillna_method]

            if fill_values.isna().any():
                logger.warning(
                    "Some columns are entirely NaN; defaulting those to 0."
                )
                fill_values = fill_values.fillna(0)

            return df.fillna(fill_values)

        except Exception as exc:
            logger.error("_fillna error (%s); falling back to zero-fill.", exc)
            return df.fillna(0)

    # ------------------------------------------------------------------
    # Array helpers
    # ------------------------------------------------------------------

    def _to_float32(self, df: pd.DataFrame) -> np.ndarray:
        """Return a contiguous float32 array, L2-normalised if requested."""
        arr = np.ascontiguousarray(df.values.astype("float32"))
        if self.normalize:
            faiss.normalize_L2(arr)  # in-place
        return arr

    def _faiss_metric_const(self) -> int:
        return (
            faiss.METRIC_INNER_PRODUCT
            if self.metric == "ip"
            else faiss.METRIC_L2
        )

    # ------------------------------------------------------------------
    # Index construction
    # ------------------------------------------------------------------

    def _build_index(self, search_np: np.ndarray) -> faiss.Index:
        """Construct, configure, and (where needed) train the FAISS index."""
        n, d = search_np.shape
        faiss_metric = self._faiss_metric_const()

        # ── flat ───────────────────────────────────────────────────────
        if self.index_type == "flat":
            index = (
                faiss.IndexFlatIP(d)
                if self.metric == "ip"
                else faiss.IndexFlatL2(d)
            )
            logger.info("Built IndexFlat (metric=%s, d=%d)", self.metric, d)

        # ── ivf_flat ───────────────────────────────────────────────────
        elif self.index_type == "ivf_flat":
            nlist = self._resolve_nlist(n)
            quantizer = (
                faiss.IndexFlatIP(d) if self.metric == "ip"
                else faiss.IndexFlatL2(d)
            )
            index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss_metric)
            logger.info(
                "Training IVFFlat (nlist=%d, metric=%s) on %d vectors…",
                nlist, self.metric, n,
            )
            index.train(search_np)
            index.nprobe = min(self.nprobe, nlist)
            logger.info("IVFFlat trained — nprobe=%d", index.nprobe)

        # ── ivf_pq ─────────────────────────────────────────────────────
        elif self.index_type == "ivf_pq":
            nlist = self._resolve_nlist(n)
            pq_m = self._resolve_pq_m(d)
            # IVF-PQ always uses an L2 coarse quantizer internally
            quantizer = faiss.IndexFlatL2(d)
            index = faiss.IndexIVFPQ(quantizer, d, nlist, pq_m, self.pq_nbits)
            logger.info(
                "Training IVFPQ (nlist=%d, pq_m=%d, pq_nbits=%d) on %d vectors…",
                nlist, pq_m, self.pq_nbits, n,
            )
            index.train(search_np)
            index.nprobe = min(self.nprobe, nlist)
            logger.info("IVFPQ trained — nprobe=%d", index.nprobe)

        # ── hnsw ───────────────────────────────────────────────────────
        elif self.index_type == "hnsw":
            index = faiss.IndexHNSWFlat(d, self.hnsw_m, faiss_metric)
            index.hnsw.efConstruction = self.ef_construction
            # efSearch must be >= k to avoid empty result slots
            index.hnsw.efSearch = max(self.ef_search, self.k_neighbors)
            logger.info(
                "Built HNSW (M=%d, efConstruction=%d, efSearch=%d, metric=%s)",
                self.hnsw_m,
                self.ef_construction,
                index.hnsw.efSearch,
                self.metric,
            )

        return index

    def _resolve_nlist(self, n: int) -> int:
        """Return *nlist*, auto-computing and warning about training-size requirements."""
        nlist = self.nlist if self.nlist is not None else max(1, int(np.sqrt(n)))
        # FAISS recommends at least 39 × nlist training vectors for stable results
        min_required = 39 * nlist
        if n < min_required:
            logger.warning(
                "search_df has %d vectors but nlist=%d needs ≥ %d for stable IVF "
                "training. Consider lowering nlist or switching to index_type='flat'.",
                n, nlist, min_required,
            )
        return nlist

    def _resolve_pq_m(self, d: int) -> int:
        """Return *pq_m*, falling back to the largest valid divisor of *d* if needed."""
        if d % self.pq_m == 0:
            return self.pq_m
        fallback = max(m for m in range(1, d + 1) if d % m == 0 and m <= self.pq_m)
        logger.warning(
            "pq_m=%d does not divide feature dim=%d; using pq_m=%d instead.",
            self.pq_m, d, fallback,
        )
        return fallback

    # ------------------------------------------------------------------
    # Search
    # ------------------------------------------------------------------

    def indexing(self) -> Optional[List[int]]:
        """Build the FAISS index and search for neighbours.

        Returns:
            Deduplicated, sorted list of matched row indices into
            ``search_df``, or ``None`` on failure.
        """
        t = time.time()
        try:
            query_np = self._to_float32(self.query_df)
            search_np = self._to_float32(self.search_df)

            index = self._build_index(search_np)
            index.add(search_np)

            effective_k = min(self.k_neighbors, len(self.search_df))
            _, I = index.search(query_np, effective_k)

            # I can contain –1 for unfilled slots (approximate indexes)
            self.indices = np.unique(I[I >= 0]).tolist()

            logger.info(
                "Search complete — index_type='%s', metric='%s', "
                "found %d unique neighbours in %.4f s",
                self.index_type, self.metric,
                len(self.indices), time.time() - t,
            )
            return self.indices

        except Exception as exc:
            logger.error("FastKNN indexing failed: %s", exc)
            return None

    def run(self) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
        """Run the full KNN pipeline and return matched rows from ``search_df``.

        Returns:
            A 2-tuple:

            * **sourced_search_df** — feature-columns-only subset of matched
              rows.
            * **sourced_search_df_with_label** — full original rows (all
              columns) of the matched rows from ``search_df_raw``.
        """
        if self.indexing() is None:
            return pd.DataFrame(), None

        if self.indices:
            self.sourced_search_df = self.search_df.iloc[self.indices]
            self.sourced_search_df_with_label = self.search_df_raw.iloc[self.indices]

        return self.sourced_search_df, self.sourced_search_df_with_label

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    def validate(self) -> None:
        """Report label-recall statistics after a :meth:`run` call.

        Compares positive-label (``label == 1``) counts in the original
        ``search_df`` against those retained after KNN retrieval, giving a
        concise recall summary.
        """
        if self.label not in self.search_df_raw.columns:
            logger.warning(
                "Label column '%s' not found in search_df — skipping validation.",
                self.label,
            )
            return

        if not self.indices:
            logger.warning("No results found. Call run() before validate().")
            return

        try:
            raw_label_count = int((self.search_df_raw[self.label] == 1).sum())
            sourced_label_count = 0
            if self.sourced_search_df_with_label is not None:
                sourced_label_count = int(
                    (self.sourced_search_df_with_label[self.label] == 1).sum()
                )

            recall = (
                sourced_label_count / raw_label_count
                if raw_label_count > 0
                else float("nan")
            )

            logger.info("── FastKNN Validation ──────────────────────────────")
            logger.info("  Index type : %s  (metric=%s, normalize=%s)",
                        self.index_type, self.metric, self.normalize)
            logger.info("  Rows        before / after : %d / %d",
                        len(self.search_df_raw), len(self.sourced_search_df))
            logger.info("  Labels      before / after : %d / %d  (recall=%.1f%%)",
                        raw_label_count, sourced_label_count, recall * 100)
            logger.info("────────────────────────────────────────────────────")

        except Exception as exc:
            logger.error("Validation error: %s", exc)

    # ------------------------------------------------------------------
    # Dunder helpers
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"FastKNN("
            f"index_type={self.index_type!r}, "
            f"metric={self.metric!r}, "
            f"normalize={self.normalize}, "
            f"k_neighbors={self.k_neighbors}, "
            f"n_search={len(self.search_df_raw)})"
        )
