from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.retrievers import BaseRetriever
from langchain_community.vectorstores import FAISS  #Facebook AI Similarity Search
from sentence_transformers import CrossEncoder
from langchain_core.documents import Document
from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI
from typing import List
import pandas as pd

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",     # Ruta a modelo Pre entrenado
    model_kwargs={'device':'cpu'}, # Opciones de configuracion del modelo
    encode_kwargs={'normalize_embeddings': False}) # Opciones de Encoding

try:
    vectorstore = FAISS.load_local("cache", embeddings)
except:    
    loader = PyPDFDirectoryLoader("data/")
    data = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=450, chunk_overlap=100, length_function=len)
    docs = text_splitter.split_documents(data)

    #DB y retriever
    vectorstore = FAISS.from_documents(docs, embeddings)  # Create a retriever object from the 'db' with a search configuration where it retrieves up to 4 relevant splits/documents.
    vectorstore.save_local("cache")

#Renranker para mejorar respuestas
model = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-2-v2', max_length=512)  #Por lejos el mejor, los otros no sirven
class Reranking_retriever(BaseRetriever):
  def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
      busqueda = vectorstore.similarity_search_with_score(query,  k=10, fetch_k=15)       # k = 10 numero total de documento a traer previo al re ranking

      df = pd.DataFrame({ # Funciones lambda toman la ultima variable como input y la previa como iteracionm la primera x es que se retornara
        'scores':  list(map(lambda x : x[-1], busqueda)),
        'respuestas': list(map(lambda x : x[0].page_content, busqueda)),
        'metadata': list(map(lambda x : x[0].metadata ,busqueda))})

      print(df.scores)
      respuestas = df.respuestas.to_list()  #lista de respuestas

      sentence_combinations = [[query, respuesta] for respuesta in respuestas]      # So we create the respective sentence combinations

      scores = model.predict(sentence_combinations)  #Aplica cross encoding para ver que para de q y a tienen mayor relacion, en este caso se manda la pregunta en cada una de ellas y se compara una a una con las respuestas
      scores = scores.argsort()[::-1] #Ordena puntajes de mas relevate a menos relevante siendo indice 0 el mas relevante

      docs = []
      for i in scores[:3]:  #Los 3 resulados mas relevantes
        docs.append(Document(page_content=df.respuestas[i], metadata=df.metadata[i]))
      return docs

retriever = Reranking_retriever()  #Mi retriever personalizado, de 10 elementos retorna 3 filtrando por un cross encoder

QA_CHAIN_PROMPT = PromptTemplate.from_template(""" 
Usa el siguiente contexto para responder la pregunta.

Contexto
{contexto}

Pregunta: {question}
Respuesta Util:"""
)

def get_chain():    
 
    llm = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0)

    llm_chain = LLMChain(llm=llm, prompt=QA_CHAIN_PROMPT, callbacks=None, verbose=True)

    document_prompt = PromptTemplate( input_variables=["page_content"], template="Contexto:\n{page_content}")
    
    combine_documents_chain = StuffDocumentsChain( llm_chain=llm_chain, document_variable_name="contexto", document_prompt=document_prompt, callbacks=None)
    
    chain = RetrievalQA(combine_documents_chain=combine_documents_chain, callbacks=None, verbose=True, retriever=retriever)
    
    return(chain)