import torch
import torch.nn as nn
import json
from tqdm import tqdm, trange

# Model parameters
parameters = json.loads(open("parameter.json").read())
model_path = parameters["model_path"]

# Define the simple RNN model
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleRNN, self).__init__()
        self.hidden_size = hidden_size
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden):
        x = torch.nn.functional.one_hot(x, num_classes=input_size).float()
        out, hidden = self.rnn(x.unsqueeze(0), hidden)
        out = self.fc(out[:, -1, :])  # Take last time step's output
        return out, hidden

model = torch.load(model_path, weights_only=False)
with open("vocab.json", "r") as f:
    chars = json.loads(f.read())
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}
print("Loaded pre-trained model.")

input_size = len(chars)
hidden_size = parameters["hidden_size"]
output_size = len(chars)

# Text generation function
def generate_text(start_text, length):
    model.eval()
    hidden = torch.zeros(1, 1, hidden_size)
    input_seq = torch.tensor([char_to_idx[ch] for ch in start_text])

    generated_text = start_text
    for _ in trange(length):
        output, hidden = model(input_seq, hidden)
        predicted_idx = output.argmax().item()
        generated_text += idx_to_char[predicted_idx]
        input_seq = torch.cat((input_seq[1:], torch.tensor([predicted_idx])))

    return generated_text

# Generate some text
while True:
    prompt = input("Ask LLM: ")
    length = int(input("Length of text: "))
    print("LLM Output: ", generate_text(prompt, length))