import gradio as gr
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration, pipeline
import torch
from gtts import gTTS


# Cargamos el modelo para el chat
model_name = 'facebook/blenderbot-400M-distill'
tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
model = BlenderbotForConditionalGeneration.from_pretrained(model_name)


# Cargamos el traductor de ingles a español
english_model_name = "Helsinki-NLP/opus-mt-en-es"
translator_en_es = pipeline("translation", model=english_model_name)

# Cargamos el traductor de español a ingles
spanish_model_name = "Helsinki-NLP/opus-mt-es-en"
translator_es_en = pipeline("translation", model=spanish_model_name)


def take_last_tokens(inputs, note_history, history):
    """Filtrar los últimos 128 tokens"""
    if inputs['input_ids'].shape[1] > 128:
        inputs['input_ids'] = torch.tensor([inputs['input_ids'][0][-128:].tolist()])
        inputs['attention_mask'] = torch.tensor([inputs['attention_mask'][0][-128:].tolist()])
        note_history = ['</s> <s>'.join(note_history[0].split('</s> <s>')[2:])]
        history = history[1:]
    return inputs, note_history, history

def add_note_to_history(note, note_history):
    """Añadir una nota a la información histórica del chat"""
    note_history.append(note)
    note_history = '</s> <s>'.join(note_history)
    return [note_history]


def predict(text, history):
      history = history or []
      if history: 
          history_useful = ['</s> <s>'.join([str(a[0])+'</s> <s>'+str(a[1]) for a in history])]
      else:
          history_useful = []
      # Traducimos el texto ingresado a ingles
      text_input = translator_es_en(text)[0]['translation_text']

      # comparamos con el historial y codificamos la nueva entrada del usuario
      history_useful = add_note_to_history(text_input, history_useful)  
      inputs = tokenizer(history_useful, return_tensors="pt")
      inputs, history_useful, history = take_last_tokens(inputs, history_useful, history)

      # Generar una respuesta
      reply_ids = model.generate(**inputs)
      response = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
      # sumamos la respuesta al historial del chat
      history_useful = add_note_to_history(response, history_useful)
      list_history = history_useful[0].split('</s> <s>')
      history.append((list_history[-2], list_history[-1]))

      # pasamos el resultado a gtts para obtener el audio
      spanish_text = translator_en_es(response)
      result_es = spanish_text[0]['translation_text']
      sound_file = 'output.wav'
      tts = gTTS(result_es, lang="es", tld='com.mx')
      tts.save(sound_file)
      return sound_file, history

description = """
<h2 style="text-align:center">Inicia el chat con la IA que ha sido entrenada para hablar contigo sobre lo que quieras.</h2>
<h2 style="text-align:center">¡Hablemos!</h2>
"""
article = """Instrucciones: 
\n1. Inserte el texto en la casilla de texto 
\n2. Presionar 'Enviar' y esperar la respuesta
\n4. Para enviar otro texto borrar el actual y volver al punto 1.

El modelo usa: 
- Modelo conversacional [facebook/blenderbot-400M-distill](https://huggingface.co/facebook/blenderbot-400M-distill?text=Hey+my+name+is+Julien%21+How+are+you%3F), 
- Para las traducciones [Helsinki-NLP](https://huggingface.co/Helsinki-NLP) 
- Para la respuesta de voz [gTTS](https://pypi.org/project/gTTS/)
\n... y mucha magia ☺
"""

gr.Interface(fn=predict,
              title="ChatBot Text-to-Speach en Español",
              inputs= [gr.Textbox("", max_lines = 5, label = "Inserte su texto aqui") , 'state'],
              outputs = [gr.Audio(type='file', label="Respuesta de IA en forma de audio"), 'state'],
              description = description ,
              article = article).launch(debug=True)