import gradio as gr
import json
from huggingface_hub import InferenceClient

client = InferenceClient(
    "mistralai/Mistral-7B-Instruct-v0.1"
)

system_prompt = """[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer in spanish. Never answer in English. Always answer as
helpfully as possible, while being safe.

If a question does not make any sense, or is not factually coherent, explain
why instead of answering something not correct. If you don't know the answer
to a question, please don't share false information.

Your goal is to provide answers relating to a legal document, making classifications of arguments inside an allegation.<</SYS>>
"""

rag_text = """
"""

format_text = """{
    'tipo': 'clasificación del argumento (Argumento tipo X)',
    'descripcion': 'Descripcion breve del argumento',
    'texto_referencia': 'texto original completo perteneciente a dicho argumento'
}"""

prompt_template_text = """
                          --- Lista Argumentos ---
    * Argumento tipo 49 (Desvirtuación de los hechos, el denunciado no está de acuerdo con el relato de los hechos)
    * Argumento tipo 304 (El denunciado solicita medios de prueba)
    * Argumento tipo 994 (El denunciado solicita fotografía como medio de prueba, desafiando el relato del agente denunciante)
    * Argumento tipo 1002 (Prescripción, el delito ha prescrito)
    * Argumento tipo 2014 (Principio de proporcionalidad, denunciado considera sanción excesiva)
    * Argumento tipo 2027 (Niega desobediencia de señales de tráfico de prohibición)
    * Argumento tipo 2002 (Denuncia no notificada en el momento en el que se formuló)
    * Argumento tipo 1001 (Defecto de forma en la denuncia, la sanción contiene algún error)
"""

recorte = """
PRIMERA.- Que los hechos descritos en la denuncia no se corresponden con la realidad de lo
sucedido, por lo que los niego expresamente. La carga de la prueba corresponde al órgano instructor
del procedimiento, en virtud del principio de presunción de inocencia recogido en el art.24 de la
Constitución Española y en el art. 53 de la Ley 39/2015 de 1 de Octubre del Procedimiento
Administrativo Común de las Administraciones Públicas.
"""

recorte2 = """
"PRIMERO.- Que entendemos que la actuación policial debe regirse en todo momento por los 
"principios .de oportunidad, congruencia y proporcionalidad, entendiendo esta parte que ha existido 
una confusión de apreciación en el momento de la denuncia, pues no existe más prueba en contra que 
la foto adjuntada a la denuncia, apreciando claramente como venimos diciendo que las luces de freno 
están activadas por el conductor.
"""

def format_prompt(message):
  prompt = "<s>"
  prompt += f"[INST] {message} [/INST]"
  return prompt

def generate(prompt):

    generate_kwargs = dict(
        temperature=0.9,
        max_new_tokens=1024,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt)

    output = client.text_generation(formatted_prompt, **generate_kwargs)

    return output


def process_input(text, rag, prompt_template):

    if not(rag):
        rag_text = " "
    #if not(prompt_template):
    #    prompt_template_text = " "
    
    prompt = f"""
    {system_prompt}

    Teniendo en cuenta que los argumentos se clasifican de la siguiente manera:
    {prompt_template_text}
    
    Partiendo del siguiente fragmento de texto:

    Fragmento:
    ---------------------------------------------------------------
      {text}
    --------------------------------------------------------------

    Identifica y clasifica los argumentos expuestos en el anterior fragmento de texto siguiendo estos pasos:
     - Primero: Determina cúantos argumentos se exponen en el fragmento proporcionado.
     - Segundo: Determina el contenido de cada uno de los argumentos encontrados.
     - Tercero: Clasifica cada uno de los argumentos encontrados en el fragmento con el tipo de argumento que le corresponda según la 'Lista Argumentos''
     - Cuarto: Este es el paso más importante. Devuelve una lista de JSON, cuya longitud será igual al número de argumentos encontrados, que contenga los argumentos clasificados con el siguiente formato:
        {format_text}


    Recuerda responder en español. Los argumentos que no se puedan clasificar o que *claramente* no encajen con ningún tipo predefinido serán clasificados como 'Desconocido'. Recuerda que los tipos de argumentos de representan con un número que puedes encontrar en la información anterior.
    Respira profundamente y piensa paso a paso.
    [/INST]


    """
    output = generate(prompt)

    # Generamos el archivo JSON
    with open('output2.json', 'w') as f:
        json.dump(output, f)
    
    return output

def create_interface():
    # Definimos los componentes de la interfaz
    input_text = gr.Textbox(label="Input")
    rag_checkbox = gr.Checkbox(label="RAG")
    prompt_template = gr.Checkbox(label="PromptTemplate")
    output_text = gr.Textbox(label="Output")
    classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo")

    # Definimos la función que se ejecutará cuando se envíe la entrada
    def fn(text, rag, prompt_template):
        
        output = process_input(text, rag, prompt_template)
        classification_types = [item['clasificacion_tipo'] for item in output]
        classification_types_options = [(option, option) for option in classification_types]
        classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo", choices=classification_types_options, interactive = True)
        
        return output, classification_types_checkboxes
        
    examples = [
        [recorte, False, True],
        [recorte, False, True]
    ]
        
    # Creamos la interfaz
    iface = gr.Interface(
        fn=fn,
        inputs=[input_text, rag_checkbox, prompt_template],
        outputs=[output_text, classification_types_checkboxes],
        examples=examples
    )

    return iface

iface = create_interface()
iface.launch()