# -*- 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 = [["wala manis a-ttidun?", "Tachelhit"],
            ["ġ-iḍ-an ġ-ilul-umsiggel, illas lḥal s-umdlu isemmiḍn.", "Tachelhit"]]

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 🗣️",
    allow_flagging="manual",
    flagging_options=['error', 'bad-quality', 'wrong-pronounciation'],
)
iface.launch()