#!/usr/bin/env python3
"""Experimento: mide si un embedder transformer supera al model2vec estatico.

    ./experimento_embedder.py                 prueba intfloat/multilingual-e5-base
    ./experimento_embedder.py BAAI/bge-m3      prueba otro modelo

## Por que este experimento

model2vec (potion-multilingual-128M) es un embedder ESTATICO: rapidisimo (un
vector por token, precalculado, sin red neuronal en la consulta) pero flojo en
lo semantico. Su punto ciego se ve en las dos preguntas que el RAG no resuelve:

    "detectar sql injection a ciegas basada en tiempo"  ->  no encuentra SLEEP
    "descubrir directorios ocultos en una web"          ->  no encuentra gobuster

Las dos son cross-lingual: la pregunta va en castellano coloquial y el comando
en ingles tecnico ("time-based", "content discovery"). Un transformer de verdad
—e5, bge— entiende eso; el estatico no.

Este experimento NO toca nada: coge una MUESTRA del indice vivo (para que quepa
en memoria y sea rapido), la re-embebe con el transformer, y corre la misma
evaluacion. Si gana claro, merece la pena cambiar busca.py; si no, no.

## El coste que habria que pagar si se adopta

El estatico encodea la consulta en microsegundos. Un transformer en CPU tarda
~50-150 ms por consulta. Para un asistente de voz local es asumible, pero no es
gratis, y por eso se mide antes de decidir.
"""
import json
import os
import re
import sys
import unicodedata

AQUI = os.path.dirname(os.path.abspath(__file__))
RAIZ = os.path.dirname(os.path.dirname(AQUI))
VIVO = os.path.join(RAIZ, "datos", "saber.jsonl")
EVAL = os.path.join(AQUI, "eval_set.jsonl")
MODELO_TF = sys.argv[1] if len(sys.argv) > 1 else "intfloat/multilingual-e5-base"

# Muestra: todo lo que un caso de eval podria querer + relleno, para no cargar
# 14.000 vectores de transformer en RAM. Se cogen todas las entradas de
# metodologia y pregunta (el oro) mas una parte del resto.
MAX_ENTRADAS = 6000


def _norm(t):
    t = unicodedata.normalize("NFD", (t or "").lower())
    t = "".join(c for c in t if unicodedata.category(c) != "Mn")
    return re.sub(r"[^a-z0-9 ]", " ", t)


def _muestra():
    oro, resto = [], []
    with open(VIVO, encoding="utf-8") as f:
        for l in f:
            try:
                d = json.loads(l)
            except json.JSONDecodeError:
                continue
            (oro if d.get("tipo") in ("metodologia", "pregunta") else resto).append(d)
    ent = oro + resto[: max(0, MAX_ENTRADAS - len(oro))]
    return ent


def _eval(nombre, encode_corpus, encode_query, entradas):
    import numpy as np
    textos = [(e.get("indexar") or e["texto"]) for e in entradas]
    M = encode_corpus(textos)
    M = M / np.clip(np.linalg.norm(M, axis=1, keepdims=True), 1e-9, None)
    palabras_e = [set(_norm(e["texto"]).split()) for e in entradas]
    nom = [_norm(e.get("herramienta", "")) for e in entradas]

    casos = [json.loads(l) for l in open(EVAL, encoding="utf-8") if l.strip()]
    hit1 = hit5 = 0
    fallos = []
    for c in casos:
        q = encode_query([c["pregunta"]])[0]
        q = q / max(float(np.linalg.norm(q)), 1e-9)
        sim = M @ q
        pal = {w for w in _norm(c["pregunta"]).split() if len(w) > 3}
        if pal:
            bono = np.array([0.08 * len(pal & palabras_e[i]) +
                             (0.5 if nom[i] and nom[i] in pal else 0.0)
                             for i in range(len(entradas))], dtype="float32")
            sim = sim + bono
        orden = np.argsort(-sim)[:5]
        esperado = [_norm(k).strip() for k in c["espera"]]
        tops = [_norm(entradas[i]["texto"]) for i in orden]
        ok = lambda t: any(k and k in t for k in esperado)
        if tops and ok(tops[0]):
            hit1 += 1
        if any(ok(t) for t in tops):
            hit5 += 1
        else:
            fallos.append(c["pregunta"])
    n = len(casos)
    print(f"\n  {nombre}")
    print(f"    hit@1  {hit1:3d}/{n}  ({100*hit1//n}%)")
    print(f"    hit@5  {hit5:3d}/{n}  ({100*hit5//n}%)")
    if fallos:
        print(f"    no resuelve {len(fallos)}: " + "; ".join(fallos[:6]))
    return hit1, hit5


def main():
    entradas = _muestra()
    print(f"Muestra: {len(entradas)} entradas del indice vivo")

    # 1) model2vec estatico, la referencia
    from model2vec import StaticModel
    est = StaticModel.from_pretrained("minishlab/potion-multilingual-128M")
    enc_est = lambda xs: est.encode(xs)
    h1e, h5e = _eval("model2vec (estatico, el actual)", enc_est, enc_est, entradas)

    # 2) el transformer a prueba
    print(f"\n  cargando {MODELO_TF} (puede descargar ~1-2 GB la 1a vez)...", flush=True)
    from sentence_transformers import SentenceTransformer
    tf = SentenceTransformer(MODELO_TF)
    # los modelos e5 piden prefijos "query:" y "passage:"
    es_e5 = "e5" in MODELO_TF.lower()
    enc_doc = lambda xs: tf.encode([("passage: " + x) if es_e5 else x for x in xs],
                                   show_progress_bar=False, batch_size=32)
    enc_q = lambda xs: tf.encode([("query: " + x) if es_e5 else x for x in xs],
                                 show_progress_bar=False, batch_size=32)
    h1t, h5t = _eval(f"{MODELO_TF} (transformer)", enc_doc, enc_q, entradas)

    print("\n" + "=" * 44)
    print(f"  hit@1  estatico {h1e}  ->  transformer {h1t}   ({h1t-h1e:+d})")
    print(f"  hit@5  estatico {h5e}  ->  transformer {h5t}   ({h5t-h5e:+d})")
    print("=" * 44)
    if h5t > h5e or (h5t == h5e and h1t > h1e):
        print("  El transformer gana. Merece plantear el cambio en busca.py")
        print("  (coste: ~50-150 ms por consulta en CPU, hoy es instantaneo).")
    else:
        print("  El transformer NO compensa aqui. El estatico se queda.")
    return 0


if __name__ == "__main__":
    sys.exit(main())