"""Task name embeddings for the GP success predictor.

Embeds task description strings using a lightweight sentence transformer
(all-MiniLM-L6-v2, 80MB), then reduces to N_COMPONENTS dimensions via PCA
fitted on the ground truth task names.

The PCA is fitted lazily on first use and cached for the session.
"""

from __future__ import annotations

import logging

import numpy as np

logger = logging.getLogger(__name__)

N_COMPONENTS: int = 5

_model = None
_pca = None
_pca_fitted: bool = False


def _get_model():
    """Lazy-load the sentence transformer model."""
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer("all-MiniLM-L6-v2")
        logger.info("Loaded sentence-transformers/all-MiniLM-L6-v2")
    return _model


def embed_raw(texts: list[str]) -> np.ndarray:
    """Embed a list of task name strings into 384-dim vectors."""
    model = _get_model()
    return model.encode(texts, show_progress_bar=False, convert_to_numpy=True)


def fit_pca(task_names: list[str]) -> None:
    """Fit PCA on a set of task names (call once with all GT task names)."""
    global _pca, _pca_fitted
    from sklearn.decomposition import PCA

    raw = embed_raw(task_names)
    n = min(N_COMPONENTS, raw.shape[0], raw.shape[1])
    _pca = PCA(n_components=n)
    _pca.fit(raw)
    _pca_fitted = True
    logger.info(
        "PCA fitted on %d task names -> %d components (%.1f%% variance)",
        len(task_names), n, 100 * _pca.explained_variance_ratio_.sum(),
    )


def embed_task(task_name: str) -> np.ndarray:
    """Embed a single task name into N_COMPONENTS dimensions.

    Returns zeros if PCA hasn't been fitted yet.
    """
    if not _pca_fitted or _pca is None:
        return np.zeros(N_COMPONENTS)
    raw = embed_raw([task_name])
    return _pca.transform(raw)[0]


def embed_tasks_batch(task_names: list[str]) -> np.ndarray:
    """Embed a batch of task names into (len, N_COMPONENTS) array.

    Returns zeros if PCA hasn't been fitted yet.
    """
    if not _pca_fitted or _pca is None:
        return np.zeros((len(task_names), N_COMPONENTS))
    raw = embed_raw(task_names)
    return _pca.transform(raw)
