"""
app.py - the main module for the gradio app

Usage:
    python app.py
"""
import contextlib
import gc
import logging
import os
import random
import re
import time
from pathlib import Path

os.environ["USE_TORCH"] = "1"
os.environ[
    "TOKENIZERS_PARALLELISM"
] = "false"  # parallelism on tokenizers is buggy with gradio

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

import gradio as gr
import nltk
import torch
from cleantext import clean
from doctr.models import ocr_predictor

from pdf2text import convert_PDF_to_Text
from summarize import load_model_and_tokenizer, summarize_via_tokenbatches
from utils import load_example_filenames, saves_summary, truncate_word_count

_here = Path(__file__).parent

nltk.download("stopwords", quiet=True)


MODEL_OPTIONS = [
    "pszemraj/long-t5-tglobal-base-16384-book-summary",
    "pszemraj/long-t5-tglobal-base-sci-simplify",
    "pszemraj/long-t5-tglobal-base-sci-simplify-elife",
    "pszemraj/long-t5-tglobal-base-16384-booksci-summary-v1",
    "pszemraj/pegasus-x-large-book-summary",
]  # models users can choose from


def predict(
    input_text: str,
    model_name: str,
    token_batch_length: int = 1024,
    empty_cache: bool = True,
    **settings,
) -> list:
    """
    predict - helper fn to support multiple models for summarization at once

    :param str input_text: the input text to summarize
    :param str model_name: model name to use
    :param int token_batch_length: the length of the token batches to use
    :param bool empty_cache: whether to empty the cache before loading a new= model
    :return: list of dicts with keys "summary" and "score"
    """
    if torch.cuda.is_available() and empty_cache:
        torch.cuda.empty_cache()

    model, tokenizer = load_model_and_tokenizer(model_name)
    summaries = summarize_via_tokenbatches(
        input_text,
        model,
        tokenizer,
        batch_length=token_batch_length,
        **settings,
    )

    del model
    del tokenizer
    gc.collect()

    return summaries


def proc_submission(
    input_text: str,
    model_name: str,
    num_beams: int,
    token_batch_length: int,
    length_penalty: float,
    repetition_penalty: float,
    no_repeat_ngram_size: int,
    max_input_length: int = 4096,
):
    """
    proc_submission - a helper function for the gradio module to process submissions

    Args:
        input_text (str): the input text to summarize
        model_name (str): the hf model tag 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 2048.

    Returns:
        str in HTML format, string of the summary, str of score
    """

    settings = {
        "length_penalty": float(length_penalty),
        "repetition_penalty": float(repetition_penalty),
        "no_repeat_ngram_size": int(no_repeat_ngram_size),
        "encoder_no_repeat_ngram_size": 4,
        "num_beams": int(num_beams),
        "min_length": 4,
        "max_length": int(token_batch_length // 4),
        "early_stopping": True,
        "do_sample": False,
    }
    st = time.perf_counter()
    history = {}
    clean_text = clean(input_text, lower=False)
    processed = truncate_word_count(clean_text, max_words=max_input_length)

    if processed["was_truncated"]:
        tr_in = processed["truncated_text"]
        # create elaborate HTML warning
        input_wc = re.split(r"\s+", input_text)
        msg = f"""
        <div style="background-color: #FFA500; color: white; padding: 20px;">
        <h3>Warning</h3>
        <p>Input text was truncated to {max_input_length} words. That's about {100*max_input_length/len(input_wc):.2f}% of the submission.</p>
        </div>
        """
        logging.warning(msg)
        history["WARNING"] = msg
    else:
        tr_in = input_text
        msg = None

    if len(input_text) < 50:
        # this is essentially a different case from the above
        msg = f"""
        <div style="background-color: #880808; color: white; padding: 20px;">
        <h3>Warning</h3>
        <p>Input text is too short to summarize. Detected {len(input_text)} characters.
        Please load text by selecting an example from the dropdown menu or by pasting text into the text box.</p>
        </div>
        """
        logging.warning(msg)
        logging.warning("RETURNING EMPTY STRING")
        history["WARNING"] = msg

        return msg, "", []

    _summaries = predict(
        input_text=tr_in,
        model_name=model_name,
        token_batch_length=token_batch_length,
        **settings,
    )
    sum_text = [
        f"Batch {i}:\n\t" + s["summary"][0] for i, s in enumerate(_summaries, start=1)
    ]
    sum_scores = [
        f" - Batch Summary {i}: {round(s['summary_score'],4)}"
        for i, s in enumerate(_summaries)
    ]

    sum_text_out = "\n".join(sum_text)
    history["Summary Scores"] = "<br><br>"
    scores_out = "\n".join(sum_scores)
    rt = round((time.perf_counter() - st) / 60, 2)
    logging.info(f"Runtime: {rt} minutes")
    html = ""
    html += f"<p>Runtime: {rt} minutes with model: {model_name}</p>"
    if msg is not None:
        html += msg

    html += ""

    # save to file
    settings["model_name"] = model_name
    saved_file = saves_summary(_summaries, **settings)

    return html, sum_text_out, scores_out, saved_file


def load_single_example_text(
    example_path: str or Path,
    max_pages: int = 20,
) -> str:
    """
    load_single_example_text - loads a single example text file

    :param strorPath example_path: name of the example to load
    :param int max_pages: the maximum number of pages to load from a PDF
    :return str: the text of the example
    """
    global name_to_path
    full_ex_path = name_to_path[example_path]
    full_ex_path = Path(full_ex_path)
    if full_ex_path.suffix == ".txt":
        with open(full_ex_path, "r", encoding="utf-8", errors="ignore") as f:
            raw_text = f.read()
        text = clean(raw_text, lower=False)
    elif full_ex_path.suffix == ".pdf":
        logging.info(f"Loading PDF file {full_ex_path}")
        conversion_stats = convert_PDF_to_Text(
            full_ex_path,
            ocr_model=ocr_model,
            max_pages=max_pages,
        )
        text = conversion_stats["converted_text"]
    else:
        logging.error(f"Unknown file type {full_ex_path.suffix}")
        text = "ERROR - check example path"

    return text


def load_uploaded_file(file_obj, max_pages: int = 20, lower: bool = False) -> str:
    """
    load_uploaded_file - loads a file uploaded by the user

    :param file_obj (POTENTIALLY list): Gradio file object inside a list
    :param int max_pages: the maximum number of pages to load from a PDF
    :param bool lower: whether to lowercase the text
    :return str: the text of the file
    """
    # 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:
        logging.info(f"Loading file:\t{file_path}")
        if file_path.suffix == ".txt":
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                raw_text = f.read()
            text = clean(raw_text, lower=lower)
        elif file_path.suffix == ".pdf":
            logging.info(f"loading as PDF file {file_path}")
            conversion_stats = convert_PDF_to_Text(
                file_path,
                ocr_model=ocr_model,
                max_pages=max_pages,
            )
            text = conversion_stats["converted_text"]
        else:
            logging.error(f"Unknown file type {file_path.suffix}")
            text = "ERROR - check file - unknown file type"

        return text
    except Exception as e:
        logging.error(f"Trying to load file:\t{file_path},\nerror:\t{e}")
        return "Error: Could not read file. Ensure that it is a valid text file with encoding UTF-8 if text, and a PDF if PDF."


if __name__ == "__main__":
    logging.info("Starting app instance")
    logging.info("Loading OCR model")
    with contextlib.redirect_stdout(None):
        ocr_model = ocr_predictor(
            "db_resnet50",
            "crnn_mobilenet_v3_large",
            pretrained=True,
            assume_straight_pages=True,
        )
    name_to_path = load_example_filenames(_here / "examples")
    logging.info(f"Loaded {len(name_to_path)} examples")
    demo = gr.Blocks()
    _examples = list(name_to_path.keys())
    with demo:
        gr.Markdown("# Document Summarization with Long-Document Transformers")
        gr.Markdown(
            "This is an example use case for fine-tuned long document transformers. The model is trained on book summaries (via the BookSum dataset). The models in this demo are [LongT5-base](https://huggingface.co/pszemraj/long-t5-tglobal-base-16384-book-summary) and [Pegasus-X-Large](https://huggingface.co/pszemraj/pegasus-x-large-book-summary)."
        )
        with gr.Column():
            gr.Markdown("## Load Inputs & Select Parameters")
            gr.Markdown(
                "Enter text below in the text area. The text will be summarized [using the selected parameters](https://huggingface.co/blog/how-to-generate). Optionally load an example below or upload a file. (`.txt` or `.pdf` - _[link to guide](https://i.imgur.com/c6Cs9ly.png)_)"
            )
            with gr.Row(variant="compact"):
                with gr.Column(scale=0.5, variant="compact"):
                    model_name = gr.Dropdown(
                        choices=MODEL_OPTIONS,
                        value=MODEL_OPTIONS[0],
                        label="Model Name",
                    )
                    num_beams = gr.Radio(
                        choices=[2, 3, 4],
                        label="Beam Search: # of Beams",
                        value=2,
                    )
                    load_examples_button = gr.Button(
                        "Load Example in Dropdown",
                    )
                    load_file_button = gr.Button("Load an Uploaded File")
                with gr.Column(variant="compact"):
                    example_name = gr.Dropdown(
                        _examples,
                        label="Examples",
                        value=random.choice(_examples),
                    )
                    uploaded_file = gr.File(
                        label="File Upload",
                        file_count="single",
                        type="file",
                    )
            with gr.Row():
                input_text = gr.Textbox(
                    lines=4,
                    label="Input Text (for summarization)",
                    placeholder="Enter text to summarize, the text will be cleaned and truncated on Spaces. Narrative, academic (both papers and lecture transcription), and article text work well. May take a bit to generate depending on the input text :)",
                )

        with gr.Column():
            gr.Markdown("## Generate Summary")
            gr.Markdown(
                "Summarization should take ~1-2 minutes for most settings, but may extend up to 5-10 minutes in some scenarios."
            )
            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 summary will appear here"
            )
            gr.Markdown(
                "The summary scores can be thought of as representing the quality of the summary. less-negative numbers (closer to 0) are better:"
            )
            summary_scores = gr.Textbox(
                label="Summary Scores", placeholder="Summary scores will appear here"
            )

            text_file = gr.File(
                label="Download as Text File",
                file_count="single",
                type="file",
                interactive=False,
            )

        gr.Markdown("---")
        with gr.Column():
            gr.Markdown("### Advanced Settings")
            with gr.Row(variant="compact"):
                length_penalty = gr.Slider(
                    minimum=0.5,
                    maximum=1.0,
                    label="length penalty",
                    default=0.7,
                    step=0.05,
                )
                token_batch_length = gr.Radio(
                    choices=[512, 1024, 1536, 2048],
                    label="token batch length",
                    value=1536,
                )

            with gr.Row(variant="compact"):
                repetition_penalty = gr.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.Column():
            gr.Markdown("### About")
            gr.Markdown(
                "- Models are fine-tuned on the [BookSum dataset](https://arxiv.org/abs/2105.08209). The goal was to create a model that generalizes well and is useful for summarizing text in academic and everyday use."
            )
            gr.Markdown(
                "- _Update April 2023:_ Additional models fine-tuned on the [PLOS](https://huggingface.co/datasets/pszemraj/scientific_lay_summarisation-plos-norm) and [ELIFE](https://huggingface.co/datasets/pszemraj/scientific_lay_summarisation-elife-norm) subsets of the [scientific lay summaries](https://arxiv.org/abs/2210.09932) dataset are available (see dropdown at the top)."
            )
            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,
                model_name,
                num_beams,
                token_batch_length,
                length_penalty,
                repetition_penalty,
                no_repeat_ngram_size,
            ],
            outputs=[output_text, summary_text, summary_scores, text_file],
        )

    demo.launch(enable_queue=True)