"""
File: app.py

Description: Translate text...

Author: Didier Guillevic
Date: 2024-09-07
"""
import spaces
import gradio as gr
import langdetect
from typing import List

import logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

from deep_translator import GoogleTranslator
from model_spacy import nlp_xx

import model_translation as translation


def detect_language(text):
    lang = langdetect.detect(text)
    return lang


def build_text_chunks(text, src_lang, sents_per_chunk):
    """
    Given a text:
    - Split the text into sentences.
    - Build text chunks:
        - Consider up to sents_per_chunk
        - Ensure that we do not exceed translation.max_words_per_chunk
    """
    # Split text into sentences...
    sentences = [
        sent.text.strip() for sent in nlp_xx(text).sents if sent.text.strip()]
    logger.info(f"LANG: {src_lang}, TEXT: {text[:20]}, NB_SENTS: {len(sentences)}")
    
    # Create text chunks of N sentences
    chunks = []
    chunk = ''
    chunk_nb_sentences = 0
    chunk_nb_words = 0

    for i in range(0, len(sentences)):
        # Get sentence
        sent = sentences[i]
        sent_nb_words = len(sent.split())

        # If chunk already 'full', save chunk, start new chunk
        if (
                (chunk_nb_words + sent_nb_words > translation.max_words_per_chunk) or
                (chunk_nb_sentences + 1 > sents_per_chunk)
           ):
            chunks.append(chunk)
            chunk = ''
            chunk_nb_sentences = 0
            chunk_nb_words = 0
        
        # Append sentence to current chunk. One sentence per line.
        chunk = (chunk + '\n' + sent) if chunk else sent
        chunk_nb_sentences += 1
        chunk_nb_words += sent_nb_words

    # Append last chunk
    if chunk:
        chunks.append(chunk)

    return chunks


def translate_with_Helsinki(
        chunks, src_lang, tgt_lang, input_max_length, output_max_length) -> str:
    """Translate the chunks with the Helsinki model
    """
    if src_lang not in translation.src_langs:
        return (
            f"ISSUE: currently no model for language '{src_lang}'. "
             "If wrong language, please specify language."
        )
    logger.info(f"LANG: {src_lang}, TEXT: {chunks[0][:50]}...")
    tokenizer, model = translation.get_tokenizer_model_for_src_lang(src_lang)

    translated_chunks = []
    for chunk in chunks:
        # NOTE: The 'fa' (Persian) model has multiple target languages to choose from.
        # We need to specifiy the desired languages among: fra ita por ron spa
        #   https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-fa-itc
        # Prepend text with >>fra<< in order to translate in French.
        if src_lang == 'fa':
            chunk = ">>fra<< " + chunk

        inputs = tokenizer(
            chunk, return_tensors="pt", max_length=input_max_length,
            truncation=True, padding="longest").to(model.device)
        outputs = model.generate(**inputs, max_length=output_max_length)
        translated_chunk = tokenizer.batch_decode(
            outputs, skip_special_tokens=True)[0]
        #logger.info(f"Text: {chunk}")
        #logger.info(f"Translation: {translated_chunk}")
        translated_chunks.append(translated_chunk)

    return '\n'.join(translated_chunks)


@spaces.GPU
def translate_with_m2m100(
        chunks: List[str],
        src_lang: str,
        tgt_lang: str) -> str:
    """Translate with the m2m100 model
    """
    m2m100 = translation.ModelM2M100()
    m2m100.tokenizer.src_lang = src_lang

    translated_chunks = []
    for chunk in chunks:
        input_ids = m2m100.tokenizer(
            chunk, return_tensors="pt").input_ids.to(m2m100.model.device)
        outputs = m2m100.model.generate(
            input_ids=input_ids,
            forced_bos_token_id=m2m100.tokenizer.get_lang_id(tgt_lang))
        translated_chunk = m2m100.tokenizer.batch_decode(
            outputs, skip_special_tokens=True)[0]
        translated_chunks.append(translated_chunk)
    
    return '\n'.join(translated_chunks)


@spaces.GPU
def translate_with_MADLAD(
        chunks: List[str],
        tgt_lang: str,
        input_max_length: int=512,
        output_max_length: int=512) -> str:
    """Translate with Google MADLAD model
    """
    madlad = translation.ModelMADLAD()

    for chunk in chunks:
        input_text = f"<2{tgt_lang}> {text}"
        #logger.info(f" Translating: {input_text[:30]}")
        input_ids = madlad.tokenizer(
            input_text, return_tensors="pt", max_length=input_max_length,
            truncation=True, padding="longest").input_ids.to(madlad.model.device)
        outputs = madlad.model.generate(
            input_ids=input_ids, max_length=output_max_length)
        translated_chunk = madlad.tokenizer.decode(
            outputs[0], skip_special_tokens=True)
        translated_chunks.append(translated_chunk)
    
    return '\n'.join(translated_chunks)


def translate_text(
        text: str,
        src_lang: str=None,
        sents_per_chunk: int=5,
        input_max_length: int=512,
        output_max_length: int=512):
    """
    Translate the given text into English (default "easy" language)  
    """
    src_lang = src_lang if (src_lang and src_lang != "auto") else detect_language(text)
    tgt_lang = 'en' # Default "easy" language

    chunks = build_text_chunks(text, src_lang, sents_per_chunk)

    translated_text_Helsinki = translate_with_Helsinki(
        chunks, src_lang, tgt_lang, input_max_length, output_max_length)
    translated_text_m2m100 = translate_with_m2m100(chunks, src_lang, tgt_lang)
    translated_text_MADLAD = translate_with_MADLAD(chunks, tgt_lang)
    translated_text_google_translate = GoogleTranslator(
        source='auto', target='en').translate(text=text)

    return (
        translated_text_Helsinki,
        translated_text_m2m100,
        translated_text_MADLAD,
        translated_text_google_translate
    )


#
# User interface
#
with gr.Blocks() as demo:

    gr.Markdown("""
        ## Text translation v0.0.2 (small paragraph, multilingual)
    """)
    input_text = gr.Textbox(
        lines=15,
        placeholder="Enter text to translate",
        label="Text to translate",
        render=False
    )
    output_text_Helsinki = gr.Textbox(
        lines=6,
        label="Bilingual translation model (Helsinki NLP)",
        render=False
    )
    output_text_m2m100 = gr.Textbox(
        lines=6,
        label="Facebook m2m100 translation model (**small**)",
        render=False
    )
    output_text_MADLAD = gr.Textbox(
        lines=6,
        label="Google MADLAD translation model (**small**)",
        render=False
    )
    output_text_google_translate = gr.Textbox(
        lines=6,
        label="Google Translate",
        render=False
    )

    # Extra (additional) input parameters
    sentences_per_chunk = gr.Slider(
        minimum=1, maximum=10, value=5, step=1,
        label="nb sentences per context",
        render=False
    )
    src_lang = gr.Radio(
        choices=["auto", "ar", "en", "fa", "fr", "he", "zh"], value="auto",
        label="Source language",
        render=False
    )

    # Examples
    examples = [
        ["ریچارد مور، رئیس سازمان مخفی اطلاعاتی بریتانیا (ام‌آی‌۶) در دیدار ویلیام برنز، رئیس سازمان اطلاعات مرکزی آمریکا (سیا) گفت همچنان احتمال اقدام ایران علیه اسرائیل در واکنش به ترور اسماعیل هنیه، رهبر حماس وجود دارد. آقای برنز نیز در این دیدار فاش کرد که در سال اول جنگ اوکراین، «خطر واقعی» وجود داشت که روسیه به استفاده از «تسلیحات هسته‌ای تاکتیکی» متوسل شود. این دو مقام امنیتی هشدار دادند که «نظم جهانی» از زمان جنگ سرد تا کنون تا این حد «در معرض تهدید» نبوده است.", "fa"],
        ["Clément Delangue est, avec Julien Chaumond et Thomas Wolf, l’un des trois Français cofondateurs de Hugging Face, une start-up d’intelligence artificielle (IA) de premier plan. Valorisée à 4,2 milliards d’euros après avoir levé près de 450 millions d’euros depuis sa création en 2016, cette société de droit américain est connue comme la plate-forme de référence où développeurs et entreprises publient des outils et des modèles pour faire de l’IA en open source, c’est-à-dire accessible gratuitement et modifiable.", "fr"],
        ["يُعد تفشي مرض جدري القردة قضية صحية عالمية خطيرة، ومن المهم محاولة منع انتشاره للحفاظ على سلامة الناس وتجنب العدوى. د. صموئيل بولاند، مدير الحوادث الخاصة بمرض الجدري في المكتب الإقليمي لمنظمة الصحة العالمية في أفريقيا، يتحدث من كينشاسا في جمهورية الكونغو الديمقراطية، ولديه بعض النصائح البسيطة التي يمكن للناس اتباعها لتقليل خطر انتشار المرض.", "ar"],
        ["張先生稱，奇瑞已經凖備在西班牙生產汽車，並決心採取「本地化」的方式進入歐洲市場。此外，他也否認該公司的出口受益於不公平補貼。奇瑞成立於1997年，是中國最大的汽車公司之一。它已經是中國最大的汽車出口商，並且制定了進一步擴張的野心勃勃的計劃。", "zh"],
        ["ברוכה הבאה, קיטי: בית הקפה החדש בלוס אנג'לס החתולה האהובה והחברים שלה מקבלים בית קפה משלהם בשדרות יוניברסל סיטי, שם תוכלו למצוא מגוון של פינוקים מתוקים – החל ממשקאות ועד עוגות", "he"],
    ]
    
    gr.Interface(
        fn=translate_text,
        inputs=[input_text, src_lang,],
        outputs=[
            output_text_Helsinki,
            output_text_m2m100,
            output_text_MADLAD,
            output_text_google_translate,
        ],
        additional_inputs=[sentences_per_chunk,],
        #clear_btn=None,  # Unfortunately, clear_btn also reset the additional inputs. Hence disabling for now.
        allow_flagging="never",
        examples=examples,
        cache_examples=True
    )

    with gr.Accordion("Documentation", open=False):
        gr.Markdown("""
            - Models: serving Helsinki NLP, Facebook M2M100 and Google MADLAD models.
            - Basic: processing of long paragraph / document to be enhanced.
            - Most examples are copy/pasted from BBC news international web sites.
        """)

if __name__ == "__main__":
    demo.launch()