# ENUNCIADO
# rag_engine.py
#
# Este archivo contendrá toda la lógica del motor RAG. Se deben
# implementar obligatoriamente las siguientes funciones (con los nombres y parámetros
# exactos que se indican).
# Al inicio del script se deben cargar:
#   • El modelo de embeddings: SentenceTransformer("MongoDB/mdbr-leaf-ir")
#   • El modelo de lenguaje: PleIAs/Pleias-RAG350M (usando AutoTokenizer y AutoModelForCausalLM de transformers).
#   • Los documentos desde documents.json.
# Función recuperar_documentos(consulta, top_k=2, umbral=0.4)
# Dada una consulta en inglés, recupera los documentos más relevantes de la base de conocimiento.
#   • Parámetros:
#       o consulta (str): pregunta del usuario.
#       o top_k (int): número máximo de documentos a retornar.
#       o umbral (float): valor mínimo de similitud (coseno) para considerar un
#         documento relevante. Los documentos con similitud inferior a este
#         umbral se descartan.
#   • Proceso:
#       1. Calcular el embedding de la consulta y de todos los documentos
#          (preferiblemente una sola vez al cargar el script y almacenarlos para
#          evitar recalcular).
#       2. Calcular la similitud del coseno entre el embedding de la consulta y los
#          embeddings de los documentos.
#       3. Ordenar los documentos de mayor a menor similitud.
#       4. Recorrer en ese orden y seleccionar aquellos cuya similitud sea mayor o
#          igual al umbral, hasta un máximo de top_k documentos.
#   • Retorno: Lista con los textos de los documentos seleccionados.
# Función generar_respuesta(consulta, documentos_recuperados)
# Genera una respuesta usando el modelo de lenguaje, inyectando los documentos
# recuperados como contexto.
# Parámetros:
#   o consulta (str): pregunta original del usuario.
#   o documentos_recuperados (list): lista de textos con los documentos
# relevantes.
# Proceso:
#   1. Se concatenan todos los documentos en un solo string (por ejemplo,
#   separados por espacios).
#   2. Se construye un prompt con el siguiente formato:
#   “””
#   Answer the question based only on the context provided
#   Context: <" ".join(documentos_recuperados)>
#   Question: <consulta>
#   Answer:
#   “””
#   3. Se genera la respuesta con el modelo
# Retorno: Cadena con la respuesta generada.
# Función preguntar(consulta, top_k=2, umbral=0.4)
#   • Descripción:
#       o Función de alto nivel que une la lógica de recuperar_documentos y
#         generar_respuestas
#   • Parámetros: los mismos que recuperar_documentos.
#   • Retorno: La respuesta generada (cadena).

import json
import torch

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

from sklearn.metrics.pairwise import cosine_similarity


# -----------------------------
# Cargar documentos
# -----------------------------

with open("documents.json", "r") as f:
    documents = json.load(f)

# convertir a lista de textos
docs_text = list(documents.values())


# -----------------------------
# Modelo de embeddings
# -----------------------------

embed_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")

# calcular embeddings una sola vez
doc_embeddings = embed_model.encode(docs_text)


# -----------------------------
# Modelo de lenguaje (LLM)
# -----------------------------

model_name = "microsoft/phi-2"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)


# -------------------------------------------------
# FUNCION 1
# recuperar_documentos
# -------------------------------------------------


def recuperar_documentos(consulta, top_k=2, umbral=0.4):

    # embedding de la consulta
    query_embedding = embed_model.encode([consulta])

    # calcular similitud coseno
    similitudes = cosine_similarity(query_embedding, doc_embeddings)[0]

    # ordenar índices por similitud
    indices_ordenados = similitudes.argsort()[::-1]

    docs_relevantes = []

    for idx in indices_ordenados:

        if similitudes[idx] >= umbral:
            docs_relevantes.append(docs_text[idx])

        if len(docs_relevantes) >= top_k:
            break

    return docs_relevantes


# -------------------------------------------------
# FUNCION 2
# generar_respuesta
# -------------------------------------------------


def generar_respuesta(consulta, documentos_recuperados):

    contexto = " ".join(documentos_recuperados)

    prompt = f"""
        Answer the question based only on the context provided

        Context: {contexto}

        Question: {consulta}

        Answer:
        """

    inputs = tokenizer(prompt, return_tensors="pt")

    outputs = model.generate(**inputs, max_new_tokens=100)

    respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return respuesta


# -------------------------------------------------
# FUNCION 3
# preguntar
# -------------------------------------------------


def preguntar(consulta, top_k=2, umbral=0.4):

    docs = recuperar_documentos(consulta, top_k, umbral)

    respuesta = generar_respuesta(consulta, docs)

    return respuesta