import os
from gensim.models import KeyedVectors
from sentence_transformers import SentenceTransformer

MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
# STRATEGY G: We load the exact Golden Dataset (presanitized 50k)
VOCAB_FILE = os.path.join(MODEL_DIR, "es_freq_50k.txt")

def build_bert_vectors():
    print(f"Loading pre-sanitized vocabulary from {VOCAB_FILE}...")
    
    if not os.path.exists(VOCAB_FILE):
        print(f"ERROR: {VOCAB_FILE} not found. Please generate the golden dataset first.")
        return

    lemmatized_vocab = []
    seen = set()
    with open(VOCAB_FILE, 'r', encoding='utf-8') as f:
        for line in f:
            word = line.strip().lower()
            if word and word not in seen:
                lemmatized_vocab.append(word)
                seen.add(word)

    vocab_size_k = round(len(lemmatized_vocab) / 1000)
    print(f"Loaded {len(lemmatized_vocab)} exact golden lemmas ({vocab_size_k}k model).")
    
    versioned_filename = f"es_bert_vectors_{vocab_size_k}k.kv"
    NEW_MODEL_PATH = os.path.join(MODEL_DIR, versioned_filename)

    print("Loading BERT sentence transformer model...")
    model_name = "hiiamsid/sentence_similarity_spanish_es"
    encoder = SentenceTransformer(model_name)
    
    # Determine vector size
    sample_vec = encoder.encode("hola")
    vector_size = len(sample_vec)
    print(f"Model vectors have {vector_size} dimensions.")

    print(f"Encoding {len(lemmatized_vocab)} lemmas in batches... this might take a few minutes.")
    embeddings = encoder.encode(lemmatized_vocab, batch_size=256, show_progress_bar=True, normalize_embeddings=True)

    print("Creating new Gensim KeyedVectors...")
    bert_kv = KeyedVectors(vector_size=vector_size)
    bert_kv.add_vectors(lemmatized_vocab, embeddings)

    vocab_size_k = round(len(lemmatized_vocab) / 1000)
    versioned_filename = f"es_bert_vectors_{vocab_size_k}k.kv"
    versioned_filepath = os.path.join(MODEL_DIR, versioned_filename)

    print(f"Saving BERT embeddings to {versioned_filepath}...")
    bert_kv.save(versioned_filepath)
    print(f"Done! You can now configure app.py and embeddings.py to use {versioned_filename}")

if __name__ == "__main__":
    build_bert_vectors()