import gradio as gr
from gradio.themes.base import Base
import tensorflow as tf
import numpy as np
import pickle
import os
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Load the trained model
model = tf.keras.models.load_model("RomanUrduPoetryModel.h5")

# Load tokenizer from saved file
with open("UrduRomanPoetryTokenizer.pkl", "rb") as f:
    tokenizer = pickle.load(f)

# Define max_sequence_length (Same as in Jupyter)
max_sequence_length = 50

# Optimized Poetry Generation Function with Download Feature
def generate_poetry(seed_text, next_words=20, temperature=0.8):
    token_list = tokenizer.texts_to_sequences([seed_text])[0]
    
    for _ in range(next_words):
        token_list = pad_sequences([token_list], maxlen=max_sequence_length - 1, padding="pre")

        predictions = model.predict(token_list, verbose=0)[0]  # Faster prediction
        predictions = np.log(predictions + 1e-7) / temperature
        exp_preds = np.exp(predictions)
        probabilities = exp_preds / np.sum(exp_preds)

        predicted = np.random.choice(len(probabilities), p=probabilities)

        output_word = next((word for word, index in tokenizer.word_index.items() if index == predicted), "")
        if not output_word:
            break  # Stop if no valid word is found

        seed_text += " " + output_word
        token_list = np.append(token_list, predicted)  # Update token list

    # Save poetry to a text file for downloading
    file_path = "generated_poetry.txt"
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(seed_text)

    return seed_text, file_path

# Customize UI
class Seafoam(Base):
    pass
seafoam = Seafoam(font=gr.themes.GoogleFont("Plus Jakarta Sans"))

style ="""
    .gradio-audio {
        border-radius: 15px !important;
    }
    .gradio-button {
        background: #007bff;
        color: white;
        font-weight: bold;
        border: none;
        border-radius: 20px;
    }
    .gradio-button:hover {
        background: #0056b3;
    }
    .gradio-dropdown {
        background: transparent;
    }
    .gradio-button2 {
        background: transparent;
        border:1.5px solid var(--input-border-color);
        font-weight: bold;
        border-radius: 20px;
    }
    .gradio-button2:hover {
        background: var(--input-border-color);
    }
    label.container.show_textbox_border.svelte-173056l textarea.svelte-173056l {
        background:transparent;
        border-radius: 20px;
    }
    div.svelte-633qhp {
        border-radius: 15px;
        overflow-y: hidden;
    }
    span.svelte-1gfkn6j {
        padding-left: 20px,
        font-size:16px;
        font-weight: bold;
    }
    .gradio-container.gradio-container-5-16-0 .contain span.svelte-1gfkn6j {
        padding-left: 12px;
    }
    .icon-button-wrapper.hide-top-corner.svelte-1jx2rq3 {
        border-radius: 20px;
        margin: 5px 6.09px 0px 0px;
        padding: 6px 5.5px 5px 5.5px;
    }
"""

# Gradio Interface with Better UI
with gr.Blocks(theme=seafoam, css=style) as app:
    gr.Markdown("# Shaayer", elem_classes=["gradio-title"])

    with gr.Row():
        seed_input = gr.Textbox(label="Poetry Seed", placeholder="Enter your poetry seed here ...")
        num_words = gr.Slider(10, 50, step=5, label="Number of Words", value=20)
        temp = gr.Slider(0.2, 1.0, step=0.1, label="Creativity (Temperature)", value=0.5)

    poetry_output = gr.Textbox(label="Generated Poetry", elem_classes=["gradio-textbox"])
    download_btn = gr.DownloadButton("Download Generated Poetry", value="generated_poetry.txt", visible=False, elem_classes=["gradio-button2"])

    generate_button = gr.Button("Generate", variant="primary", elem_classes=["gradio-button"])

    def generate_download_links(seed_input, num_words, temp):
        poetry_output, text_file = generate_poetry(seed_input, num_words, temp)
        return poetry_output, gr.update(value=text_file, visible=True)

    generate_button.click(generate_download_links, inputs=[seed_input, num_words, temp], outputs=[poetry_output, download_btn])

# Launch Gradio App with Better Layout
app.launch(share=True, inbrowser=True)