import torch
import gradio as gr
from gtts import gTTS
from transformers import pipeline
from transformers import GPT2LMHeadModel, GPT2Tokenizer
css = """
#input {background-color: #FFCCCB} 
"""
# Utility Functions
flatten = lambda l: [item for sublist in l for item in sublist]

def to_data(x):
    return x.data.numpy()

def to_var(x):
    if not torch.is_tensor(x):
        x = torch.Tensor(x)
    return x

def clear():
    return None,[],[]

def append(text, history,dialog_hx,personas):
    history.append([text,None])
    history , audio,dialog_hx= bot.respond(history,dialog_hx,personas)
    return history, audio, None,dialog_hx

class AI_Companion:
    """
    Class that Implements AI Companion.
    """

    def __init__(self, asr = "openai/whisper-tiny", chatbot = "af1tang/personaGPT"):
        """
        Create an Instance of the Companion.
        Parameters:
        asr: Huggingface ASR Model Card. Default: openai/whisper-tiny
        chatbot: Huggingface Conversational Model Card. Default: af1tang/personaGPT
        """

        self.device = torch.device('cpu')
        self.asr = pipeline("automatic-speech-recognition",model = asr,device= -1 if self.device == "cpu" else 0)
        self.model = GPT2LMHeadModel.from_pretrained(chatbot).to(self.device)
        self.tokenizer = GPT2Tokenizer.from_pretrained(chatbot)
        self.personas=[]
        self.sett={
            "do_sample":True,
            "top_k":10,
            "top_p":0.92,
            "max_length":1000,
        }

    def listen(self, audio, history):
        """
        Convert Speech to Text.

        Parameters:
        audio: Audio Filepath
        history: Chat History

        Returns:
        history : history with recognized text appended
        Audio : empty gradio component to clear gradio voice input
        """
        text = self.asr(audio)["text"]
        history.append([text,None])
        return history , None
    
    def add_fact(self,audio,personas,msg):
        '''
        Add fact to Persona.
        Takes in Audio, converts it into text and adds it to the facts list.

        Parameters:
        audio : audio of the spoken fact
        '''
        if audio is not None:
            text=self.asr(audio)
            personas.append(text['text']+self.tokenizer.eos_token)
        else:
            personas.append(msg+self.tokenizer.eos_token)
        return None,personas,None
    
    def respond(self, history,dialog_hx,personas,**kwargs):
        """
        Generates Response to User Input.

        Parameters:
        history: Chat History
        
        Returns:
        history: history with response appended
        audio: audio of the spoken response
        """

        person = self.tokenizer.encode(''.join(['<|p2|>'] + personas + ['<|sep|>'] + ['<|start|>']))
        user_inp= self.tokenizer.encode(history[-1][0]+self.tokenizer.eos_token)
        dialog_hx.append(user_inp)
        bot_input_ids = to_var([person + flatten(dialog_hx)]).long()
        with torch.no_grad():

            full_msg = self.model.generate(bot_input_ids,
                                        repetition_penalty=1.4,
                                        top_k = 10,
                                        top_p = 0.92,
                                        max_new_tokens = 256,
                                        num_beams=2,
                                        pad_token_id = self.tokenizer.eos_token_id)
        

        response = to_data(full_msg.detach()[0])[bot_input_ids.shape[-1]:]
        dialog_hx.append(response)
        history[-1][1] = self.tokenizer.decode(response, skip_special_tokens=True)
        self.speak(history[-1][1])
        return history, "out.mp3",dialog_hx
    
    def talk(self, audio, history,dialog_hx,personas,text):
        if audio is not None:
            history, _ = self.listen(audio, history)
        else:
            history.append([text,None])
        history, audio,dialog_hx = self.respond(history,dialog_hx,personas)
        return history, None, audio,dialog_hx,None

    def speak(self, text):
        """
        Speaks the given text using gTTS,
        Parameters:
        text: text to be spoken
        """
        tts = gTTS(text, lang='en')
        tts.save('out.mp3')

# Initialize AI Companion
bot = AI_Companion()
personas=[]
for i in ['I\'m a 19 year old girl','I study at IIT Indore','I am an easy-going and fun loving person','I love to swim','I am friendly, nice ,fun and kind','I am studious and get good grades']:
    response = i+ bot.tokenizer.eos_token
    personas.append(response)


# Create the Interface
with gr.Blocks() as demo:
    dialog_hx=gr.State([])
    personas=gr.State(personas)
    chatbot = gr.Chatbot([], elem_id = "chatbot").style(height = 300)
    audio = gr.Audio(source = "microphone", type = "filepath", label = "Input")
    msg = gr.Textbox()
    audio1 = gr.Audio(type = "filepath", label = "Output",elem_id="input")
    with gr.Row():
        b1 = gr.Button("Submit")
        b2 = gr.Button("Clear")
        b3=  gr.Button("Add Fact")
    b1.click(bot.talk, [audio, chatbot,dialog_hx,personas,msg], [chatbot, audio, audio1,dialog_hx,msg])
    msg.submit(append, [msg, chatbot,dialog_hx,personas], [chatbot, audio1, msg,dialog_hx])
    b2.click(clear, [] , [audio,chatbot,dialog_hx])
    b3.click(bot.add_fact, [audio,personas,msg], [audio,personas,msg])
demo.launch(share=True)