import os
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

en2ko = "KoJLabs/nllb-finetuned-en2ko"
ko2en = "KoJLabs/nllb-finetuned-ko2en"
style = "KoJLabs/bart-speech-style-converter"

en2ko_model = AutoModelForSeq2SeqLM.from_pretrained(en2ko)
ko2en_model = AutoModelForSeq2SeqLM.from_pretrained(ko2en)
style_model = AutoModelForSeq2SeqLM.from_pretrained(style)

en2ko_tokenizer = AutoTokenizer.from_pretrained(en2ko)
ko2en_tokenizer = AutoTokenizer.from_pretrained(ko2en)
style_tokenizer = AutoTokenizer.from_pretrained(style)

def translation(source, target, text):
    formats = {"English":"eng_Latn", "Korean":"kor_Hang"}
    src = formats[source]
    tgt = formats[target]

    if src == "eng_Latn":
        translator = pipeline(
            'translation',
            model=en2ko_model,
            tokenizer=en2ko_tokenizer,
            src_lang=src, 
            tgt_lang=tgt, 
        )

    if src == "kor_Hang":
        translator = pipeline(
            'translation',
            model=ko2en_model,
            tokenizer=ko2en_tokenizer,
            src_lang=src, 
            tgt_lang=tgt
        )

    output = translator(text)
    translated_text = output[0]['translation_text']

    return translated_text

def augmentation(text):
    ko2en_translator = pipeline(
            'translation',
            model=ko2en_model,
            tokenizer=ko2en_tokenizer,
            src_lang="kor_Hang", 
            tgt_lang="eng_Latn"
        )
    
    output = ko2en_translator(text)
    ko2en_text = output[0]['translation_text']
    
    en2ko_translator = pipeline(
            'translation',
            model=en2ko_model,
            tokenizer=en2ko_tokenizer,
            src_lang="eng_Latn", 
            tgt_lang="kor_Hang", 
        )
    
    output = en2ko_translator(ko2en_text)
    en2ko_text = output[0]['translation_text']

    return en2ko_text
    

def conversion(source, text):
    formats = {
                "formal":"문어체",
                "informal":"구어체",
                "android":"안드로이드",
                "azae":"아재",
                "chat":"채팅",
                "choding":"초등학생",
                "emoticon":"이모티콘",
                "enfp":"enfp",
                "gentle":"신사",
                "halbae":"할아버지",
                "halmae":"할머니",
                "joongding":"중학생",
                "king":"왕",
                "naruto":"나루토",
                "seonbi":"선비",
                "sosim":"소심한",
                "translator":"번역기",
            }
    style = formats[source]

    input_text = f"{style} 형식으로 변환:" + text

    converter = pipeline(
            'text2text-generation',
            model=style_model,
            tokenizer=style_tokenizer,
        )

    output = converter(input_text)
    generated_text = output[0]['generated_text']

    return generated_text



lang = ['English','Korean']
style = ['formal', 'informal', 'android', 'azae', 'chat', 'choding', 'emoticon', 'enfp', \
         'gentle', 'halbae', 'halmae', 'joongding', 'king', 'naruto', 'seonbi', 'sosim', 'translator']

    
translation_app = gr.Interface(
    fn=translation,
    inputs=[gr.inputs.Dropdown(choices=lang, label='Source Language'), gr.inputs.Dropdown(choices=lang, label='Target Language'), gr.inputs.Textbox(lines=5, label='Text to Translate')],
    outputs=[gr.outputs.Textbox(label='Translated Text')],
    title=title,
    enable_queue=True,
)

augmentation_app = gr.Interface(
    fn=augmentation,
    inputs=[gr.inputs.Textbox(lines=5, label='Korean Text to Augmentation')],
    outputs=[gr.outputs.Textbox(label='Augmented Text')],
    title=title,
    enable_queue=True,
)
    
conversion_app = gr.Interface(
    fn=conversion,
    inputs=[gr.inputs.Dropdown(choices=style, label='Speech Style'), gr.inputs.Textbox(lines=5, label='Text to style conversion')],
    outputs=[gr.outputs.Textbox(label='Converted Text')],
    title=title,
    enable_queue=True,
)

demo = gr.TabbedInterface([translation_app, augmentation_app, conversion_app], \
                          ["Translation", "Augmentation", "Speech Style conversion"],\
                          title = 'KoTAN Translator & Speech-style converter. \n \
                            🔥If you want to download as pip package, \
                            please visit our github. (https://github.com/KoJLabs/KoTAN) 🔥'
                            )

demo.launch()