"""
DnD Room Generator App
Only uses sampling, not beams or greedy

https://huggingface.co/blog/how-to-generate
"""
import gradio as gr
import transformers as tr

SPECIAL_TOKENS = {
    'eos_token': '<|EOS|>',
    'bos_token': '<|endoftext|>',
    'pad_token': '<pad>',
    'sep_token': '<|body|>'
}
MPATH = "./mdl_roomgen2"
MODEL = tr.GPT2LMHeadModel.from_pretrained(MPATH)

# ToDo: Will save tokenizer next time so can replace this with a load
TOK = tr.GPT2Tokenizer.from_pretrained("gpt2")
TOK.add_special_tokens(SPECIAL_TOKENS)


def generate_room(room_name, max_length, top_k, temperature, top_p):
    """
    Uses pretrained model to generate text for a dungeon room

    Args:
        room_name:
        method:

    Returns:

    """
    prompt = " ".join(
        [
            SPECIAL_TOKENS["bos_token"],
            room_name,
            SPECIAL_TOKENS["sep_token"]
        ]
    )
    ids = TOK.encode(prompt, return_tensors="pt")
    # Sample
    output = MODEL.generate(
        ids,
        max_length=max_length,
        do_sample=True,
        top_k=top_k,
        temperature=temperature,
        top_p=top_p
    )
    output = TOK.decode(output[0][ids.shape[1]:], clean_up_tokenization_spaces=True).replace("  ", " ")
    return output


if __name__ == "__main__":
    iface = gr.Interface(
        title="RPG Room Generator",
        fn=generate_room,
        inputs=[
            gr.inputs.Textbox(lines=1, label="Room Name"),
            gr.inputs.Slider(minimum=50, maximum=250, default=175, label="Length"),
            gr.inputs.Slider(minimum=0, maximum=100, default=50, label="Top K"),
            gr.inputs.Slider(minimum=0, maximum=1.0, default=0.7, step=0.1, label="Temperature"),
            gr.inputs.Slider(minimum=0, maximum=1.0, default=0.5, step=0.1, label="Top P")
        ],
        outputs="text",
        layout="horizontal",
        allow_flagging=None,
        theme="dark"
    )
    app, local_url, share_url = iface.launch(share=True)