# Author: Simon-Pierre Boucher — contact@spboucher.ai
#
"""Sentence embeddings and cosine-similarity features."""

import numpy as np

from . import config
from .references import REFERENCES, SIM_COLS


def encode_references(model=None):
    """Encode the 20 reference descriptions. Returns (names, embeddings)."""
    if model is None:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(config.EMBEDDING_MODEL)
    names = list(REFERENCES.keys())
    embeddings = model.encode(list(REFERENCES.values()), normalize_embeddings=True)
    return names, embeddings


def encode_remarks(texts, model=None, show_progress=True):
    """Encode listing descriptions with the paper's embedding model."""
    if model is None:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(config.EMBEDDING_MODEL)
    return model.encode(
        list(texts),
        batch_size=config.ENCODE_BATCH_SIZE,
        show_progress_bar=show_progress,
        normalize_embeddings=True,
    )


def load_or_encode_remarks(texts, cache_path=config.EMBEDDINGS_NPY, force=False):
    """Load cached embeddings if they match the sample size, else encode.

    Embeddings are deterministic for a given model version, so the cache is a
    pure speed-up; pass force=True to re-encode from scratch.
    """
    if not force and cache_path.exists():
        cached = np.load(cache_path)
        if len(cached) == len(texts):
            return cached, True
    embeddings = encode_remarks(texts)
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    np.save(cache_path, embeddings)
    return embeddings, False


def similarity_features(prop_embeddings, ref_embeddings):
    """Cosine similarities (dot product of normalized vectors): n x 20 matrix."""
    return prop_embeddings @ ref_embeddings.T


def add_similarity_columns(df, sim_matrix):
    """Attach the 20 ``sim_<slug>`` columns to the DataFrame (in place)."""
    for i, col in enumerate(SIM_COLS):
        df[col] = sim_matrix[:, i]
    return df