from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from datetime import datetime as dt
import streamlit as st
from streamlit_tags import st_tags
import beam_search
import top_sampling
from pprint import pprint
import json

with open("config.json") as f:
    cfg = json.loads(f.read())

st.set_page_config(layout="wide")


@st.cache(allow_output_mutation=True)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"])
    model = AutoModelForSeq2SeqLM.from_pretrained(cfg["model_name_or_path"])
    generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
    return generator, tokenizer


def sampling_changed(obj):
    print(obj)


with st.spinner('Loading model...'):
    generator, tokenizer = load_model()
# st.image("images/chef-transformer.png", width=400)
st.header("Chef Transformer 👩‍🍳 / 👨‍🍳")
st.markdown(
    "This demo uses [T5 trained on RecipeNLG](https://huggingface.co/flax-community/t5-recipe-generation) "
    "to generate recipe from a given set of ingredients"
)
img = st.sidebar.image("images/chef-transformer-transparent.png", width=310)
add_text_sidebar = st.sidebar.title("Popular Recipes")
add_text_sidebar = st.sidebar.text("Recipe preset(example#1)")
add_text_sidebar = st.sidebar.text("Recipe preset(example#2)")

add_text_sidebar = st.sidebar.title("Mode")
sampling_mode = st.sidebar.selectbox("Select a Mode", index=0, options=["Top Sampling", "Beam Search"])

original_keywords = st.multiselect(
    "Choose ingredients",
    cfg["first_100"],
    ["parmesan cheese", "fresh oregano", "basil", "whole wheat flour"]
)

st.write("Add custom ingredients here:")
custom_keywords = st_tags(
    label="",
    text='Press enter to add more',
    value=['salt'],
    suggestions=["z"],
    maxtags=15,
    key='tags')

# def custom_keywords_on_change():
#     pass


# custom_keywords = st.text_input(
#     'Add custom ingredients here (separated by `,`): ',
#     ", ".join(["salt", "pepper"]),
#     key="custom_keywords",
#     # on_change=custom_keywords_on_change,
#     max_chars=1000)
# custom_keywords = list(set([x.strip() for x in custom_keywords.strip().split(',') if len(x.strip()) > 0]))

all_ingredients = []
all_ingredients.extend(original_keywords)
all_ingredients.extend(custom_keywords)
all_ingredients = ", ".join(all_ingredients)
st.markdown("**Generate recipe for:** " + all_ingredients)

submit = st.button('Get Recipe!')
if submit:
    with st.spinner('Generating recipe...'):
        if sampling_mode == "Beam Search":
            generated = generator(
                all_ingredients,
                return_tensors=True,
                return_text=False,
                **beam_search.generate_kwargs)
            outputs = beam_search.post_generator(generated, tokenizer)
        elif sampling_mode == "Top Sampling":
            generated = generator(
                all_ingredients,
                return_tensors=True,
                return_text=False,
                **top_sampling.generate_kwargs)
            outputs = top_sampling.post_generator(generated, tokenizer)
    output = outputs[0]
    output['title'] = " ".join([w.capitalize() for w in output['title'].split()])
    markdown_output = ""
    markdown_output += f"## {output['title']}\
"
    markdown_output += f"#### Ingredients:\
"
    for o in output["ingredients"]:
        markdown_output += f"- {o}\
"
    markdown_output += f"#### Directions:\
"
    for o in output["directions"]:
        markdown_output += f"- {o}\
"
    st.markdown(markdown_output)