import gradio as gr
import torch
from torch import nn
import lightning.pytorch as pl
from torch.nn import functional as F
from utils import GPTLM

newmodel = GPTLM.load_from_checkpoint('shakespeare_gpt.pth')

chars = ['\n', ' ', '!', '$', '&', "'", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
vocab_size = len(chars)
# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }

encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string


def generate_dialogue(character_dropdown, seed_slider):

  
  if character_dropdown == "NONE":
    context = torch.zeros((1, 1), dtype=torch.long)
    return decode(newmodel.model.generate(context, max_new_tokens=100)[0].tolist())
  else:
    context = torch.tensor([encode(character_dropdown)], dtype=torch.long)
    return decode(newmodel.model.generate(context, max_new_tokens=100)[0].tolist())
      

HTML_TEMPLATE = """    
<style>
    
    #app-header {
        text-align: center;
        background: rgba(255, 255, 255, 0.3); /* Semi-transparent white */
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        position: relative; /* To position the artifacts */
    }
    #app-header h1 {
        color: #FF0000;
        font-size: 2em;
        margin-bottom: 10px;
    }
    .concept {
        position: relative;
        transition: transform 0.3s;
    }
    .concept:hover {
        transform: scale(1.1);
    }
    .concept img {
        width: 100px;
        border-radius: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .concept-description {
        position: absolute;
        bottom: -30px;
        left: 50%;
        transform: translateX(-50%);
        background-color: #4CAF50;
        color: white;
        padding: 5px 10px;
        border-radius: 5px;
        opacity: 0;
        transition: opacity 0.3s;
    }
    .concept:hover .concept-description {
        opacity: 1;
    }
    /* Artifacts */
    
</style>
<div id="app-header">
    <!-- Artifacts -->
    <div class="artifact large"></div>
    <div class="artifact large"></div>
    <div class="artifact large"></div>
    <div class="artifact large"></div>
    <!-- Content -->
    <h1>SHAKESPEARE  DIALOGUE  GENERATOR</h1>
    <p>Generate dialogue for Shakespearean character by selecting character from dropdown.</p>
"""

with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background: url('file=https://github.com/santule/ERA/assets/20509836/b6b4031a-265d-43f6-bd59-813097c0022b')}") as interface:
    gr.HTML(value=HTML_TEMPLATE, show_label=False)
    with gr.Row():
        character_dropdown = gr.Dropdown(
            label="Select a Character",
            choices=["NONE","ROMEO","JULIET","MENENIUS","ANTONIO"],
            value='Dream'
        )
        seed_slider = gr.Slider(
            label="Random Seed",
            minimum=0,
            maximum=1000,
            step=1,
            value=42
        )
        inputs = [character_dropdown, seed_slider]

    with gr.Row():
        outputs = gr.Textbox(
            label="Generated Dialogue"
        )

    with gr.Row():
        button = gr.Button("Generate Dialogue")
        button.click(generate_dialogue, inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    interface.launch(enable_queue=True)