import gradio as gr
import torch
import torch.nn as nn
import json
import os


# Define parameters and directory paths
SEQUENCE_LENGTH = 40
LEARNING_RATE = 0.001
NUM_EPOCHS = 10
GRAD_CLIP = 5.0  # Gradient clipping to prevent exploding gradients
CHECKPOINT_FOLDER = "checkpoints"
os.makedirs(CHECKPOINT_FOLDER, exist_ok=True)
model_name = "Talk-model-1"

# Suppose we have a saved ‘model_config.json’ and a ‘vocab.txt’ in the CHECKPOINT_FOLDER
CONFIG_PATH = os.path.join(CHECKPOINT_FOLDER, f"{model_name}_config.json")
VOCAB_PATH = os.path.join(CHECKPOINT_FOLDER, f"{model_name}_vocab.txt")

# Load model configuration if available
if os.path.exists(CONFIG_PATH):
    with open(CONFIG_PATH, "r") as config_file:
        model_config = json.load(config_file)
    VOCAB_SIZE = model_config["vocab_size"]
    HIDDEN_SIZE = model_config["hidden_size"]
    NUM_LAYERS = model_config["num_layers"]
    BATCH_SIZE = model_config["batch_size"]

# Load vocabulary if available
if os.path.exists(VOCAB_PATH):
    with open(VOCAB_PATH, "r") as vocab_file:
        vocab = vocab_file.read().rstrip()
    VOCAB_SIZE = len(vocab)  # Update VOCAB_SIZE based on loaded vocabulary

# Define the sophisticated RNN model architecture
class AdvancedRNN(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super(AdvancedRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.rnn = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden):
        x = self.embedding(x)
        out, hidden = self.rnn(x, hidden)
        out = self.fc(out.reshape(out.size(0) * out.size(1), out.size(2)))
        return out, hidden

    def init_hidden(self, batch_size):
        weight = next(self.parameters())
        hidden = (weight.new_zeros(self.num_layers, batch_size, self.hidden_size),
                  weight.new_zeros(self.num_layers, batch_size, self.hidden_size))
        return hidden

# Initialize model, loss function, and optimizer
model = AdvancedRNN(VOCAB_SIZE, HIDDEN_SIZE, NUM_LAYERS)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Predict the next character
def predict(model, char, hidden=None, temperature=0.1):
    # Convert to ASCII and then tensor
    input_tensor = torch.tensor([ord(char)], dtype=torch.long).unsqueeze(0)
    if hidden is None:
        hidden = model.init_hidden(1)
    output, hidden = model(input_tensor, hidden)
    output_dist = output.view(-1).div(temperature).exp()
    top_i = torch.multinomial(output_dist, 1)[0]
    predicted_char = chr(top_i)
    return predicted_char, hidden

def generate_text_interface(start_text, length=200):
    model.eval()  # Use model instead of trained_model
    
    # Check if length is acceptable
    if length < len(start_text) or length <= 0:
        return "Length must be greater than the length of the starting text and greater than 0."
    
    # Preparing the hidden state
    hidden = model.init_hidden(1)
    for char in start_text[:-1]:
        _, hidden = predict(model=model, char=char, hidden=hidden)
    
    # Start generating text from the last character of the provided text
    input_char = start_text[-1]
    generated_text = start_text
    
    # Generate text
    for i in range(length - len(start_text)):
        input_char, hidden = predict(model=model, char=input_char, hidden=hidden)
        generated_text += input_char
    
    return generated_text

# Build and launch the Gradio web app
interface = gr.Interface(
    fn=generate_text_interface,
    inputs=[
        gr.components.Textbox(label="Starting Text", placeholder="The universe "),
        gr.components.Slider(minimum=1, maximum=1000, value=200, label="Length of Generated Text")
    ],
    outputs=gr.components.Textbox(label="Generated Text"),
    title="Text Generation App",
    description="Generate text based on a starting sequence."
)


# You can choose not to train the model every time you run this script by checking for the existence
# of a saved model. For this example, we’ll assume you’ve trained and saved it as ‘model_final.pth’.
# Load the saved model weights before launching the app if it’s already trained.
CHECKPOINT_PATH = os.path.join(CHECKPOINT_FOLDER, f"{model_name}.pth")
if os.path.exists(CHECKPOINT_PATH):
    model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=torch.device("cpu")))
    model.eval()  # Set the model to evaluation mode

# Launch the app
interface.launch(share=True)