from transformers import pipeline
import wikipedia
wikipedia.set_lang("es")
import random
import gradio as gr
import textwrap

model_name = "PlanTL-GOB-ES/roberta-base-bne-sqac"

nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)


def get_wiki_article(topic):
    topic=topic
    try:
        search = wikipedia.search(topic, results = 1)[0]
    except wikipedia.DisambiguationError as e:
        choices = [x for x in e.options if ('disambiguation' not in x) and ('All pages' not in x) and (x!=topic)]
        search = random.choice(choices)
    try:
        p = wikipedia.page(search)
    except wikipedia.exceptions.DisambiguationError as e:
        choices = [x for x in e.options if ('disambiguation' not in x) and ('All pages' not in x) and (x!=topic)]
        s = random.choice(choices)
        p = wikipedia.page(s)
    return p.content, p.url

def get_answer(topic, question):
    w_art, w_url=get_wiki_article(topic)
    if len(w_art) < 7000:
        #w_art = w_art[:3000]
        qa = {'question': question, 'context': w_art}
        res = nlp(qa)
        res['type'] = "one-pass"
        if res['answer'] not in ['\n',' ','']:
            ans = res['answer']
            if res["start"] < 1000:
                st = 0
            else:
                st = res["start"] - 1000
            en = res["end"]+1000
            contextual = w_art[st:res["start"]]+ " --> "+ans+ " <-- "+w_art[res["end"]:en]
            #thetext = w_art[res['start']-500:res['start']+500:]
            res['text']  = contextual
        #thetext = w_art[res['start']-200:res['start']+200:]
        #res['text']  = thetext
    else:
        res = recursive_rank(w_art,question)
        res['type'] = "recursive"
    return res['answer'], res['text'],'<a href="'+w_url+'">'+w_url+'</a>', {'confidence':res['score']},res['type']

def recursive_rank(w_art,question):
    chk = 1000
    parts = round(len(w_art)/1000)
    size = round(len(w_art)/parts)
    trozos = textwrap.wrap(w_art, size)
    
    answers = []
    for chunk in trozos:
        if len(chunk) < 1000:
            pass
        else:
            qa = {'question': question, 'context': chunk}
            res = nlp(qa)
            if res['answer'] not in ['\n',' ','']:
                ans = res['answer']
                if res["start"] < chk:
                    st = 0#0newres["start"]
                else:
                    st = res["start"] - chk
                # if newres["end"] < 500:
                #     en = newres["end"]
                # else:
                en = res["end"]+chk
                contextual = chunk[st:res["start"]]+ " --> "+ans+ " <-- "+chunk[res["end"]:en]
                thetext = contextual#chunk[res['start']-300:res['start']+300:]
                res['text']  = thetext
                answers.append(res)
    answers.sort(key=lambda x: x['score'], reverse=True) 
    return answers[0]


inputs = [
          gr.inputs.Textbox(lines=5, label="Tema"),
          gr.inputs.Textbox(lines=5, label="Pregunta")
]
outputs = [
            gr.outputs.Textbox(type='str',label="Respuesta"),
            gr.outputs.Textbox(type='str',label="Contexto de la Respuesta"),
            gr.outputs.HTML(label="Articulo de Referencia"),
            gr.outputs.Label(type="confidences",label="Confianza en la respuesta (asumiendo que la página de referencia es correcta para el tema)"),
            gr.outputs.Textbox(type='str',label="Tipo de Búsqueda"),
]

title = "Pregunta/Respuesta en la Wikipedia en Castellano"
description = """
<center>
<img src="https://huggingface.co/spaces/crodri/squacdemo/resolve/main/plantl.png" width=200px> 
<bold>
Pregunta a la Wikipedia en castellano lo que quieras saber. Proponle un tema, y haz una pregunta sobre es tema y el sistema buscará la página más relevante donde señalarte la respuesta.
</bold>

El sistema ha aprendido a contestar preguntas entrenando un <a href="https://huggingface.co/PlanTL-GOB-ES/roberta-large-bne-sqac">modelo neuronal</a> con el conjunto 
<a href="https://huggingface.co/datasets/PlanTL-GOB-ES/SQAC"> SQAC </a> creado por el Grupo de Minería del BSC/CNS, como parte del PlanTL de impulso a las tecnologías del lenguaje.

Puedes hacer preguntas acerca del qué, quien, cuando, cuales, etc.

Si la respuesta no es correcta o la página seleccionada no és relevante, intente replantear la pregunta o buscar un término más específico como tema.

Los articulos de más de 7 mil caracteres pueden tardar en procesarse ya que se realiza una búsqueda recursiva.


</center>

"""
article = """<center>


</center>
"""
ejemplos = [
    ['Búsqueda de respuestas','¿Qué es el Question Answering?'],
    ['Invasión de Yugoslavia', '¿Cómo denominaron las fuerzas del Eje la invasión de Yugoslavia?'],
    ['Duna de Pilat', '¿Cómo son las pequeñas dunas que se encuentran junto a la Duna de Pilat?'],
    ['Fútbol', '¿Cuándo empezó a celebrarse la Copa Mundial de Fútbol?'],
    ['Batalla del Ebro', '¿Quién era jefe del Estado Mayor republicano?'],
    ['Batalla del Ebro', '¿Qué bando ganó la Batalla?'],


]

gr.Interface(get_answer, inputs, outputs, title=title, description=description, examples=ejemplos, article=article, 
 theme="default", flagging_options=["incorrect", "correct"]).launch(share=False,enable_queue=False)