import torch
from mapper import load_mapping_data
import random

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
map_vocab = load_mapping_data('mapping_vocab.json')

word2idx = map_vocab['word2idx']
idx2word = map_vocab['idx2word']

# Generating the next word
def predict_next_word(model, context, n_contex = 5):
    with torch.no_grad():
        context = [word2idx[word] for word in context if word in word2idx]
        if len(context) < n_contex:
            context = [0] * (5 - len(context)) + context
        context_tensor = torch.tensor([context], dtype=torch.long).to(device)
        output = model(context_tensor).squeeze()
        probabilities = torch.softmax(output, dim=0)
        weight, top5 = torch.topk(probabilities, k=5, dim=0)
        
        predicted_idx = random.choices(top5, weight)[0].item()

    return idx2word[str(predicted_idx)]

# generate sentence
def generate_sentence(model, context, max_length=10, n_context=5):
    with torch.no_grad():
        sentence = context.copy()
        for _ in range(max_length):
            next_word = predict_next_word(model, sentence[-n_context:])
            if next_word in ['<s>', '</s>']:
                context.append(' ')
                sentence.append(next_word)
            else:
                sentence.append(next_word)
                context.append(next_word)

    return ''.join([cont if cont not in ['<s>', '</s>'] else ' ' for cont in context])

if __name__ == '__main__':
    from load_model import *
    nwp = load_model('models/nwp_scratch_model.pth')
    sn = generate_sentence(nwp, [], max_length=50)
    print(sn)