from dotenv import load_dotenv
import gradio as gr
import random

from utils.model import Model
from utils.data import dataset

import gc
import torch

load_dotenv()

__model_on_gpu__ = Model.__model_list__[0]
model = {model_name: Model(model_name) for model_name in Model.__model_list__}

random_label = '🔀 Random dialogue from dataset'
examples = {
    "example 1": """Boston's injury reporting for Kristaps Porziņģis has been fairly coy. He missed Game 3, but his coach told reporters just before Game 4 that was technically available, but with a catch.
Joe Mazzulla said Porziņģis would "only be used in specific instances, if necessary." That sounds like the team doesn't want to risk further injury to his dislocated Posterior Tibialis (or some other body part, due to overcompensation for the ankle), unless it's in a desperate situation.
Being up 3-1, with Game 5 at home, doesn't qualify as desperate. So, expect the Celtics to continue slow-playing KP's return.
It'd obviously be nice for Boston to have his rim protection and jump shooting back. It was missed in the Game 4 blowout, but the Celtics have also demonstrated they can win without the big man throughout this campaign.
On top of winning Game 3 of this series, Boston is plus-10.9 points per 100 possessions when Porziņģis has been off the floor this regular and postseason.""",
    
    "example 2": """Prior to the Finals, we predicted that Dereck Lively II's minutes would swell over the course of the series, and that's starting to play out.
He averaged 18.8 minutes in Games 1 and 2 and was up to 26.2 in Games 3 and 4. That's with the regulars being pulled long before the final buzzer in Friday's game, too.
Expect the rookie's playing time to continue to climb in Game 5. It seems increasingly clear that coach Jason Kidd trusts him over the rest of Dallas' bigs, and it's not hard to see why.
Lively has been absolutely relentless on the offensive glass all postseason. He makes solid decisions as a passer when his rolls don't immediately lead to dunks. And he's not a liability when caught defending guards or wings outside.
All of that has led to postseason averages of 8.2 points, 7.6 rebounds, 1.4 assists and 1.0 blocks in just 21.9 minutes, as well as a double-double in 22 minutes of Game 4.
Back in Boston, Kidd is going to rely on Lively even more. He'll play close to 30 minutes and reach double-figures in both scoring and rebounding again.""",

    random_label: ""
}

def generate_answer(sources, model, model_name, prompt):
    content = prompt + '\n' + sources + '\n\n'
    global __model_on_gpu__

    if __model_on_gpu__ != model_name:
        if not __model_on_gpu__:
            model[__model_on_gpu__].cpu()
            gc.collect()
            torch.cuda.empty_cache()

        model[model_name].cuda()
        __model_on_gpu__ = model_name

    answer = model[model_name].gen(content)

    return answer

def process_input(input_text, model_selection, prompt):
    if input_text:
        response = generate_answer(input_text, model, model_selection, prompt)
        return f"## Original Article:\n\n{input_text}\n\n## Summarization:\n\n{response}"
    else:
        return "Please fill the input to generate outputs."

def update_input(example):
    if example == random_label:
        return random.choice(dataset)['dialogue']
    return examples[example]

def create_summarization_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## This is a playground to test summarizations")
        
        with gr.Row():
            example_dropdown = gr.Dropdown(choices=list(examples.keys()), label="Choose an example", value=random_label)
            model_dropdown = gr.Dropdown(choices=Model.__model_list__, label="Choose a model", value=Model.__model_list__[0])
        
        Template_text = gr.Textbox(value="""Summarize the following dialogue""", label='Input Prompting Template', lines=8, placeholder='Input your prompts')

        input_text = gr.Textbox(label="Input Text", lines=10, placeholder="Enter text here...")
        submit_button = gr.Button("✨ Submit ✨")
        output = gr.Markdown()

        example_dropdown.change(update_input, inputs=[example_dropdown], outputs=[input_text])
        submit_button.click(process_input, inputs=[input_text, model_dropdown, Template_text], outputs=[output])

    return demo

if __name__ == "__main__":
    demo = create_summarization_interface()
    demo.launch()