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_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):
  if flagging_callback.flag(list(args), flag_option = None):
      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_2",
                                                 organization = "fsdlredteam",
                                                 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,
                                       show_progress=True,
                                       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) 
                            
#demo.launch(debug=True)
if __name__ == "__main__":
    demo.launch(enable_queue=False, debug=True)