# test in gradio application
from assets.CharacterDataset import CharacterDataset # change in native
from assets.PredictionNetwork import PredictionNetwork, generate_text # change in native
from timeit import default_timer as timer
import numpy as np
import torch.nn.functional as F 
import torch 
import gradio as gr

# samples path
samples_path = "advs.txt" # change in native
# pretrained weights path
pretrained_weights_path = "sherlock_holmes_lstm_text_generator_state_dict.pth" # change in native

# hyperparameters
vocab_size = 50
window_size = 40
embedding_dim = 512
linear_units = 64
lstm_units = 32
n_layers = 64
max_norm = 2

# get the text from the textfile
with open(samples_path,"r") as f:
  text = "\n".join(f.readlines())

# instance the dataset
dataset = CharacterDataset(text,window_size=window_size,vocab_size=vocab_size)

# instance the model
model = PredictionNetwork(
  vocab_size = vocab_size,
  embedding_dimention = embedding_dim, 
  linear_units = linear_units,
  lstm_units = lstm_units, 
  max_norm = max_norm,
  n_layers = n_layers
)

# load the pretrained weights 
model.load_state_dict(torch.load(f=pretrained_weights_path,map_location=torch.device('cpu')))

# text generation function
def generate_text(n_chars,initial_text):
  res = initial_text
  model.eval()
  h,c = None,None
  begin_time = timer()
  with torch.inference_mode():
    for _ in range(int(n_chars)):
      prev_chars = initial_text if res == initial_text else res[-1]
      features = torch.LongTensor([[dataset.ch2idx[c] for c in prev_chars]])
      logits, h,c = model(features,h,c)
      probs = F.softmax(logits[0],dim = 0).to("cpu").detach().numpy()
      new_ch = np.random.choice(dataset.vocabulary,p = probs)
      res += new_ch
  end_time = timer()
  total_time = end_time - begin_time
  return res,total_time

# create gradio application
title = "Sherlock Holmes text generator"
description = "An LSTM model trained to generate text using the anthology: [The Adventures Of Sherlock Holmes](https://sherlock-holm.es/stories/plain-text/advs.txt)."
article = "The model was created using Pytorch. It incorporates embedding matrices to code the features of the vocabulary extracted from the book, and it is a partial improvement performance wise from my last project based on [Shakespheare's plays](https://huggingface.co/spaces/Q-b1t/shakespeare_text_generation). I plan on developing another version that incorporates a full transformer in the course of the next few weeks."

initial_text = "In the year 1878 I took my degree of Doctor of Medicine of the University of London,"

demo = gr.Interface(
    fn = generate_text,
    inputs = [gr.Number(value = 200,label = "Sequence Length",info = "Length of the sample sequence you wish to generate."),gr.TextArea(lines = 5,label="Initial Text",value = initial_text)],
    outputs = [gr.TextArea(lines = 5,label="Sequence Output"),gr.Number(label = "Execution Time (seconds)")],    title = title,
    description = description,
    article = article
)

demo.launch()