import datetime
import gradio as gr
from langdetect import detect, DetectorFactory, detect_langs
from transformers import pipeline

models = {'en': 'Narsil/deberta-large-mnli-zero-cls', # English
          'de': 'Sahajtomar/German_Zeroshot', # German
          'es': 'Recognai/zeroshot_selectra_medium', # Spanish
          'it': 'joeddav/xlm-roberta-large-xnli', # Italian
          'ru': 'DeepPavlov/xlm-roberta-large-en-ru-mnli', # Russian
          'no': 'NbAiLab/nb-bert-base-mnli'} # Norsk

hypothesis_templates = {'en': 'This example is {}.', # English
                        'de': 'Dieses beispiel ist {}.', # German
                        'es': 'Este ejemplo es {}.', # Spanish
                        'it': 'Questo esempio è {}.', # Italian
                        'ru': 'Этот пример {}.', # Russian
                        'no': 'Dette eksempelet er {}.'} # Norsk

def detect_lang(sequence, labels):
    DetectorFactory.seed = 0
    seq_lang = 'en'

    try:
        seq_lang = detect(sequence)
        lbl_lang = detect(labels)
    except:
        print("Language detection failed!",
              "Date:{}, Sequence:{}, Labels:{}".format(
                  str(datetime.datetime.now()),
                  labels))

    if seq_lang != lbl_lang:
        print("Different languages detected for sequence and labels!",
              "Date:{}, Sequence:{}, Labels:{}, Sequence Language:{}, Label Language:{}".format(
                  str(datetime.datetime.now()),
                  sequence,
                  labels,
                  seq_lang,
                  lbl_lang))

    if seq_lang in models:
        print("Sequence Language detected:",
              "Date:{}, Sequence:{}, Sequence Language:{}".format(
                  str(datetime.datetime.now()),
                  sequence,
                  labels))
    else:
        print("Language not supported. Defaulting to English!",
              "Date:{}, Sequence:{}, Sequence Language:{}".format(
                  str(datetime.datetime.now()),
                  sequence,
                  seq_lang))
        seq_lang = 'en'

    return seq_lang


def sequence_to_classify(sequence, labels):
    label_clean = str(labels).split(",")

    lang = detect_lang(sequence, labels)
    classifier = pipeline("zero-shot-classification",
                          #hypothesis_template=hypothesis_templates[lang],
                          model=models[lang])
    response = classifier(sequence, label_clean, multi_class=True)

    predicted_labels = response['labels']
    predicted_scores = response['scores']
    clean_output = {idx: float(predicted_scores.pop(0)) for idx in predicted_labels}
    print("Date:{}, Sequence:{}, Labels: {}".format(
        str(datetime.datetime.now()),
        sequence,
        predicted_labels))

    return clean_output

example_text1="Folkehelseinstituttets mest optimistiske anslag er at alle voksne er ferdigvaksinert innen midten av september."
example_labels1="politikk,helse,sport,religion"
example_text2="Kutt smør i terninger, og la det temperere seg litt mens deigen elter. Ha hvetemel, sukker, gjær, salt og kardemomme i en bakebolle til kjøkkenmaskin. Bruker du fersk gjær kan du smuldre gjæren i bollen, eller røre den ut i melken. Alt vil ettehvert blande seg godt, så begge deler er like bra."
example_labels2="helse,sport,religion, mat"

iface = gr.Interface(
    title="Multilingual Multi-label Zero-shot Classification",
    description="Currently supported languages are English, German, Spanish, Italian, Russian, Norsk.",
    fn=sequence_to_classify,
    inputs=[gr.inputs.Textbox(lines=20,
        label="Please enter the text you would like to classify...",
        placeholder="Text here..."),
        gr.inputs.Textbox(lines=5,
        label="Possible candidate labels (separated by comma)...",
        placeholder="laLels here...")],
    outputs=gr.outputs.Label(num_top_classes=5),
    capture_session=True,
    #interpretation="default",
    examples=[
        [example_text1, example_labels1],
        [example_text2, example_labels2]
    ])
iface.launch()