"""Vectorization engine — converts text responses to identity vectors.

Uses sentence-transformers all-MiniLM-L6-v2 (384 dimensions).
Model is downloaded lazily on first use.
"""

from __future__ import annotations

from typing import Optional

import numpy as np


class Vectorizer:
    """Converts text responses into embedding vectors using sentence-transformers.

    The model is loaded lazily on first call to avoid slow imports and
    unnecessary downloads when not needed.
    """

    DEFAULT_MODEL = "all-MiniLM-L6-v2"
    VECTOR_DIM = 384

    def __init__(self, model_name: Optional[str] = None):
        self.model_name = model_name or self.DEFAULT_MODEL
        self._model = None

    def _load_model(self) -> None:
        """Load the sentence-transformers model (lazy, on first use)."""
        if self._model is not None:
            return
        try:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_name)
        except ImportError:
            raise RuntimeError(
                "sentence-transformers is required for vectorization. "
                "Install it with: pip install sentence-transformers"
            )

    def encode(self, text: str) -> np.ndarray:
        """Encode a single text string into a 384-dimensional vector.

        Args:
            text: The text to encode.

        Returns:
            A numpy array of shape (384,) with float32 values.
        """
        self._load_model()
        # SentenceTransformer.encode returns ndarray
        vector = self._model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
        return vector.astype(np.float32)

    def encode_batch(self, texts: list[str]) -> np.ndarray:
        """Encode multiple texts into vectors.

        Args:
            texts: List of text strings to encode.

        Returns:
            A numpy array of shape (len(texts), 384) with float32 values.
        """
        if not texts:
            return np.empty((0, self.VECTOR_DIM), dtype=np.float32)
        self._load_model()
        vectors = self._model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
        return vectors.astype(np.float32)

    def average_vectors(self, vectors: np.ndarray) -> np.ndarray:
        """Compute the average of multiple vectors, then normalize.

        This produces a single representative vector from multiple responses.

        Args:
            vectors: Array of shape (n, 384).

        Returns:
            A normalized vector of shape (384,).
        """
        if vectors.ndim == 1:
            return self._normalize(vectors)
        if len(vectors) == 0:
            return np.zeros(self.VECTOR_DIM, dtype=np.float32)
        avg = vectors.mean(axis=0)
        return self._normalize(avg)

    def weighted_average_vectors(self, vectors: np.ndarray, weights: np.ndarray) -> np.ndarray:
        """Compute weighted average of vectors, then normalize.

        # WHY weighted average instead of replacing average_vectors: The
        # unweighted average_vectors() is still used by external callers and
        # must remain unchanged. This method applies prompt-type weights so
        # scenario/tradeoff responses (behavioral, harder to fake) contribute
        # more to the identity vector than direct self-report prompts.
        # The L2 normalization after averaging ensures the result is a unit
        # vector comparable to other identity vectors via cosine similarity.

        Args:
            vectors: Array of shape (n, 384).
            weights: Array of shape (n,) with per-vector weights.

        Returns:
            A normalized vector of shape (384,).
        """
        if vectors.ndim == 1:
            return self._normalize(vectors)
        if len(vectors) == 0:
            return np.zeros(self.VECTOR_DIM, dtype=np.float32)
        # Weighted mean: sum(w_i * v_i) / sum(w_i)
        w = weights.reshape(-1, 1)  # (n, 1) for broadcasting
        weighted_sum = (vectors * w).sum(axis=0)
        total_weight = weights.sum()
        if total_weight < 1e-10:
            return self._normalize(vectors.mean(axis=0))
        avg = weighted_sum / total_weight
        return self._normalize(avg)

    @staticmethod
    def _normalize(vector: np.ndarray) -> np.ndarray:
        """L2-normalize a vector."""
        norm = np.linalg.norm(vector)
        if norm < 1e-10:
            return vector
        return (vector / norm).astype(np.float32)
