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

MODEL = "flax-community/t5-base-dutch-demo"
# MODEL: str = "."

DEFAULT_TEXT: str = """De Japanse onderzoekers gebruikten geavanceerde glasvezeltechnologie om de hoge snelheid te halen. Hiervoor werd onder meer gebruikgemaakt van zeldzame mineralen en een zogeheten microkam, waarbij laserstralen in bundels worden afgeschoten.
Het NICT simuleerde de data op zo'n manier dat de downloadsnelheid over een gebied van 3.000 kilometer kon worden gemeten. Over die afstand werd door de onderzoekers geen vertraging in de snelheid of verzwakking in het signaal gemeten.
Met een snelheid van 319 Tb/s zouden 57.000 films per seconde kunnen worden gedownload of de gehele Spotify-bibliotheek in drie seconden, schrijft Yahoo News.
Voorlopig lijkt de techniek overigens niet meer te zijn dan een experiment, vanwege de hoge kosten die ermee zijn gemoeid. In Nederland zijn providers bezig met het aanleggen van internet met een hoge downloadsnelheid. Deze glasvezelnetwerken bereiken snelheden tot 1 gigabit per seconde.
"""

generator_kwargs_beam = {
    "max_length": 175,
    "min_length": 30,
    "no_repeat_ngram_size": 2,
    "early_stopping": True,
    "num_beams": 5,
    "length_penalty": 1.5,
    "num_return_sequences": 1,
}

generator_kwargs_top_k = {
    "max_length": 175,
    "min_length": 30,
    "no_repeat_ngram_size": 2,
    "do_sample": True,
    "top_k": 60,
    "top_p": 0.95,
    "num_return_sequences": 1,
}


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

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

    def summarize(self, input_text, generate_kwargs) -> str:
        if not self.generator:
            self.load()
        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 demo the [T5 Dutch finetuned on CNN DM](https://huggingface.co/flax-community/t5-base-dutch-demo)
transformer model. Please refer to the [model page](https://huggingface.co/flax-community/t5-base-dutch-demo) for more information.
"""
    )
    st.sidebar.title("Mode:")

    if sampling_mode := st.sidebar.selectbox(
        "select a Mode", index=0, options=["Beam Search", "Top-k Sampling"]
    ):
        if sampling_mode == "Beam Search":
            params = generator_kwargs_beam
        else:
            params = generator_kwargs_top_k
        st.sidebar.json(params)

    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(input_text, params)

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


if __name__ == "__main__":
    main()