# -*- 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'
             }
VARIANTS = list(ISO_CODES.values())
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]
    audio = model.synthesis(text)
    return (audio['sampling_rate'], audio['x'])

examples = [["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.inputs.Textbox(
            label="Text",
            default="Text to synthesize.",
        ),
        gr.inputs.Dropdown(label="Variant", choices=list(ISO_CODES.keys()), default="Tachelhit")
    ],
    outputs=gr.outputs.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()