"""
GovBridge India — INDRA Poincaré Projection Engine (Sprint 29)

PROJECT INDRA Phase 1: Non-Euclidean Civic Topologies

This engine implements the HyEm (Hyperbolic Embedding) paradigm:
1. Receives oversampled Euclidean candidates from HNSW (400 vectors)
2. Projects them into the Poincaré ball model via conformal exponential map
3. Computes hyperbolic distances (arcosh-based) between query and candidates
4. Returns candidates sorted by hyperbolic distance (ascending)

ARCHITECTURAL CONSTRAINTS:
- ALL NumPy buffers are pre-allocated at __init__ time.
- ZERO heap allocations inside compute methods (GC pressure = 0).
- IEEE 754 boundary enforcement: all vectors clamped to R_MAX = 1 - 1e-3.
- Denominator floor: DENOM_MIN = 1e-15 prevents division-by-zero.
- Negative norm² clamping: np.maximum(diff_norm2, 0.0) prevents NaN from
  floating-point variance underflow on identical vectors.

MATHEMATICAL FOUNDATION:
  exp_0(v) = tanh(||v||) * (v / ||v||)    [Conformal exponential map at origin]
  d_H(u,v) = arcosh(1 + 2 * ||u-v||² / ((1-||u||²)(1-||v||²)))

USAGE:
  from indra_engine import IndraProjectionEngine
  engine = IndraProjectionEngine(batch_size=400, dimensions=768)
  distances = engine.project_and_rank(query_vec, candidate_matrix)
"""

import numpy as np
from numpy.typing import NDArray
from typing import Optional
import os

# ── Optional Learned Adapter Weights (Sprint 31) ────────────
# If W.npy and b.npy exist, use learned projection instead of naïve expmap0.
# Falls back to original tanh(||v||)/||v|| if weights not found.
_WEIGHTS_DIR = os.path.join(os.path.dirname(__file__), "..", "training", "weights")
_W_PATH = os.path.join(_WEIGHTS_DIR, "W.npy")
_B_PATH = os.path.join(_WEIGHTS_DIR, "b.npy")


class IndraProjectionEngine:
    """
    Zero-allocation Poincaré ball projection engine.

    Pre-allocates all working buffers at construction time.
    No heap allocations occur during compute_poincare_distances()
    or map_to_poincare_ball_inplace().
    """

    # ── Numerical Constants ──────────────────────────────────
    # Maximum radius in Poincaré ball. Vectors beyond this are
    # clipped to prevent arcosh(1 + 2 * inf / 0) = NaN.
    R_MAX: float = 1.0 - 1e-3   # 0.999

    # Minimum denominator value. Prevents div-by-zero when a
    # vector lies exactly on the ball boundary (||v|| ≈ 1.0).
    DENOM_MIN: float = 1e-15

    # Minimum norm for normalization. Vectors with ||v|| < this
    # are treated as zero vectors and mapped to the origin.
    NORM_MIN: float = 1e-8

    def __init__(self, batch_size: int = 400, dimensions: int = 768):
        """
        Pre-allocate all working memory.

        Args:
            batch_size: Maximum number of candidate vectors per query.
            dimensions: Embedding dimensionality (768 for Nomic).
        """
        self.batch_size = batch_size
        self.dimensions = dimensions

        # ── Pre-allocated buffers ────────────────────────────
        # Candidate matrix projected into Poincaré ball
        self._poincare_candidates = np.zeros(
            (batch_size, dimensions), dtype=np.float32
        )

        # Query vector projected into Poincaré ball
        self._poincare_query = np.zeros(dimensions, dtype=np.float32)

        # Scratch buffers for distance computation
        self._norms_sq = np.zeros(batch_size, dtype=np.float64)
        self._query_norm_sq = np.float64(0.0)
        self._diff = np.zeros(
            (batch_size, dimensions), dtype=np.float32
        )
        self._diff_norm_sq = np.zeros(batch_size, dtype=np.float64)
        self._distances = np.zeros(batch_size, dtype=np.float64)
        self._denom = np.zeros(batch_size, dtype=np.float64)

        # Buffer for norms during projection
        self._proj_norms = np.zeros(batch_size, dtype=np.float32)

    # ── Learned Adapter (Sprint 31) ───────────────────────────
    _adapter_W: Optional[NDArray[np.float32]] = None
    _adapter_b: Optional[NDArray[np.float32]] = None
    _adapter_loaded: bool = False

    def _load_adapter_weights(self) -> bool:
        """
        Attempt to load learned W.npy and b.npy from training/weights/.

        Returns True if weights loaded successfully, False otherwise.
        Falls back to naïve projection on failure.
        """
        if self._adapter_loaded:
            return self._adapter_W is not None

        self._adapter_loaded = True

        try:
            if os.path.exists(_W_PATH) and os.path.exists(_B_PATH):
                W = np.load(_W_PATH)
                b = np.load(_B_PATH)

                # Validate shapes
                if W.shape != (768, 128) or b.shape != (128,) or self.dimensions != 768:
                    return False

                self._adapter_W = W.astype(np.float32)
                self._adapter_b = b.astype(np.float32)

                # Reallocate Poincaré buffers for 128-dim
                self._poincare_candidates = np.zeros(
                    (self.batch_size, 128), dtype=np.float32
                )
                self._poincare_query = np.zeros(128, dtype=np.float32)
                self._diff = np.zeros(
                    (self.batch_size, 128), dtype=np.float32
                )

                return True
        except Exception:
            pass

        return False

    def map_to_poincare_ball_inplace(
        self,
        vectors: NDArray[np.float32],
        out: NDArray[np.float32],
        norms_buf: Optional[NDArray[np.float32]] = None,
    ) -> None:
        """
        Conformal exponential map at the origin: exp_0(v) = tanh(||v||) * (v / ||v||)

        Projects Euclidean vectors into the Poincaré ball IN-PLACE
        (writes to the `out` buffer). No allocations.

        Args:
            vectors: Input Euclidean vectors, shape (N, D) or (D,).
            out: Output buffer, same shape as vectors.
            norms_buf: Pre-allocated buffer for norms. If None, uses self._proj_norms.
        """
        is_1d = vectors.ndim == 1

        # ── Sprint 31: Use learned adapter if available ───────
        if self._load_adapter_weights() and self._adapter_W is not None:
            if is_1d:
                # Single vector: y = x @ W + b, then expmap0
                projected = vectors @ self._adapter_W + self._adapter_b  # (128,)
                norm = np.linalg.norm(projected).astype(np.float32)
                if norm < self.NORM_MIN:
                    out[:128] = 0.0
                else:
                    scale = np.float32(np.tanh(norm) / norm)
                    result = projected * scale
                    out_norm = np.linalg.norm(result).astype(np.float32)
                    if out_norm > self.R_MAX:
                        result *= np.float32(self.R_MAX / out_norm)
                    out[:128] = result
            else:
                n = vectors.shape[0]
                projected = vectors[:n] @ self._adapter_W + self._adapter_b  # (n, 128)
                for i in range(n):
                    norm_val = np.linalg.norm(projected[i]).astype(np.float32)
                    if norm_val < self.NORM_MIN:
                        out[i, :128] = 0.0
                    else:
                        scale = np.float32(np.tanh(norm_val) / norm_val)
                        result = projected[i] * scale
                        out_norm = np.linalg.norm(result).astype(np.float32)
                        if out_norm > self.R_MAX:
                            result *= np.float32(self.R_MAX / out_norm)
                        out[i, :128] = result
            return

        if is_1d:
            # Single vector case (query)
            norm = np.linalg.norm(vectors).astype(np.float32)
            if norm < self.NORM_MIN:
                out[:] = 0.0
                return
            scale = np.float32(np.tanh(norm) / norm)
            np.multiply(vectors, scale, out=out)
            # Radial boundary clip
            out_norm = np.linalg.norm(out).astype(np.float32)
            if out_norm > self.R_MAX:
                np.multiply(out, np.float32(self.R_MAX / out_norm), out=out)
        else:
            # Batch case (candidates)
            n = vectors.shape[0]
            if norms_buf is None:
                norms_buf = self._proj_norms

            # Compute L2 norms: ||v_i||
            np.einsum('ij,ij->i', vectors[:n], vectors[:n], out=norms_buf[:n])
            np.sqrt(norms_buf[:n], out=norms_buf[:n])

            for i in range(n):
                norm_val = norms_buf[i]
                if norm_val < self.NORM_MIN:
                    out[i, :] = 0.0
                else:
                    scale = np.float32(np.tanh(norm_val) / norm_val)
                    np.multiply(vectors[i], scale, out=out[i])

                    # Radial boundary clip per vector
                    out_norm = np.linalg.norm(out[i]).astype(np.float32)
                    if out_norm > self.R_MAX:
                        np.multiply(
                            out[i],
                            np.float32(self.R_MAX / out_norm),
                            out=out[i]
                        )

    def compute_poincare_distances(
        self,
        query_vec: NDArray[np.float32],
        candidate_matrix: NDArray[np.float32],
        n_candidates: int,
    ) -> NDArray[np.float64]:
        """
        Compute hyperbolic distances between query and candidates in
        the Poincaré ball model.

        Formula:
          d_H(u, v) = arcosh(1 + 2 * ||u - v||² / ((1 - ||u||²) * (1 - ||v||²)))

        All computations use pre-allocated buffers. Zero heap allocations.

        Args:
            query_vec: Euclidean query vector, shape (D,).
            candidate_matrix: Euclidean candidate vectors, shape (N, D).
            n_candidates: Actual number of candidates (may be < batch_size).

        Returns:
            Hyperbolic distances array, shape (n_candidates,).
        """
        n = min(n_candidates, self.batch_size)

        # ── Step 1: Project query into Poincaré ball ─────────
        self.map_to_poincare_ball_inplace(
            query_vec, self._poincare_query
        )

        # ── Step 2: Project candidates into Poincaré ball ────
        self.map_to_poincare_ball_inplace(
            candidate_matrix[:n],
            self._poincare_candidates[:n],
            norms_buf=self._proj_norms
        )

        # ── Step 3: Compute ||u||² (query norm squared) ──────
        self._query_norm_sq = np.float64(
            np.dot(self._poincare_query, self._poincare_query)
        )

        # ── Step 4: Compute ||v_i||² (candidate norms squared) ─
        np.einsum(
            'ij,ij->i',
            self._poincare_candidates[:n].astype(np.float64),
            self._poincare_candidates[:n].astype(np.float64),
            out=self._norms_sq[:n]
        )

        # ── Step 5: Compute ||u - v_i||² ─────────────────────
        np.subtract(
            self._poincare_candidates[:n],
            self._poincare_query,
            out=self._diff[:n]
        )
        np.einsum(
            'ij,ij->i',
            self._diff[:n].astype(np.float64),
            self._diff[:n].astype(np.float64),
            out=self._diff_norm_sq[:n]
        )

        # GUARD: Floating-point variance underflow.
        # Identical vectors can produce tiny negative values like -1e-7
        # due to IEEE 754 rounding. Clamp to zero.
        np.maximum(self._diff_norm_sq[:n], 0.0, out=self._diff_norm_sq[:n])

        # ── Step 6: Compute denominator ──────────────────────
        # denom = (1 - ||u||²) * (1 - ||v_i||²)
        np.subtract(1.0, self._norms_sq[:n], out=self._denom[:n])
        np.multiply(
            self._denom[:n],
            (1.0 - self._query_norm_sq),
            out=self._denom[:n]
        )

        # GUARD: Denominator floor. If a vector lies exactly on the
        # boundary (||v|| ≈ 1.0), denom → 0, causing infinity.
        np.maximum(self._denom[:n], self.DENOM_MIN, out=self._denom[:n])

        # ── Step 7: Compute arcosh argument ──────────────────
        # arg = 1 + 2 * ||u - v||² / denom
        np.divide(self._diff_norm_sq[:n], self._denom[:n], out=self._distances[:n])
        np.multiply(self._distances[:n], 2.0, out=self._distances[:n])
        np.add(self._distances[:n], 1.0, out=self._distances[:n])

        # GUARD: arcosh domain enforcement. arcosh(x) requires x >= 1.
        # Clamp to exactly 1.0 (distance = 0) if numerical error
        # pushes below 1.0.
        np.maximum(self._distances[:n], 1.0, out=self._distances[:n])

        # ── Step 8: Final hyperbolic distance ────────────────
        np.arccosh(self._distances[:n], out=self._distances[:n])

        return self._distances[:n].copy()  # Return a copy (safe for caller)

    def project_and_rank(
        self,
        query_vec: NDArray[np.float32],
        candidate_matrix: NDArray[np.float32],
        n_candidates: int,
    ) -> NDArray[np.intp]:
        """
        Full pipeline: project + compute distances + return sorted indices.

        Args:
            query_vec: Euclidean query embedding, shape (D,).
            candidate_matrix: Euclidean candidate embeddings, shape (N, D).
            n_candidates: Actual number of candidates.

        Returns:
            Array of indices sorted by ascending hyperbolic distance.
        """
        distances = self.compute_poincare_distances(
            query_vec, candidate_matrix, n_candidates
        )
        return np.argsort(distances)