import torch
import scipy
import os
import streamlit as st
from transformers import set_seed, pipeline
from transformers import VitsTokenizer, VitsModel
from datasets import load_dataset, Audio
from IPython.display import Audio as Aud
from src import *

from huggingface_hub import login
from dotenv import load_dotenv

#load_dotenv()
#HUGGINGFACE_KEY = os.environ.get("HUGGINGFACE_KEY")
#login(HUGGINGFACE_KEY)


########################
language_list = ['mos', 'fra', 'eng']


st.title("Demo: Automated Tools for Mooré Language")
tts, stt, trans, lid = st.tabs(["Text to speech", "Speech to text", "Translation", "Language ID"])

########################
with tts:
    
    tts_text = st.text_area(label = "Please enter your text here:", value="", placeholder="ne y wĩndga")

    tts_col1, tts_col2,  = st.columns(2)

    with tts_col1:
        tts_lang = st.selectbox('Language of text', (language_list), format_func = decode_iso)
    
    

    if st.button("Speak"):
        st.divider()
        with st.spinner(":rainbow[Synthesizing, please wait...]"):
            synth = synthesize_facebook(tts_text, tts_lang)
            st.audio(synth, sample_rate=16_000)

########################
with stt:

    stt_file = st.file_uploader("Please upload an audio file:", type=['mp3', 'm4a'], key = "stt_uploader")
    stt_lang = st.selectbox("Please select the language:" , (language_list), format_func = decode_iso)


    if st.button("Transcribe"):
        st.divider()
        with st.spinner("rainbow[Received your file, please wait while I process it...]"):
            stt = transcribe(stt_file, stt_lang)
            ":violet[The transcription is:]" 
            ':violet[ "' + stt + '"]'

########################
with trans:
    
    trans_text = st.text_area(label = "Please enter your translation text here:", value="", placeholder="ne y wĩndga")
    #trans_col1, trans_col2, trans_col3 = st.columns([.25, .25, .5])
    trans_col1, trans_col2 = st.columns(2)

    with trans_col1:
        src_lang = st.selectbox('Translate from:', (language_list), format_func = decode_iso)
    with trans_col2:
        target_lang = st.selectbox('Translate to:', (language_list), format_func = decode_iso, index=1)
    #with trans_col3:
    #    trans_model = st.selectbox("Translation model:",
    #                            ("Facebook (nllb-200-distilled-600M)", 
    #                             "Helsinki NLP (opus-mt-mos-en)", 
    #                             "Masakhane (m2m100_418m_mos_fr_news)")
    #                           )
    
    
    if st.button("Translate"):
        st.divider()
        with st.spinner(":rainbow[Translating from " + decode_iso(src_lang) + " into " + decode_iso(target_lang) + ", please wait...]"):
            translation = translate(trans_text, src_lang, target_lang) #, trans_model)
            translation

########################
with lid:
    langid_file = st.file_uploader("Please upload an audio file:", type=['mp3', 'm4a'], key = "lid_uploader")

    if st.button("Identify"):
        st.divider()
        with st.spinner(":rainbow[Received your file, please wait while I process it...]"):
            lang = identify_language(langid_file)
            lang = decode_iso(lang)
            ":violet[The detected language is " + lang + "]"


# supported colors: blue, green, orange, red, violet, gray/grey, rainbow.
# https://docs.streamlit.io/library/api-reference/text/st.markdown