# -*- coding: utf-8 -*-
from ttsmms import TTS
import gradio as gr

ISO_CODES = {'Tachelhit': 'shi',
             'Tarifit (Latin script)': 'rif-script_latin',
             'Tarifit (Arabic script)': 'rif-script_arabic',
             'Taqbaylit': 'kab',
             'Tamasheq': 'taq',
             'Tamajaq, Tawallammat (Tifinagh script)': 'ttq-script_tifinagh'
             }

mapping = {'ɣ': 'ġ',
           'v': 'ġ',
           'c': 'š',
           'x': 'ḫ',
           'T': 'ṭ',
           'S': 'ṣ',
           'D': 'ḍ',
           'H': 'ḥ',
           'Z': 'ẓ',
           '3': 'ε',
           '7': 'ḥ',
           '9': 'q',
           'gh': 'ġ',
           'kh': 'ḫ'
          }

MODELS = {}

def tts(text, variant):
    variant_code = ISO_CODES[variant]
    if variant_code not in MODELS:
        MODELS[variant_code] = TTS(variant_code)
    model = MODELS[variant_code]

    if variant_code == 'shi':
        for key, value in mapping.items():
            text = text.replace(key, value)
    audio = model.synthesis(text)
    return (audio['sampling_rate'], audio['x'])

examples = [["Riɣ ad cceɣ s ufus kḍuɣ s inxar", "Tachelhit"],
            ["arraw n lhem yukr aɣ ihdumn nɣ", "Tachelhit"],
            ["wa tamɣart ma d ukan teskart ?", "Tachelhit"],
            ["ar d iṭṭar unẓar, ffuɣn d igḍaḍ, mmɣin d ijjign", "Tachelhit"],
            ["Egg lxir di timura, ad tafed di tiwwura.", "Tarifit (Latin script)"],
            ["Aqemmum iqnen ur ṯ-ttidfen izan.", "Tarifit (Latin script)"]]

description = "Text-to-speech for Tachelhit, Tarifit, Taqbaylit, Tamasheq and Tamajaq (Tawallammat)."

iface = gr.Interface(
    fn=tts,
    inputs=[
        gr.Textbox(
            label="Text",
            default="Text to synthesize.",
        ),
        gr.Dropdown(label="Variant", choices=list(ISO_CODES.keys()), value="Tachelhit")
    ],
    outputs=gr.Audio(label="Output", type="numpy"),
    examples=examples,
    title="🗣️ Tamazight Text-to-Speech with MMS (Massively Multilingual Speech) 🗣️",
    description=description,
    allow_flagging="manual",
    flagging_options=['error', 'bad-quality', 'wrong-pronounciation'],
)
iface.launch()