import logging
import time
from pathlib import Path

import gradio as gr
import nltk
from cleantext import clean

from summarize import load_model_and_tokenizer, summarize_via_tokenbatches
from utils import load_example_filenames, truncate_word_count

_here = Path(__file__).parent

nltk.download("stopwords")  # TODO=find where this requirement originates from

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


def proc_submission(
    input_text: str,
    model_type: str,
    summary_type: str,
    num_beams,
    token_batch_length,
    length_penalty,
    #repetition_penalty,
    #no_repeat_ngram_size: int = 3,
    max_input_length: int = 768,
):
    """
    proc_submission - a helper function for the gradio module to process submissions
    
    Args:
        input_text (str): the input text to summarize
        model_size (str): the size of the model to use
        num_beams (int): the number of beams to use
        token_batch_length (int): the length of the token batches to use
        length_penalty (float): the length penalty to use
        repetition_penalty (float): the repetition penalty to use
        no_repeat_ngram_size (int): the no repeat ngram size to use
        max_input_length (int, optional): the maximum input length to use. Defaults to 768.
        
    Returns:
        str in HTML format, string of the summary, str of compression rate in %
    """

    settings = {
        "length_penalty": float(length_penalty),
        "repetition_penalty": 3.5,#float(repetition_penalty),
        "no_repeat_ngram_size": 3,
        "encoder_no_repeat_ngram_size": 4,
        "num_beams": int(num_beams),
        "min_length": 11,
        "max_length": int(token_batch_length // 4),
        "early_stopping": True,
    }
    st = time.perf_counter()
    history = {}
    clean_text = clean(input_text, lower=False)
    #max_input_length = 2048 if model_type == "tldr" else max_input_length
    processed = truncate_word_count(clean_text, max_input_length)

    if processed["was_truncated"]:
        tr_in = processed["truncated_text"]
        msg = f"Input text was truncated to {max_input_length} words (based on whitespace)"
        logging.warning(msg)
        history["WARNING"] = msg
    else:
        tr_in = input_text
        msg = None

    _summaries = summarize_via_tokenbatches(
        tr_in,
        model_led_det if (model_type == "LED" and summary_type == "Detailed") else model_det,
        tokenizer_led_det if (model_type == "LED" and summary_type == "Detailed") else tokenizer_det,
        model_led_tldr if (model_type == "LED" and summary_type == "TLDR") else model_tldr,
        tokenizer_led_tldr if (model_type == "LED" and summary_type == "TLDR") else tokenizer_tldr,
        batch_length=token_batch_length,
        **settings,
    )
    sum_text = [f"Section {i}: " + s["summary"][0] for i, s in enumerate(_summaries)]
    compression_rate = [
        f" - Section {i}: {round(s['compression_rate'],3)}"
        for i, s in enumerate(_summaries)
    ]

    sum_text_out = "\n".join(sum_text)
    history["compression_rate"] = "<br><br>"
    rate_out = "\n".join(compression_rate)
    rt = round((time.perf_counter() - st) / 60, 2)
    print(f"Runtime: {rt} minutes")
    html = ""
    html += f"<p>Runtime: {rt} minutes on CPU inference</p>"
    if msg is not None:
        html += f"<h2>WARNING:</h2><hr><b>{msg}</b><br><br>"

    html += ""

    return html, sum_text_out, rate_out


def load_single_example_text(
    example_path: str or Path,
):
    """
    load_single_example - a helper function for the gradio module to load examples
    Returns:
        list of str, the examples
    """
    global name_to_path
    full_ex_path = name_to_path[example_path]
    full_ex_path = Path(full_ex_path)
    # load the examples into a list
    with open(full_ex_path, "r", encoding="utf-8", errors="ignore") as f:
        raw_text = f.read()
        text = clean(raw_text, lower=False)
    return text


def load_uploaded_file(file_obj):
    """
    load_uploaded_file - process an uploaded file
    Args:
        file_obj (POTENTIALLY list): Gradio file object inside a list
    Returns:
        str, the uploaded file contents
    """

    # file_path = Path(file_obj[0].name)

    # check if mysterious file object is a list
    if isinstance(file_obj, list):
        file_obj = file_obj[0]
    file_path = Path(file_obj.name)
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            raw_text = f.read()
        text = clean(raw_text, lower=False)
        return text
    except Exception as e:
        logging.info(f"Trying to load file with path {file_path}, error: {e}")
        return "Error: Could not read file. Ensure that it is a valid text file with encoding UTF-8."


if __name__ == "__main__":

    model_det, tokenizer_det = load_model_and_tokenizer("Blaise-g/longt5_tglobal_large_sumpubmed")
    model_tldr, tokenizer_tldr = load_model_and_tokenizer("Blaise-g/longt5_tglobal_large_scitldr")
    model_led_det, tokenizer_led_det = load_model_and_tokenizer("Blaise-g/led_pubmed_sumpubmed_1")
    model_led_tldr, tokenizer_led_tldr = load_model_and_tokenizer("Blaise-g/led_large_sumpbumed_scitldr")

    name_to_path = load_example_filenames(_here / "examples")
    logging.info(f"Loaded {len(name_to_path)} examples")
    demo = gr.Blocks()

    with demo:

        gr.Markdown("# Automatic summarization of biomedical research papers with neural abstractive methods into a long and comprehensive synopsis or extreme TLDR summary version")
        gr.Markdown(
            "A rather simple demo (developed for my Master Thesis project) using an ad-hoc fine-tuned LongT5 or LED model to summarize long biomedical articles (or any scientific text related to the biomedical domain) into a detailed, explanatory synopsis or extreme TLDR summary."
        )
        with gr.Column():

            gr.Markdown("### Load Text Inputs, Select Model & Summary Type")
            gr.Markdown(
                "Enter text below in the text area. The text will be summarized [using the selected text generation parameters](https://huggingface.co/blog/how-to-generate). Optionally load an available example below or upload a file."
            )
            with gr.Row():
                summary_type = gr.Radio(
                    choices=["TLDR", "Detailed"], label="Summary Type", value="Detailed"
                )
                model_type = gr.Radio(
                    choices=["LongT5", "LED"], label="Model Architecture", value="LongT5"
                )
                num_beams = gr.Radio(
                    choices=[2, 3, 4, 5, 6],
                    label="Beam Search: # of Beams",
                    value=2,
                )
            gr.Markdown(
                "_The LED model is less performant than the LongT5 model, but it's smaller in terms of size and therefore all other parameters being equal allows for a longer input sequence._"
            )
            with gr.Row():
                length_penalty = gr.inputs.Slider(
                    minimum=0.5,
                    maximum=1.0,
                    label="length penalty",
                    default=0.7,
                    step=0.05,
                )
                token_batch_length = gr.Radio(
                    choices=[512, 768, 1024],
                    label="token batch length",
                    value=512,
                )

            #with gr.Row():
                #repetition_penalty = gr.inputs.Slider(
                    #minimum=1.0,
                    #maximum=5.0,
                    #label="repetition penalty",
                    #default=3.5,
                    #step=0.1,
                #)
                #no_repeat_ngram_size = gr.Radio(
                    #choices=[2, 3, 4],
                    #label="no repeat ngram size",
                    #value=3,
               # )
            with gr.Row():
                example_name = gr.Dropdown(
                    list(name_to_path.keys()),
                    label="Choose an Example",
                )
                load_examples_button = gr.Button(
                    "Load Example",
                )
            input_text = gr.Textbox(
                lines=6,
                label="Input Text (for summarization)",
                placeholder="Enter any scientific text to be condensed into a long and comprehensive digested format or an extreme TLDR summary version, the text will be preprocessed and truncated if necessary to fit within the computational constraints. The models were trained to handle long scientific papers but generalize reasonably well also to shorter text documents like abstracts with an appropriate. Might take a while to produce long summaries :)",
            )
            gr.Markdown("Upload your own file:")
            with gr.Row():
                uploaded_file = gr.File(
                    label="Upload a text file",
                    file_count="single",
                    type="file",
                )
                load_file_button = gr.Button("Load Uploaded File")

            gr.Markdown("---")

        with gr.Column():
            gr.Markdown("## Generate Summary")
            gr.Markdown(
                "Summary generation should take approximately less than 2 minutes for most settings."
            )
            summarize_button = gr.Button(
                "Summarize!",
                variant="primary",
            )

            output_text = gr.HTML("<p><em>Output will appear below:</em></p>")
            gr.Markdown("### Summary Output")
            summary_text = gr.Textbox(
                label="Summary 📝", placeholder="The generated 📝 will appear here"
            )
            gr.Markdown(
                "The compression rate indicates the ratio between the machine-generated summary length and the input text (from 0% to 100%). The higher the compression rate the more extreme the summary is."
            )
            compression_rate = gr.Textbox(
                label="Compression rate 🗜", placeholder="The 🗜 will appear here"
            )

            gr.Markdown("---")

        with gr.Column():
            gr.Markdown("## About the Models")
            gr.Markdown(
                "- [Blaise-g/longt5_tglobal_large_sumpubmed](https://huggingface.co/Blaise-g/longt5_tglobal_large_sumpubmed) is a fine-tuned checkpoint of [Stancld/longt5-tglobal-large-16384-pubmed-3k_steps](https://huggingface.co/Stancld/longt5-tglobal-large-16384-pubmed-3k_steps) on the [SumPubMed dataset](https://aclanthology.org/2021.acl-srw.30/). [Blaise-g/longt5_tglobal_large_scitldr](https://huggingface.co/Blaise-g/longt5_tglobal_large_scitldr) is a fine-tuned checkpoint of [Blaise-g/longt5_tglobal_large_sumpubmed](https://huggingface.co/Blaise-g/longt5_tglobal_large_sumpubmed) on the [Scitldr dataset](https://arxiv.org/abs/2004.15011). The goal was to create two models capable of handling the complex information contained in long biomedical documents and subsequently producing scientific summaries according to one of the two possible levels of conciseness: 1) A long explanatory synopsis that retains the majority of domain-specific language used in the original source text. 2)A one sentence long, TLDR style summary."
            )
            gr.Markdown(
                "- The two most important parameters-empirically-are the `num_beams` and `token_batch_length`. However, increasing these will also increase the amount of time it takes to generate a summary. The `length_penalty` and `repetition_penalty` parameters are also important for the model to generate good summaries."
            )
            gr.Markdown("---")

        load_examples_button.click(
            fn=load_single_example_text, inputs=[example_name], outputs=[input_text]
        )

        load_file_button.click(
            fn=load_uploaded_file, inputs=uploaded_file, outputs=[input_text]
        )

        summarize_button.click(
            fn=proc_submission,
            inputs=[
                input_text,
                summary_type,
                model_type,
                num_beams,
                token_batch_length,
                length_penalty,
            ],
            outputs=[output_text, summary_text, compression_rate],
        )

    demo.launch(enable_queue=True, share=False)