import streamlit as st
import os

from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM
from transformers import pipeline
from transformers import set_seed

debug = False

MODELS = [
    "flax-community/t5-base-dutch-demo",
    "yhavinga/t5-v1.1-base-dutch-cnn-test",
    "yhavinga/t5-v1.1-large-dutch-cnn-test",
]

DEFAULT_TEXT: str = """Het Van Goghmuseum in Amsterdam heeft vier kostbare prenten verworven van Mary Cassatt, de Amerikaanse impressionistische kunstenaar en tijdgenoot van Vincent van Gogh. Dat heeft het museum woensdagmiddag op een persconferentie bekendgemaakt.

Het gaat om drie grote kleurenetsen en een zwart-wit litho met voorstellingen van vrouwen. Voor deze prenten, die afkomstig zijn van een Amerikaanse verzamelaar, betaalde het museum ruim 1,4 miljoen euro. Drie grote fondsen en een aantal particulieren hebben samen de aankoopsom beschikbaar gesteld.

Mary Stevenson Cassatt (1844-1926) woonde en werkte lange tijd in Frankrijk. Ze staat met haar impressionistische schilderijen en tekeningen te boek als een van de vernieuwers van de Parijse kunstwereld in de late negentiende eeuw. Het Van Goghmuseum rekent haar prenten „tot het mooiste wat op grafisch gebied in het fin de siècle is geproduceerd”.

De drie aangekochte kleurenetsen – Het doorpassen, De brief en Badende vrouw – komen uit een serie van tien waarmee Cassatt haar naam als (prent)kunstenaar definitief vestigde. Ze maakte de etsen na een bezoek in 1890 aan een tentoonstelling van Japanse prenten in Parijs. Over die expositie schreef de Amerikaanse aan haar vriendin Berthe Morisot, een andere vrouwelijke impressionist: „We kunnen de Japanse prenten in de Beaux-Arts gaan bekijken. Echt, die mag je niet missen. Als je kleurenprenten wilt maken, is er niets mooiers voorstelbaar. Ik droom ervan en denk nergens anders meer aan dan aan kleur op koper.”"""


class TextSummarizer:
    def __init__(self):
        self.tokenizer = None
        self.model = None
        self.generator = None
        self.model_loaded = None
        set_seed(42)

    def load(self, model_name):
        os.environ["TOKENIZERS_PARALLELISM"] = "false"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        self.generator = pipeline(
            "text2text-generation", model=self.model, tokenizer=self.tokenizer
        )
        self.model_loaded = model_name

    def summarize(self, model_name, input_text, generate_kwargs) -> str:
        if not self.generator or self.model_loaded != model_name:
            with st.spinner(f"Loading model {model_name}..."):
                self.load(model_name)
        return self.generator(
            input_text, return_tensors=False, return_text=True, **generate_kwargs
        )[0].get("generated_text")


@st.cache(allow_output_mutation=True)
def instantiate_generator():
    summarizer = TextSummarizer()
    return summarizer


def main():
    st.set_page_config(  # Alternate names: setup_page, page, layout
        page_title="Netherformer",  # String or None. Strings get appended with "• Streamlit".
        layout="wide",  # Can be "centered" or "wide". In the future also "dashboard", etc.
        initial_sidebar_state="expanded",  # Can be "auto", "expanded", "collapsed"
        page_icon="📰",  # String, anything supported by st.image, or None.
    )

    with open("style.css") as f:
        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

    generator = instantiate_generator()

    st.markdown(
        """
        <style>
        [data-testid="stSidebar"][aria-expanded="true"] > div:first-child {
            width: 500px;
        }
        [data-testid="stSidebar"][aria-expanded="false"] > div:first-child {
            width: 500px;
            margin-left: -500px;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    st.sidebar.image("NewsCovid-19-512.png", width=200)
    st.sidebar.markdown(
        """# Netherformer
* Create summaries of Dutch news stories.
* Copy paste any Dutch news text and press the Generate Summary botton.
* This is a demo of several Dutch T5 models finetuned for summarization:
    * [T5 Base Dutch finetuned on CNN DM and XSUM](https://huggingface.co/flax-community/t5-base-dutch-demo) (Jul '21)
    * [T5 v1.1 Base Dutch finetuned on CNN DM](https://huggingface.co/yhavinga/t5-v1.1-base-dutch-cnn-test) (Jan '22)
    * [T5 v1.1 Large Dutch finetuned on CNN DM](https://huggingface.co/yhavinga/t5-v1.1-large-dutch-cnn-test) (Jan '22)
"""
    )
    st.sidebar.title("Parameters:")

    MODEL = st.sidebar.selectbox("Choose model", index=1, options=MODELS)

    min_length = st.sidebar.number_input(
        "Min length", min_value=10, max_value=150, value=75
    )
    max_length = st.sidebar.number_input(
        "Max length", min_value=50, max_value=250, value=142
    )
    no_repeat_ngram_size = st.sidebar.number_input(
        "No repeat NGram size", min_value=1, max_value=5, value=3
    )

    if sampling_mode := st.sidebar.selectbox(
        "select a Mode", index=0, options=["Beam Search", "Top-k Sampling"]
    ):
        if sampling_mode == "Beam Search":
            num_beams = st.sidebar.number_input(
                "Num beams", min_value=1, max_value=10, value=10
            )
            length_penalty = st.sidebar.number_input(
                "Length penalty", min_value=0.0, max_value=5.0, value=1.5, step=0.1
            )
            params = {
                "min_length": min_length,
                "max_length": max_length,
                "no_repeat_ngram_size": no_repeat_ngram_size,
                "num_beams": num_beams,
                "early_stopping": True,
                "length_penalty": length_penalty,
                "num_return_sequences": 1,
            }
        else:
            top_k = st.sidebar.number_input(
                "Top K", min_value=0, max_value=100, value=40
            )
            top_p = st.sidebar.number_input(
                "Top P", min_value=0.0, max_value=1.0, value=0.95, step=0.05
            )
            temperature = st.sidebar.number_input(
                "Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.05
            )
            params = {
                "min_length": min_length,
                "max_length": max_length,
                "no_repeat_ngram_size": no_repeat_ngram_size,
                "do_sample": True,
                "top_k": top_k,
                "top_p": top_p,
                "temperature": temperature,
                "num_return_sequences": 1,
            }

    st.sidebar.markdown(
        """For an explanation of the parameters, please to the [Huggingface blog post about text generation](https://huggingface.co/blog/how-to-generate)
        and the [Huggingface text generation interface doc](https://huggingface.co/transformers/main_classes/model.html?highlight=generate#transformers.generation_utils.GenerationMixin.generate).
        """
    )

    input_text = st.text_area("Enter a Dutch news text", DEFAULT_TEXT, height=500)

    if st.button("Generate summary"):

        with st.spinner("Generating summary ..."):
            if debug:
                from time import sleep

                sleep(2)
                response = """Een bedrijventerrein aan de Prins Willem Alexanderhaven in Roermond heeft de komende dagen weg te blijven uit de regio. De burgemeester van Roermond roept alle ramptoeristen er opnieuw toe op meer dan 10.00 uur weg."""
            else:
                response = generator.summarize(MODEL, input_text, params)

            st.header("Summary:")
            st.markdown(response)


if __name__ == "__main__":
    main()