import torch
import gradio as gr
from threading import Thread
from typing import Iterator
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.use_default_system_prompt = False

def generar_respuesta(
    mensaje: str,
    historial_chat: list[tuple[str, str]],
    max_tokens_nuevos: int = 512,
    temperatura: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    penalización_repetición: float = 1.2,
) -> Iterator[str]:

    # Preparamos la conversación en un formato adecuado
    conversacion = []

    mensaje_sistema = '''Eres un chatbot llamado EDUAI. Fuiste desarrollado por la universidad Yachay Tech. Tu misión es dar ayuda a los estudiantes de matemáticas de colegios y universidad.

    Debes de seguir las siguientes reglas de manera estricta.

    1. Eres un asistente, no eres un usuario. Responde como asistente.
    2. Tus respuestas son amables. Concisas y rápidas.
    3. No saludas en cada respuesta, solo responde la pregunta que realiza el usuario.
    4. Anima al estudiante a seguir aprendiendo continuamente.
    5. Responde en español siempre.
    6. Tus respuestas son en formato markdown y con viñetas.
    7. Usa los siguientes ejercicios y material audiovisual para complementar tu respuesta, de no existir algo que mostrar no escribas nada.

    '''

    # Si hay un mensaje del sistema, lo agregamos a la conversación
    if mensaje_sistema:
        conversacion.append({"role": "system", "content": mensaje_sistema})

    # Luego, agregamos los mensajes del usuario y el asistente del historial de chat
    for usuario, asistente in historial_chat:
        conversacion.extend([{"role": "user", "content": usuario}, {"role": "assistant", "content": asistente}])

    # Por último, agregamos el mensaje del usuario actual
    conversacion.append({"role": "user", "content": mensaje})

    # Convertimos la conversación en tokens de entrada utilizando el tokenizador
    input_ids = tokenizer.apply_chat_template(conversacion, return_tensors="pt")

    # Si la longitud de los tokens de entrada excede un límite, los recortamos
    if input_ids.shape[1] > 2048:
        input_ids = input_ids[:, -2048:]
        #gr.Advertencia(f"Se recortó la entrada de la conversación, ya que era más larga que {2048} tokens.")

    # Transferimos los tokens de entrada al dispositivo del modelo
    input_ids = input_ids.to(model.device)

    # Configuramos un "streamer" para obtener las respuestas en tiempo real
    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)

    # Definimos los parámetros para generar la respuesta
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_tokens_nuevos,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperatura,
        num_beams=1,
        repetition_penalty=penalización_repetición,
    )

    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    # Inicializamos una lista para almacenar las salidas
    salidas = []

    # Recorremos las respuestas obtenidas del "streamer" y las agregamos a la lista
    for texto in streamer:
        salidas.append(texto)
        # Generamos la salida acumulada hasta el momento
        yield "".join(salidas)

def vote(data: gr.LikeData):
    if data.liked:
        print("Fue util: " + data.value)
    else:
        print("No fue util: " + data.value)

# Crear instancias de ChatInterface para cada materia
chat_interface_math = gr.ChatInterface(
    generar_respuesta,
    title="EDUAI-BOT: Asistente de Chat (Matemáticas)",
    description="¡Bienvenido a EDUAI-BOT! Estoy aquí para ayudarte con preguntas de matemáticas.",
    theme="dark",
    examples=[["¿Puedes ayudarme con álgebra?"]],
)

chat_interface_fisica = gr.ChatInterface(
    generar_respuesta,
    title="EDUAI-BOT: Asistente de Chat (Física)",
    description="¡Bienvenido a EDUAI-BOT! Estoy aquí para ayudarte con preguntas de física.",
    theme="dark",
    examples=[["¿Qué es la ley de Newton?"]],
)

chat_interface_quimica = gr.ChatInterface(
    generar_respuesta,
    title="EDUAI-BOT: Asistente de Chat (Química)",
    description="¡Bienvenido a EDUAI-BOT! Estoy aquí para ayudarte con preguntas de química.",
    theme="dark",
    examples=[["¿Qué es una reacción de síntesis?"]],
)

chat_interface_biologia = gr.ChatInterface(
    generar_respuesta,
    title="EDUAI-BOT: Asistente de Chat (Biología)",
    description="¡Bienvenido a EDUAI-BOT! Estoy aquí para ayudarte con preguntas de biología.",
    theme="dark",
    examples=[["Explícame la fotosíntesis."]],
)

chat_interface_ingles = gr.ChatInterface(
    generar_respuesta,
    title="EDUAI-BOT: Asistente de Chat (Inglés)",
    description="¡Bienvenido a EDUAI-BOT! Estoy aquí para ayudarte con preguntas de inglés.",
    theme="dark",
    examples=[["Translate 'hello' to Spanish."]],
)

with gr.Blocks() as demo:
    chat_interface_math.render()

if __name__ == "__main__":
    demo.queue(max_size=20).launch(debug=True)