# -*- coding: utf-8 -*-
"""MtGPT2

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1HMq9Cp_jhqc9HlUipLXi8SC1mHQhSgGn
"""

#@title Setup + Download Model
# Install and load dependencies
import os

os.system("python3 -m pip install pytorch-lightning==1.7.0 aitextgen")


# Hacky workaround because "_compare version" was moved and renamed into "lightning_utilities.core.imports.compare_version"
with open("/home/user/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/progress/rich_progress.py", "r") as file:
    lines = file.readlines()

    #replace line 20 "from torchmetrics.utilities.imports import _compare_version"
    #with
    #"from lightning_utilities.core.imports import compare_version as _compare_version"
    lines[19] = "from lightning_utilities.core.imports import compare_version as _compare_version\n"

with open("/home/user/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/progress/rich_progress.py", "w") as file:
    file.writelines(lines)

import sys
from jinja2 import Template
from aitextgen import aitextgen

try:
    from google.colab import files
except ImportError:
    pass

# Download and load the model. Set to_gpu=False if running on a CPU.
ai = aitextgen(model="minimaxir/magic-the-gathering", to_gpu=False)

# This template is similar to Scryfall card formatting
TEMPLATE = Template(
    """{{ c.name }}{% if c.manaCost %} {{ c.manaCost }}{% endif %}
{{ c.type }}
{{ c.text }}{% if c.power %}
{{ c.power }}/{{ c.toughness }}{% endif %}{% if c.loyalty %}
Loyalty: {{ c.loyalty }}{% endif %}"""
)

def render_card(card_dict):
    card = TEMPLATE.render(c=card_dict)
    if card_dict["name"]:
        card = card.replace("~", card_dict["name"])
    return card

def generate_cards(
    n_cards: int = 8,
    temperature: float = 0.75,
    name: str = "",
    manaCost: str = "",
    type: str = "",
    text: str = "",
    power: str = "",
    toughness: str = "",
    loyalty: str = "",
    do_replacements: bool = True
):
    #ensure n_cards is never 0 or negative
    n_cards = int(n_cards)
    if n_cards < 1:
        n_cards = 1


    #change manaCost from Format:
    # 2UG
    #to:
    #{2}{U}{G}

    manaCost_str = ""

    for char in manaCost:
        manaCost_str += "{"
        manaCost_str += char
        manaCost_str += "}"




    prompt_str = ""

    token_dict = {
        "<|name|>": name,
        "<|manaCost|>": manaCost_str,
        "<|type|>": type,
        "<|text|>": text,
        "<|power|>": power,
        "<|toughness|>": toughness,
        "<|loyalty|>": loyalty
    }

    # Convert the token_dict into a formatted prompt string
    for token, value in token_dict.items():
        if value:
            prompt_str += f"{token}{value}"

    # Generate the cards using the prompt string and other parameters
    cards = ai.generate(
        n=n_cards,
        schema=True,
        prompt=prompt_str,
        temperature=temperature,
        return_as_list=True
    )

    cards = list(map(render_card, cards))

    out_str = "\n=====\n".join(cards)

    replacements = {
        "{G}": "🌲",
        "{U}": "🌊",
        "{R}": "🔥",
        "{B}": "💀",
        "{W}": "☀️",
        "{T}": "↩️",
        #"{1}": "(1)",
        #"{2}": "(2)",
        #"{3}": "(3)",
        #"{4}": "(4)",
        #"{5}": "(5)",
        #"{6}": "(6)",
        #"{7}": "(7)",
        #"{8}": "(8)",
        #"{9}": "(9)",
    }

    if do_replacements:
        for key, value in replacements.items():
            out_str = out_str.replace(key, value)


    return out_str

# Commented out IPython magic to ensure Python compatibility.
# %pip install gradio
import gradio as gr

iface = gr.Interface(
    fn = generate_cards,
    inputs=[
        gr.Slider(minimum = 2, maximum=16, step=1, value=8),
        gr.Slider(minimum = 0.1, maximum=1.5, step=0.01, value=0.75),
        gr.Textbox(),
        gr.Textbox(),
        gr.Textbox(),
        gr.Textbox(),
        gr.Textbox(),
        gr.Textbox(),
        gr.Textbox(),
        gr.Checkbox()
    ],
    outputs=gr.Textbox(),
    title = "GPT-2 Powered MTG Card Generator",
    description = "Enter Manacost as '2UG' for 2 colorless + Blue + Green mana. \n\n Temperature is recomended between 0.4 and 0.9. Anything above 1 will lead to random Chaos and very low values will just be boring.",
    allow_flagging = "never"
)
iface.launch(show_api = True)