backend/embeddings.py

"""
Word Embeddings Engine for Contexto game.
Handles cosine similarity calculation and ranking generation.
"""

import os
import numpy as np
import unicodedata
import re
from gensim.models import KeyedVectors

MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_FILENAME = os.environ.get("PONELE_MODEL_FILE", "es_bert_vectors_50k.kv")
MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILENAME)


def normalize_word(word):
    """
    Normalize a word for lookup: lowercase, remove accents/tildes, strip whitespace.
    Example: 'CAFÉ' -> 'cafe', 'Niño' -> 'nino'
    """
    word = word.strip().lower()
    # Remove accents/diacritics
    nfkd = unicodedata.normalize('NFKD', word)
    word_clean = ''.join(c for c in nfkd if not unicodedata.combining(c))
    # Replace ñ -> n is handled by NFKD? No, ñ is not a combining char.
    # Keep ñ as-is since it's in the model vocabulary
    # Actually, let's try both versions
    return word_clean


nlp_model = None

def get_nlp():
    global nlp_model
    if nlp_model is None:
        import spacy
        nlp_model = spacy.load("es_core_news_sm", disable=["parser", "ner"])
    return nlp_model

def normalize_for_lookup(word):
    """
    Try multiple normalizations to find the word in the model.
    Returns a list of candidates to try (original, lemma, no_accents).
    """
    original = word.strip().lower()
    candidates = [original]

    # Get lemma using SpaCy
    try:
        doc = get_nlp()(original)
        if len(doc) > 0:
            lemma = doc[0].lemma_.lower()
            
            # IMPROVEMENT: Manual bridge for -aría/-arían which SpaCy misses
            if original.endswith(('aria', 'arian', 'ariais', 'ariamos', 'arias')):
                # Try to see if stripping the suffix and adding 'ar', 'er', or 'ir' works
                # For now, let's just make sure both 'llegar' and 'llegaría' are tried
                pass 

            if lemma not in candidates:
                candidates.append(lemma)
            
            # Force infinitive if it ends in 'aría' (simple heuristic for this specific bug)
            if original.endswith('aria') and len(original) > 4:
                # We try common infinitive endings
                for suffix in ['ar', 'er', 'ir']:
                    verb_candidate = original[:-4] + suffix
                    if verb_candidate not in candidates:
                        candidates.append(verb_candidate)
    except:
        pass

    # Version without accents
    nfkd = unicodedata.normalize('NFKD', original)
    no_accents = ''.join(c for c in nfkd if not unicodedata.combining(c))
    if no_accents not in candidates:
        candidates.append(no_accents)
    
    # Also try removing accents from the lemma
    for cand in list(candidates):
        nfkd_cand = unicodedata.normalize('NFKD', cand)
        clean_cand = ''.join(c for c in nfkd_cand if not unicodedata.combining(c))
        if clean_cand not in candidates:
            candidates.append(clean_cand)

    if no_accents not in candidates:
        candidates.append(no_accents)

    return candidates


class EmbeddingsEngine:
    """
    Core engine that loads word vectors and computes rankings
    based on cosine similarity.
    """

    def __init__(self):
        self.model = None
        self._rankings_cache = {}  # Cache: secret_word -> {word: rank}

    def load_model(self):
        """Load the pre-trained word vectors from disk."""
        if not os.path.exists(MODEL_PATH):
            raise FileNotFoundError(
                f"Model not found at {MODEL_PATH}. "
                "Run 'python build_bert_model.py' first."
            )

        print("🧠 Loading word embeddings model...")
        self.model = KeyedVectors.load(MODEL_PATH)
        print(f"✅ Model loaded! Vocabulary: {len(self.model.key_to_index)} words")

    def get_vector(self, word):
        """
        Get the embedding vector for a word.
        Tries multiple normalizations to find a match.
        Returns (vector, found_word) or (None, None) if not found.
        """
        candidates = normalize_for_lookup(word)

        for candidate in candidates:
            if candidate in self.model.key_to_index:
                return self.model[candidate], candidate

        return None, None

    def cosine_similarity(self, vec_a, vec_b):
        """
        Compute cosine similarity between two vectors.

        Formula: sim(A, B) = (A · B) / (||A|| * ||B||)

        Returns a value between -1 and 1, where:
          1 = identical direction
          0 = orthogonal (unrelated)
         -1 = opposite direction
        """
        dot_product = np.dot(vec_a, vec_b)
        norm_a = np.linalg.norm(vec_a)
        norm_b = np.linalg.norm(vec_b)

        if norm_a == 0 or norm_b == 0:
            return 0.0

        return float(dot_product / (norm_a * norm_b))

    def compute_ranking(self, secret_word):
        """
        Compute the full ranking of all vocabulary words against the secret word.
        The secret word always gets rank #1.

        Returns a dict of {normalized_word: rank} sorted by cosine similarity (descending).
        Results are cached per secret word.
        """
        cache_key = secret_word.lower()

        if cache_key in self._rankings_cache:
            print(f"📋 Using cached ranking for '{secret_word}'")
            return self._rankings_cache[cache_key]

        print(f"🔄 Computing ranking for secret word: '{secret_word}'...")

        # Get the secret word's vector
        secret_vec, secret_found = self.get_vector(secret_word)
        if secret_vec is None:
            raise ValueError(f"Secret word '{secret_word}' not found in vocabulary")

        # Compute similarity for ALL words in vocabulary
        vocab_words = list(self.model.key_to_index.keys())

        # Batch computation using gensim's built-in method (much faster)
        similarities = self.model.most_similar(positive=[secret_found], topn=len(vocab_words))

        # Build ranking dict: rank 1 is the secret word itself
        ranking = {secret_found: 1}
        for rank, (word, _sim) in enumerate(similarities, start=2):
            ranking[word] = rank

        self._rankings_cache[cache_key] = ranking
        print(f"✅ Ranking computed: {len(ranking)} words ranked")

        return ranking

    def get_rank_for_word(self, guess_word, secret_word):
        """
        Get the rank of a guessed word relative to the secret word.

        Returns:
            dict with 'rank', 'similarity', 'found', and 'normalized_word'
        """
        # Normalize the guess
        guess_candidates = normalize_for_lookup(guess_word)
        secret_candidates = normalize_for_lookup(secret_word)

        # Check if it's the same word (exact match)
        guess_norm = guess_candidates[0]
        secret_norm = secret_candidates[0]

        if guess_norm == secret_norm:
            return {
                "rank": 1,
                "similarity": 1.0,
                "found": True,
                "normalized_word": guess_norm,
                "is_correct": True
            }

        # Make sure both words exist in the model
        guess_vec, guess_found = self.get_vector(guess_word)
        secret_vec, secret_found = self.get_vector(secret_word)

        if guess_vec is None:
            return {
                "rank": -1,
                "similarity": 0.0,
                "found": False,
                "normalized_word": guess_norm,
                "is_correct": False,
                "error": f"Word '{guess_word}' not found in vocabulary"
            }

        if secret_vec is None:
            return {
                "rank": -1,
                "similarity": 0.0,
                "found": False,
                "normalized_word": secret_norm,
                "is_correct": False,
                "error": f"Secret word '{secret_word}' not found in vocabulary"
            }

        # Get the precomputed ranking
        ranking = self.compute_ranking(secret_word)

        # Look up the guess word in the ranking
        rank = None
        for candidate in guess_candidates:
            if candidate in ranking:
                rank = ranking[candidate]
                break

        if rank is None:
            # Word exists in model but wasn't in the ranking (shouldn't happen)
            # Fallback: compute similarity directly
            similarity = self.cosine_similarity(guess_vec, secret_vec)
            # Estimate rank based on similarity
            rank = max(1, int((1 - similarity) * len(ranking)))

        # Determine max rank for normalization
        max_rank = len(ranking) if ranking else 50000
        
        # LOG-SCALING NORMALIZATION (Phase 2 Upgrade)
        # We transform the raw cosine similarity into a 'game feel' similarity.
        # This stretches the top 1500 ranks into the 50%-100% range to avoid the "dense blob" effect of BERT.
        import math
        if rank == 1:
            stretched_sim = 1.0
        else:
            # Power law decay to make top ranks "pop"
            stretched_sim = max(0.01, 1.0 - math.pow(rank / max_rank, 0.4))
            
            # Boost the top 1500 specifically to make the transition satisfying
            if rank <= 1500:
                stretched_sim = max(stretched_sim, 0.5 + 0.5 * (1 - (rank / 1500)**0.8))

        return {
            "rank": rank,
            "similarity": round(stretched_sim, 4),
            "found": True,
            "normalized_word": guess_found or guess_norm,
            "is_correct": rank == 1
        }

    def get_hint(self, secret_word, excluded_words):
        """
        Get a hint word: a word that is semantically close but not yet guessed.
        """
        ranking = self.compute_ranking(secret_word)
        
        # Normalize excluded words
        norm_excluded = set()
        for w in excluded_words:
            candidates = normalize_for_lookup(w)
            norm_excluded.update(candidates)
            
        # Find the best rank guessed so far (excluding rank 1)
        best_rank = 50000 # Default if no guesses or all very far
        for word in norm_excluded:
            if word in ranking:
                r = ranking[word]
                if r > 1 and r < best_rank:
                    best_rank = r
        
        # Target rank: half of best rank, but at most 1000 (if best is very far)
        # and at least 2 (can't give the secret word as a hint)
        target_rank = max(2, min(1000, best_rank // 2))
        
        # We want a word around the target rank. 
        # Since ranking is {word: rank}, we can iterate.
        # But we want it to be efficient.
        
        # Items are (word, rank)
        # ranking.items() is already somewhat sorted by rank in compute_ranking? 
        # No, it's a dict. But similarities from model.most_similar were sorted.
        
        # Find a word with rank close to target_rank
        # For small vocab (50k), we can just iterate.
        sorted_ranking = sorted(ranking.items(), key=lambda x: x[1])
        
        for word, rank in sorted_ranking:
            if rank >= target_rank and word not in norm_excluded and rank > 1:
                # Add a bit of randomness so it's not always the same?
                # For now let's just return the first one available after target.
                return {"word": word, "rank": rank}
                
        return None


    def get_vocabulary_size(self):
        """Return the number of words in the loaded model."""
        if self.model is None:
            return 0
        return len(self.model.key_to_index)

    def word_exists(self, word):
        """Check if a word exists in the vocabulary."""
        _, found = self.get_vector(word)
        return found is not None

    def get_all_words(self):
        """
        Return all vocabulary words in frequency order (most common first).
        KeyedVectors stores keys in the order they appear in the original model,
        which for FastText models is typically frequency-sorted.
        """
        if self.model is None:
            return []
        return list(self.model.key_to_index.keys())