###############################################################################
# Language Modeling on Wikitext-2
#
# This file generates new sentences sampled from the language model
#
###############################################################################

# stdlib
import argparse

# third party
import torch

import data  # isort:skip

parser = argparse.ArgumentParser(description="PyTorch Wikitext-2 Language Model")

# Model parameters.
parser.add_argument(
    "--data",
    type=str,
    default="./data/wikitext-2",
    help="location of the data corpus; default: \"./data/wikitext-2\""
)
parser.add_argument(
    "--checkpoint",
    type=str,
    default="./model.pt",
    help="model checkpoint to use; default: \"./model.pt\""
)
parser.add_argument(
    "--outf",
    type=str,
    default="generated.txt",
    help="output file for generated text; default: \"generated.txt\""
)
parser.add_argument(
    "--words",
    type=int,
    default="1000",
    help="number of words to generate; default: 1000"
)
parser.add_argument(
    "--seed",
    type=int,
    default=1111,
    help="random seed; default: 1111"
)
parser.add_argument(
    "--cuda",
    action="store_true",
    help="use CUDA"
)
parser.add_argument(
    "--temperature",
    type=float,
    default=1.0,
    help="temperature - higher will increase diversity; default: 1.0"
)
parser.add_argument(
    "--log-interval",
    type=int,
    default=100,
    help="reporting interval; default: 100")
args = parser.parse_args()

# Set the random seed manually for reproducibility.
torch.manual_seed(args.seed)
if torch.cuda.is_available():
    if not args.cuda:
        print("WARNING: You have a CUDA device, so you should probably run with --cuda")

device = torch.device("cuda" if args.cuda else "cpu")

if args.temperature < 1e-3:
    parser.error("--temperature has to be greater or equal 1e-3")

with open(args.checkpoint, "rb") as f:
    model = torch.load(f).to(device)
model.eval()

corpus = data.Corpus(args.data)
ntokens = len(corpus.dictionary)

is_transformer_model = (
    hasattr(model, "model_type") and model.model_type == "Transformer"
)
if not is_transformer_model:
    hidden = model.init_hidden(1)
input = torch.randint(ntokens, (1, 1), dtype=torch.long).to(device)

with open(args.outf, "w") as outf:
    with torch.no_grad():  # no tracking history
        for i in range(args.words):
            if is_transformer_model:
                output = model(input, False)
                word_weights = output[-1].squeeze().div(args.temperature).exp().cpu()
                word_idx = torch.multinomial(word_weights, 1)[0]
                word_tensor = torch.Tensor([[word_idx]]).long().to(device)
                input = torch.cat([input, word_tensor], 0)
            else:
                output, hidden = model(input, hidden)
                word_weights = output.squeeze().div(args.temperature).exp().cpu()
                word_idx = torch.multinomial(word_weights, 1)[0]
                input.fill_(word_idx)

            word = corpus.dictionary.idx2word[word_idx]

            outf.write(word + ("\n" if i % 20 == 19 else " "))

            if i % args.log_interval == 0:
                print(f"| Generated {i}/{args.words} words")
