# Import required libraries
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
import gradio as gr

# Initialize AI pipelines
sentiment = pipeline("sentiment-analysis", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
summarize = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
# Initialize ASR with a specific model
asr = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")

# Load chatbot model and tokenizer
chatbot_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
chatbot_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# generating a response from the model

def generate_response(message, history):
    # Encode the input message
    input_ids = chatbot_tokenizer.encode(message + chatbot_tokenizer.eos_token, return_tensors="pt")

    # Generate response
    response_ids = chatbot_model.generate(
        input_ids,
        max_length=1000,
        pad_token_id=chatbot_tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
        do_sample=True,
        top_k=100,
        top_p=0.7,
        temperature=0.8
    )

    # Decode the response
    response = chatbot_tokenizer.decode(response_ids[0], skip_special_tokens=True)
    return response

# Functions for each task
def get_sentiment(text):
  sent = sentiment(text)[0]['label']
  score = sentiment(text)[0]['score']
  return sent, score

def transcribe_audio(speech):
    text = asr(speech)["text"]
    return text

def summ(text):
  text = summarize(text)
  return text

# Chatbot function
def chat_response(message, history):
    # Encode the input text
    inputs = tokenizer(message, return_tensors="pt")
    
    # Generate response
    outputs = model.generate(
        inputs.input_ids,
        max_length=100,
        num_return_sequences=1,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # Decode and return response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Create Gradio interface
interface = gr.Blocks()

# Building the block interface with tabs
with interface:
    gr.Markdown("# TrailTrek Co")
    with gr.Tabs():
        with gr.TabItem("Sentiment Analysis"):
            with gr.Row():
                text_input = gr.Textbox(label="Enter the review")
            with gr.Row():
              text_output = [gr.Textbox(label = "Sentiment :"), gr.Textbox(label = "Score :")]
            analyze_butt = gr.Button("Analyze")
        with gr.TabItem("Summarize"):
            with gr.Row():
                summ_input = gr.Textbox(label="Text to summarize")
                summ_output = gr.Textbox(label="Result")
            summ_button = gr.Button("Summarize")
        with gr.Tab("Chatbot"):
            gr.Markdown("Say hello :)")
            chatbot = gr.Chatbot(
                label="Chat History",
                height=400
            )
            msg = gr.Textbox(
                label="Type your message",
                placeholder="Type your message here...",
                show_label=False
            )
            clear = gr.Button("Clear")

            def user(user_message, history):
                return "", history + [[user_message, None]]

            def bot(history):
                user_message = history[-1][0]
                bot_message = generate_response(user_message, history)
                history[-1][1] = bot_message
                return history

            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)

        with gr.Tab("Speech Recognition"):
            gr.Markdown("Upload an audio file for transcription.")
            audio_input = gr.Audio(label="Upload Audio", type="filepath")
            transcription_output = gr.Textbox(label="Transcription", lines=3)
            transcribe_butt = gr.Button("Transcribe")

    analyze_butt.click(get_sentiment, inputs=text_input, outputs=text_output)
    summ_button.click(summarize, inputs=summ_input, outputs=summ_output)
    transcribe_butt.click(transcribe_audio, inputs=audio_input, outputs=transcription_output)

    # Launch the interface
    interface.launch()