import os
import torch

import numpy as np
import gradio as gr

from random import sample
from detoxify import Detoxify
from datasets import load_dataset
from huggingface_hub import HfApi, ModelFilter, ModelSearchArguments
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPTNeoForCausalLM
from transformers import BloomTokenizerFast, BloomForCausalLM

HF_AUTH_TOKEN = os.environ.get("hf_token" or True)

DATASET = "allenai/real-toxicity-prompts"

CHECKPOINTS = {
    "DistilGPT2 by HuggingFace 🤗": "distilgpt2",
    "GPT-Neo 125M by EleutherAI 🤖": "EleutherAI/gpt-neo-125M",
    "BLOOM 560M by BigScience 🌸": "bigscience/bloom-560m",
    "Custom Model": None,
}

MODEL_CLASSES = {
    "DistilGPT2 by HuggingFace 🤗": (GPT2LMHeadModel, GPT2Tokenizer),
    "GPT-Neo 125M by EleutherAI 🤖": (GPTNeoForCausalLM, GPT2Tokenizer),
    "BLOOM 560M by BigScience 🌸": (BloomForCausalLM, BloomTokenizerFast),
    "Custom Model": (AutoModelForCausalLM, AutoTokenizer),
}

CHOICES = sorted(list(CHECKPOINTS.keys())[:3])


def load_model(model_name, custom_model_path, token):
    try:
        model_class, tokenizer_class = MODEL_CLASSES[model_name]
        model_path = CHECKPOINTS[model_name]

    except KeyError:
        model_class, tokenizer_class = MODEL_CLASSES["Custom Model"]
        model_path = custom_model_path or model_name

    model = model_class.from_pretrained(model_path, use_auth_token=token)
    tokenizer = tokenizer_class.from_pretrained(model_path, use_auth_token=token)

    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id

    model.eval()

    return model, tokenizer


MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop


def set_seed(seed, n_gpu):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if n_gpu > 0:
        torch.cuda.manual_seed_all(seed)


def adjust_length_to_model(length, max_sequence_length):
    if length < 0 and max_sequence_length > 0:
        length = max_sequence_length
    elif 0 < max_sequence_length < length:
        length = max_sequence_length  # No generation bigger than model size
    elif length < 0:
        length = MAX_LENGTH  # avoid infinite loop
    return length


def generate(
    model_name,
    token,
    custom_model_path,
    input_sentence,
    length=75,
    temperature=0.7,
    top_k=50,
    top_p=0.95,
    seed=42,
    no_cuda=False,
    num_return_sequences=1,
    stop_token=".",
):
    # load device
    # if not no_cuda:
    device = torch.device(
        "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"
    )
    n_gpu = 0 if no_cuda else torch.cuda.device_count()

    # Set seed
    set_seed(seed, n_gpu)

    # Load model
    model, tokenizer = load_model(model_name, custom_model_path, token)
    model.to(device)

    # length = adjust_length_to_model(length, max_sequence_length=model.config.max_position_embeddings)

    # Tokenize input
    encoded_prompt = tokenizer.encode(
        input_sentence, add_special_tokens=False, return_tensors="pt"
    )

    encoded_prompt = encoded_prompt.to(device)

    input_ids = encoded_prompt

    # Generate output
    output_sequences = model.generate(
        input_ids=input_ids,
        max_length=length + len(encoded_prompt[0]),
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        do_sample=True,
        num_return_sequences=num_return_sequences,
    )
    generated_sequences = list()

    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
        generated_sequence = generated_sequence.tolist()
        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
        # remove prompt
        text = text[
            len(
                tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)
            ) :
        ]

        # remove all text after last occurence of stop_token
        text = text[: text.rfind(stop_token) + 1]

        generated_sequences.append(text)

    return generated_sequences[0]


def show_mode(mode):
    if mode == "Single Model":
        return (gr.update(visible=True), gr.update(visible=False))
    if mode == "Multi-Model":
        return (gr.update(visible=False), gr.update(visible=True))


def prepare_dataset(dataset):
    dataset = load_dataset(dataset, split="train")
    return dataset


def load_prompts(dataset):
    prompts = [dataset[i]["prompt"]["text"] for i in range(len(dataset))]
    return prompts


def random_sample(prompt_list):
    random_sample = sample(prompt_list, 10)
    return random_sample


def show_dataset(dataset):
    raw_data = prepare_dataset(dataset)
    prompts = load_prompts(raw_data)

    return (
        gr.update(
            choices=random_sample(prompts),
            label="You can find below a random subset from the RealToxicityPrompts dataset",
            visible=True,
        ),
        gr.update(visible=True),
        prompts,
    )


def update_dropdown(prompts):
    return gr.update(choices=random_sample(prompts))


def show_search_bar(value):
    if value == "Custom Model":
        return (value, gr.update(visible=True))
    else:
        return (value, gr.update(visible=False))


def search_model(model_name, token):
    api = HfApi()

    model_args = ModelSearchArguments()
    filt = ModelFilter(
        task=model_args.pipeline_tag.TextGeneration, library=model_args.library.PyTorch
    )

    results = api.list_models(filter=filt, search=model_name, use_auth_token=token)
    model_list = [model.modelId for model in results]

    return gr.update(
        visible=True,
        choices=model_list,
        label="Choose the model",
    )


def show_api_key_textbox(checkbox):
    if checkbox:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


def forward_model_choice(model_choice_path):
    return (model_choice_path, model_choice_path)


def auto_complete(input, generated):
    output = input + " " + generated
    output_spans = [{"entity": "OUTPUT", "start": len(input), "end": len(output)}]
    completed_prompt = {"text": output, "entities": output_spans}
    return completed_prompt


def process_user_input(
    model, token, custom_model_path, input, length, temperature, top_p, top_k
):
    warning = "Please enter a valid prompt."
    if input == None:
        generated = warning
    else:
        generated = generate(
            model_name=model,
            token=token,
            custom_model_path=custom_model_path,
            input_sentence=input,
            length=length,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
        )
        generated = generated.replace("\n", " ")
        generated_with_spans = auto_complete(input=input, generated=generated)

    return (
        gr.update(value=generated_with_spans),
        gr.update(visible=True),
        gr.update(visible=True),
        input,
        generated,
    )


def pass_to_textbox(input):
    return gr.update(value=input)


def run_detoxify(text):
    results = Detoxify("original").predict(text)
    json_ready_results = {cat: float(score) for (cat, score) in results.items()}
    return json_ready_results


def compute_toxi_output(output_text):
    scores = run_detoxify(output_text)
    return (gr.update(value=scores, visible=True), gr.update(visible=True))


def compute_change(input, output):
    change_percent = round(((float(output) - input) / input) * 100, 2)
    return change_percent


def compare_toxi_scores(input_text, output_scores):
    input_scores = run_detoxify(input_text)
    json_ready_results = {cat: float(score) for (cat, score) in input_scores.items()}

    compare_scores = {
        cat: compute_change(json_ready_results[cat], output_scores[cat])
        for cat in json_ready_results
        for cat in output_scores
    }

    return (
        gr.update(value=json_ready_results, visible=True),
        gr.update(value=compare_scores, visible=True),
    )


def show_flag_choices():
    return gr.update(visible=True)


def update_flag(flag_value):
    return (
        flag_value,
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=False),
    )


def upload_flag(*args):
    flags = list(args)
    flags[1] = bytes(flags[1], "utf-8")
    flagging_callback.flag(flags)
    return gr.update(visible=True)


def forward_model_choice_multi(model_choice_path):
    CHOICES.append(model_choice_path)
    return gr.update(choices=CHOICES)


def process_user_input_multi(models, input, token, length, temperature, top_p, top_k):
    warning = "Please enter a valid prompt."
    if input == None:
        generated = warning
    else:
        generated_dict = {
            model: generate(
                model_name=model,
                token=token,
                custom_model_path=None,
                input_sentence=input,
                length=length,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
            )
            for model in sorted(models)
        }
        generated_with_spans_dict = {
            model: auto_complete(input, generated)
            for model, generated in generated_dict.items()
        }

        update_outputs = [
            gr.HighlightedText.update(value=output, label=model)
            for model, output in generated_with_spans_dict.items()
        ]
        update_hide = [
            gr.HighlightedText.update(visible=False) for i in range(10 - len(models))
        ]
        return update_outputs + update_hide


def show_choices_multi(models):
    update_show = [gr.HighlightedText.update(visible=True) for model in sorted(models)]
    update_hide = [
        gr.HighlightedText.update(visible=False, value=None, label=None)
        for i in range(10 - len(models))
    ]

    return update_show + update_hide


def show_params(checkbox):
    if checkbox == True:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


CSS = """
#inside_group {
  padding-top: 0.6em;
  padding-bottom: 0.6em;
}
#pw textarea {
    -webkit-text-security: disc;
}
"""

with gr.Blocks(css=CSS) as demo:
    dataset = gr.Variable(value=DATASET)
    prompts_var = gr.Variable(value=None)
    input_var = gr.Variable(label="Input Prompt", value=None)
    output_var = gr.Variable(label="Output", value=None)
    model_choice = gr.Variable(label="Model", value=None)
    custom_model_path = gr.Variable(value=None)
    flag_choice = gr.Variable(label="Flag", value=None)

    flagging_callback = gr.HuggingFaceDatasetSaver(
        hf_token=HF_AUTH_TOKEN,
        dataset_name="fsdlredteam/flagged_3",
        private=True,
    )

    gr.Markdown("<p align='center'><img src='https://i.imgur.com/ZxbbLUQ.png>'/></p>")
    gr.Markdown("<h1 align='center'>BuggingSpace</h1>")
    gr.Markdown(
        "<h2 align='center'>FSDL 2022 Red-Teaming Open-Source Models Project</h2>"
    )
    gr.Markdown(
        "### Pick a text generation model below, write a prompt and explore the output"
    )
    gr.Markdown("### Or compare the output of multiple models at the same time")

    choose_mode = gr.Radio(
        choices=["Single Model", "Multi-Model"],
        value="Single Model",
        interactive=True,
        visible=True,
        show_label=False,
    )

    with gr.Group() as single_model:
        gr.Markdown(
            "You can upload any model from the Hugging Face hub -even private ones, \
                provided you use your private key! "
            "Write your prompt or alternatively use one from the \
                [RealToxicityPrompts](https://allenai.org/data/real-toxicity-prompts) dataset."
        )
        gr.Markdown(
            "Use it to audit the model for potential failure modes, \
                analyse its output with the Detoxify suite and contribute by reporting any problematic result."
        )
        gr.Markdown(
            "Beware ! Generation can take up to a few minutes with very large models."
        )

        with gr.Row():
            with gr.Column(scale=1):  # input & prompts dataset exploration
                gr.Markdown("### 1. Select a prompt", elem_id="inside_group")

                input_text = gr.Textbox(
                    label="Write your prompt below.",
                    interactive=True,
                    lines=4,
                    elem_id="inside_group",
                )

                gr.Markdown("— or —", elem_id="inside_group")

                inspo_button = gr.Button(
                    "Click here if you need some inspiration", elem_id="inside_group"
                )

                prompts_drop = gr.Dropdown(visible=False, elem_id="inside_group")

                randomize_button = gr.Button(
                    "Show another subset", visible=False, elem_id="inside_group"
                )

                show_params_checkbox_single = gr.Checkbox(
                    label="Set custom params", interactive=True, value=False
                )

                with gr.Box(visible=False) as params_box_single:
                    length_single = gr.Slider(
                        label="Output length",
                        visible=True,
                        interactive=True,
                        minimum=50,
                        maximum=200,
                        value=75,
                    )

                    top_k_single = gr.Slider(
                        label="top_k",
                        visible=True,
                        interactive=True,
                        minimum=1,
                        maximum=100,
                        value=50,
                    )

                    top_p_single = gr.Slider(
                        label="top_p",
                        visible=True,
                        interactive=True,
                        minimum=0.1,
                        maximum=1,
                        value=0.95,
                    )

                    temperature_single = gr.Slider(
                        label="temperature",
                        visible=True,
                        interactive=True,
                        minimum=0.1,
                        maximum=1,
                        value=0.7,
                    )

            with gr.Column(scale=1):  # Model choice & output
                gr.Markdown("### 2. Evaluate output")

                model_radio = gr.Radio(
                    choices=list(CHECKPOINTS.keys()),
                    label="Model",
                    interactive=True,
                    elem_id="inside_group",
                )

                search_bar = gr.Textbox(
                    label="Search model",
                    interactive=True,
                    visible=False,
                    elem_id="inside_group",
                )
                model_drop = gr.Dropdown(visible=False)

                private_checkbox = gr.Checkbox(
                    visible=True, label="Private Model ?", elem_id="inside_group"
                )

                api_key_textbox = gr.Textbox(
                    label="Enter your AUTH TOKEN below",
                    value=None,
                    interactive=True,
                    visible=False,
                    elem_id="pw",
                )

                generate_button = gr.Button(
                    "Submit your prompt", elem_id="inside_group"
                )

                output_spans = gr.HighlightedText(visible=True, label="Generated text")

                flag_button = gr.Button(
                    "Report output here", visible=False, elem_id="inside_group"
                )

        with gr.Row():  # Flagging
            with gr.Column(scale=1):
                flag_radio = gr.Radio(
                    choices=[
                        "Toxic",
                        "Offensive",
                        "Repetitive",
                        "Incorrect",
                        "Other",
                    ],
                    label="What's wrong with the output ?",
                    interactive=True,
                    visible=False,
                    elem_id="inside_group",
                )

                user_comment = gr.Textbox(
                    label="(Optional) Briefly describe the issue",
                    visible=False,
                    interactive=True,
                    elem_id="inside_group",
                )

            confirm_flag_button = gr.Button(
                "Confirm report", visible=False, elem_id="inside_group"
            )

        with gr.Row():  # Flagging success
            success_message = gr.Markdown(
                "Your report has been successfully registered. Thank you!",
                visible=False,
                elem_id="inside_group",
            )

        with gr.Row():  # Toxicity buttons
            toxi_button = gr.Button(
                "Run a toxicity analysis of the model's output",
                visible=False,
                elem_id="inside_group",
            )
            toxi_button_compare = gr.Button(
                "Compare toxicity on input and output",
                visible=False,
                elem_id="inside_group",
            )

        with gr.Row():  # Toxicity scores
            toxi_scores_input = gr.JSON(
                label="Detoxify classification of your input",
                visible=False,
                elem_id="inside_group",
            )
            toxi_scores_output = gr.JSON(
                label="Detoxify classification of the model's output",
                visible=False,
                elem_id="inside_group",
            )
            toxi_scores_compare = gr.JSON(
                label="Percentage change between Input and Output",
                visible=False,
                elem_id="inside_group",
            )

    with gr.Group(visible=False) as multi_model:
        model_list = list()

        gr.Markdown(
            "#### Run the same input on multiple models and compare the outputs"
        )
        gr.Markdown(
            "You can upload any model from the Hugging Face hub -even private ones, provided you use your private key!"
        )
        gr.Markdown(
            "Use this feature to compare the same model at different checkpoints"
        )
        gr.Markdown("Or to benchmark your model against another one as a reference.")
        gr.Markdown(
            "Beware ! Generation can take up to a few minutes with very large models."
        )

        with gr.Row(elem_id="inside_group"):
            with gr.Column():
                models_multi = gr.CheckboxGroup(
                    choices=CHOICES,
                    label="Models",
                    interactive=True,
                    elem_id="inside_group",
                    value=None,
                )
            with gr.Column():
                generate_button_multi = gr.Button(
                    "Submit your prompt", elem_id="inside_group"
                )

                show_params_checkbox_multi = gr.Checkbox(
                    label="Set custom params", interactive=True, value=False
                )

                with gr.Box(visible=False) as params_box_multi:
                    length_multi = gr.Slider(
                        label="Output length",
                        visible=True,
                        interactive=True,
                        minimum=50,
                        maximum=200,
                        value=75,
                    )

                    top_k_multi = gr.Slider(
                        label="top_k",
                        visible=True,
                        interactive=True,
                        minimum=1,
                        maximum=100,
                        value=50,
                    )

                    top_p_multi = gr.Slider(
                        label="top_p",
                        visible=True,
                        interactive=True,
                        minimum=0.1,
                        maximum=1,
                        value=0.95,
                    )

                    temperature_multi = gr.Slider(
                        label="temperature",
                        visible=True,
                        interactive=True,
                        minimum=0.1,
                        maximum=1,
                        value=0.7,
                    )

        with gr.Row(elem_id="inside_group"):
            with gr.Column(elem_id="inside_group", scale=1):
                input_text_multi = gr.Textbox(
                    label="Write your prompt below.",
                    interactive=True,
                    lines=4,
                    elem_id="inside_group",
                )

            with gr.Column(elem_id="inside_group", scale=1):
                search_bar_multi = gr.Textbox(
                    label="Search another model",
                    interactive=True,
                    visible=True,
                    elem_id="inside_group",
                )

                model_drop_multi = gr.Dropdown(visible=False, elem_id="inside_group")

                private_checkbox_multi = gr.Checkbox(
                    visible=True, label="Private Model ?"
                )

                api_key_textbox_multi = gr.Textbox(
                    label="Enter your AUTH TOKEN below",
                    value=None,
                    interactive=True,
                    visible=False,
                    elem_id="pw",
                )

        with gr.Row() as outputs_row:
            for i in range(10):
                output_spans_multi = gr.HighlightedText(
                    visible=False, elem_id="inside_group"
                )
                model_list.append(output_spans_multi)

    with gr.Row():
        gr.Markdown(
            "App made during the [FSDL course](https://fullstackdeeplearning.com) \
                 by Team53: Jean-Antoine, Sajenthan, Sashank, Kemp, Srihari, Astitwa"
        )

    # Single Model

    choose_mode.change(
        fn=show_mode, inputs=choose_mode, outputs=[single_model, multi_model]
    )

    inspo_button.click(
        fn=show_dataset,
        inputs=dataset,
        outputs=[prompts_drop, randomize_button, prompts_var],
    )

    prompts_drop.change(fn=pass_to_textbox, inputs=prompts_drop, outputs=input_text)

    randomize_button.click(
        fn=update_dropdown, inputs=prompts_var, outputs=prompts_drop
    ),

    model_radio.change(
        fn=show_search_bar, inputs=model_radio, outputs=[model_choice, search_bar]
    )

    search_bar.submit(
        fn=search_model,
        inputs=[search_bar, api_key_textbox],
        outputs=model_drop,
        show_progress=True,
    )

    private_checkbox.change(
        fn=show_api_key_textbox, inputs=private_checkbox, outputs=api_key_textbox
    )

    model_drop.change(
        fn=forward_model_choice,
        inputs=model_drop,
        outputs=[model_choice, custom_model_path],
    )

    generate_button.click(
        fn=process_user_input,
        inputs=[
            model_choice,
            api_key_textbox,
            custom_model_path,
            input_text,
            length_single,
            temperature_single,
            top_p_single,
            top_k_single,
        ],
        outputs=[output_spans, toxi_button, flag_button, input_var, output_var],
        show_progress=True,
    )

    toxi_button.click(
        fn=compute_toxi_output,
        inputs=output_var,
        outputs=[toxi_scores_output, toxi_button_compare],
        show_progress=True,
    )

    toxi_button_compare.click(
        fn=compare_toxi_scores,
        inputs=[input_text, toxi_scores_output],
        outputs=[toxi_scores_input, toxi_scores_compare],
        show_progress=True,
    )

    flag_button.click(fn=show_flag_choices, inputs=None, outputs=flag_radio)

    flag_radio.change(
        fn=update_flag,
        inputs=flag_radio,
        outputs=[flag_choice, confirm_flag_button, user_comment, flag_button],
    )

    flagging_callback.setup(
        [input_var, output_var, model_choice, user_comment, flag_choice],
        "flagged_data_points",
    )

    confirm_flag_button.click(
        fn=upload_flag,
        inputs=[input_var, output_var, model_choice, user_comment, flag_choice],
        outputs=success_message,
    )

    show_params_checkbox_single.change(
        fn=show_params, inputs=show_params_checkbox_single, outputs=params_box_single
    )

    # Model comparison

    search_bar_multi.submit(
        fn=search_model,
        inputs=[search_bar_multi, api_key_textbox_multi],
        outputs=model_drop_multi,
        show_progress=True,
    )

    show_params_checkbox_multi.change(
        fn=show_params, inputs=show_params_checkbox_multi, outputs=params_box_multi
    )

    private_checkbox_multi.change(
        fn=show_api_key_textbox,
        inputs=private_checkbox_multi,
        outputs=api_key_textbox_multi,
    )

    model_drop_multi.change(
        fn=forward_model_choice_multi, inputs=model_drop_multi, outputs=[models_multi]
    )

    models_multi.change(fn=show_choices_multi, inputs=models_multi, outputs=model_list)

    generate_button_multi.click(
        fn=process_user_input_multi,
        inputs=[
            models_multi,
            input_text_multi,
            api_key_textbox_multi,
            length_multi,
            temperature_multi,
            top_p_multi,
            top_k_multi,
        ],
        outputs=model_list,
        show_progress=True,
    )

if __name__ == "__main__":
    # demo.queue(concurrency_count=3)
    demo.launch(debug=True)