import gradio as gr
import os
from InferenceInterfaces.InferenceFastSpeech2 import InferenceFastSpeech2
import argparse

# If the HF Space variable or env var is set, download models from Hub
MODEL_REPO_ID = os.getenv("FastSpeechHifiGANckpt", os.getenv("MODEL_REPO_ID", ""))
MODEL_DIR = os.path.expanduser("~/.cache/austrian_tts_models") if MODEL_REPO_ID else "Models"

HF_MODEL_FILES = {
    "best.pt":         "FastSpeech2_Austrian_vox107_wav2vec_espeak/best.pt",
    "best_HiFiGAN.pt": "HiFiGAN_aridialect/best.pt",
}


def download_hf_models():
    """Download model checkpoints from HF Hub if MODEL_REPO_ID is set."""
    if not MODEL_REPO_ID:
        return
    from huggingface_hub import hf_hub_download
    os.makedirs(MODEL_DIR, exist_ok=True)
    for repo_file, local_rel in HF_MODEL_FILES.items():
        target = os.path.join(MODEL_DIR, local_rel)
        if os.path.exists(target):
            print(f"Already cached: {target}")
            continue
        os.makedirs(os.path.dirname(target), exist_ok=True)
        print(f"Downloading {repo_file} from {MODEL_REPO_ID} ...")
        cached = hf_hub_download(repo_id=MODEL_REPO_ID, filename=repo_file, repo_type="model")
        os.symlink(cached, target)
        print(f"  -> {target}")

HEADER = (
    '<center>'
    '<a href="https://www.ofai.at/">'
    '</a>'
    '</center>\n'
)

TTS_TITLE = 'Text-to-Speech'
TTS_DESCRIPTION = (
    HEADER +
    'Type in a text and read it with the chosen voice'
)
TTS_TITLE = 'Text-to-Speech <br/> For further details see: <a href="https://www.researchgate.net/publication/373359033_Neural_Speech_Synthesis_for_Austrian_Dialects_with_Standard_German_Grapheme-to-Phoneme_Conversion_and_Dialect_Embeddings#fullTextFileContent">Paper</a>'
TTS_MODELS = (
    'Text'
)

TTS_DESCRIPTION = (
    HEADER +
    #make footnote in description for FastSpeech 2
    "<style>"
    ".center-text {"
    "  text-align: center;"
    "}"
    "</style>"
    "<div class='center-text'>"
    "A demo interface for Text-to-Speech synthesis. The implementation is based on FastSpeech 2<sup>1</sup>, using an adaptation of IMS-Toucan<sup>2</sup> with standard/dialect embeddings<sup>3</sup>."
    "<br/>"
    "For further information see our publication<sup>4</sup> : https://www.isca-archive.org/sigul_2023/gutscher23_sigul.html"
    "<br/>"
    "Geben Sie einen deutschen Text ein, wählen Sie eine:n Sprecher:in sowie einen Standard/Dialekt aus und lassen Sie den Text vorlesen."
    "<br/>"
    "Type in a German text, select a speaker, and select a standard/dialect embedding to read the text."
    "</div>"
)

REFERENCES = (
    "[1] Y. Ren, C. Hu, X. Tan, T. Qin, S. Zhao, Z. Zhao, and T.-Y. Liu, “Fastspeech 2: Fast and high-quality end-to-end text to speech,” in *ICLR 2021 - 9th International Conference on Learning Representations*, 2021. <br/>"
    "[2] https://github.com/DigitalPhonetics/IMS-Toucan <br/>"
    "[3] https://huggingface.co/TalTechNLP/voxlingua107-xls-r-300m-wav2vec <br/>"
    "[4] L. Gutscher, M. Pucher, and V. García, “Neural Speech Synthesis for Austrian Dialects with Standard German Grapheme-to-Phoneme Conversion and Dialect Embeddings,” in *Proc. 2nd Annual Meeting of the ELRA/ISCA SIG on Under-resourced Languages (SIGUL 2023)*, 2023. <br/>"
    "    https://www.isca-archive.org/sigul_2023/gutscher23_sigul.html <br/>"
)

FOOTER = (
    'Security Notice: Text wird beim klicken auf "Senden" temporär gespeichert. <br/>'
    'Data is stored temporarly when clicking "Absenden".<br/>\n'
    '<br/>\n'
    '<a href="mailto:lorenz.gutscher@ofai.at">Lorenz Gutscher</a> and '
    '<a href="mailto:michael.pucher@ofai.at">Michael Pucher</a><br/>\n'
)

ARTICLE = (REFERENCES + FOOTER)
SPEAKERS = ["Austrian Standard (m)", "Viennese (f)", "Viennese (m)", "young Viennese (f)", "Goisern (f1)", "Goisern (f2)", "Goisern (m1)", "Goisern (m2)", "Innervillgraten (f1)", "Innervillgraten (f2)", "Innervillgraten (m1)", "Innervillgraten (m2)"]
DIALECT_EMBEDDINGS = ["Standard", "Viennese", "Goisern", "Innervillgraten"]
default_spk = "Viennese (m)"
default_lang_emb = "Viennese"
default_interpolation = "No interpolation"

def synthesize(input_txt: str, speaker: str, lang_emb_avg: str, interpolation: str):
    model_id="Austrian_vox107_wav2vec_espeak"
    filename="audios/german_for_HPO_2Austrian_extended_phonemes_Empfehlung.wav"
    sentence=input_txt
    print(lang_emb_avg)
    print(speaker)
    read_texts(model_id=model_id, sentence=sentence, filename=filename, language="de", speaker=speaker, lang_emb_avg=lang_emb_avg, input_is_phones=False, interpolation=interpolation)
    return filename

def read_texts(model_id, sentence, filename, speaker=None, device="cpu", language="", input_is_phones=None, lang_emb_avg=None, interpolation="No interpolation"):
    if lang_emb_avg == "Standard" and interpolation == "No interpolation":
        lang_emb = 'Preprocessing/wav2vec_embeddings/at_emb_vox107_wav2vec.pt'
    elif interpolation == "0.5 Standard, 0.5 Viennese":
        lang_emb = "Preprocessing/wav2vec_embeddings/at_to_vd_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Standard, 0.5 Goisern":
        lang_emb = "Preprocessing/wav2vec_embeddings/at_to_goi_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Standard, 0.5 Innervillgraten":
        lang_emb = "Preprocessing/wav2vec_embeddings/at_to_ivg_emb_vox107_0_50.pt"
    # Viennese
    elif lang_emb_avg == "Viennese" and interpolation == "No interpolation":
        lang_emb = "Preprocessing/wav2vec_embeddings/vd_emb_vox107_wav2vec.pt"
    elif interpolation == "0.5 Viennese, 0.5 Goisern":
        lang_emb = "Preprocessing/wav2vec_embeddings/vd_to_goi_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Viennese, 0.5 Innervillgraten":
        lang_emb = "Preprocessing/wav2vec_embeddings/vd_to_ivg_emb_vox107_0_50.pt"
    # Goisern
    elif lang_emb_avg == "Goisern" and interpolation == "No interpolation":
        lang_emb = "Preprocessing/wav2vec_embeddings/goi_emb_vox107_wav2vec.pt"
    elif interpolation == "0.5 Goisern, 0.5 Innervillgraten":
        lang_emb = "Preprocessing/wav2vec_embeddings/goi_to_ivg_emb_vox107_0_50.pt"
    # Innervillgraten
    elif lang_emb_avg == "Innervillgraten" and interpolation == "No interpolation":
        lang_emb = "Preprocessing/wav2vec_embeddings/ivg_emb_vox107_wav2vec.pt"
    print("speaker in read texts is: " + str(speaker))
    tts = InferenceFastSpeech2(device=device, model_name=model_id, language=language, Avocodo=False, model_dir=MODEL_DIR)
    tts.set_language(language)
    tts.set_phoneme_input(input_is_phones)
    print("model_id is: " + model_id)
    print("lang_emb is: " + str(lang_emb_avg))
    use_avg = True

    if speaker == "Austrian Standard (m)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to spo_at_berlin_001:")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_utterance_embedding("Utility/example_wavs/spo_at_berlin_001.wav")
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Austrian Standard (m)":
        print("utterance embedding is set to spo_vd_berlin_001:")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_utterance_embedding("Utility/example_wavs/spo_vd_vdftw_001278.wav")
        tts.set_language_embedding(lang_emb,use_avg=use_avg)

    elif speaker== "Viennese (f)":
        print("utterance embedding is set to hga_vd_berlin_003:")
        tts.set_utterance_embedding("Utility/example_wavs/hga_vd_berlin_003.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Viennese (m)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to hpo_at_wean_0002:")
        tts.set_utterance_embedding("Utility/example_wavs/hpo_at_nordwind_005.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Viennese (m)":
        print("utterance embedding is set to hpo_vd_wean_0002:")
        tts.set_utterance_embedding("Utility/example_wavs/hpo_vd_wean_0002.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "young Viennese (f)":
        print("utterance embedding is set to joe_vd_fritz_048:")
        tts.set_utterance_embedding("Utility/example_wavs/joe_vd_fritz_048.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)

    elif speaker == "Goisern (f1)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to gun_at_berlin_001:")
        tts.set_utterance_embedding("Utility/example_wavs/gun_at_berlin_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Goisern (f1)":
        print("utterance embedding is set to gun_goi_goi_001:")
        tts.set_utterance_embedding("Utility/example_wavs/gun_goi_goi_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Goisern (m1)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to hoi_at_berlin_001:")
        tts.set_utterance_embedding("Utility/example_wavs/hoi_at_berlin_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Goisern (m1)":
        print("utterance embedding is set to hoi_goi_goi_051:")
        tts.set_utterance_embedding("Utility/example_wavs/hoi_goi_goi_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Goisern (f2)":
        print("utterance embedding is set to tfe_goi_goi_051:")
        tts.set_utterance_embedding("Utility/example_wavs/tfe_goi_goi_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker == "Goisern (m2)":
        print("utterance embedding is set to wke_goi_goi_051:")
        tts.set_utterance_embedding("Utility/example_wavs/wke_goi_goi_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)

    elif speaker =="Innervillgraten (f1)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to bsc_at_berlin_001:")
        tts.set_utterance_embedding("Utility/example_wavs/bsc_at_berlin_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker =="Innervillgraten (f1)":
        print("utterance embedding is set to bsc_ivg_ivg_009:")
        tts.set_utterance_embedding("Utility/example_wavs/bsc_ivg_ivg_009.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker =="Innervillgraten (m1)" and lang_emb_avg == "Standard":
        print("utterance embedding is set to csc_at_berlin_001:")
        tts.set_utterance_embedding("Utility/example_wavs/csc_at_berlin_001.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker =="Innervillgraten (m1)":
        print("utterance embedding is set to csc_ivg_ivg_009:")
        tts.set_utterance_embedding("Utility/example_wavs/csc_ivg_ivg_009.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker =="Innervillgraten (m2)":
        print("utterance embedding is set to lsc_ivg_ivg_009:")
        tts.set_utterance_embedding("Utility/example_wavs/lsc_ivg_ivg_009.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)
    elif speaker =="Innervillgraten (f2)":
        print("utterance embedding is set to psc_ivg_ivg_009:")
        tts.set_utterance_embedding("Utility/example_wavs/psc_ivg_ivg_009.wav")
        #tts.set_utterance_embedding(utt_emb)
        tts.set_language_embedding(lang_emb,use_avg=use_avg)

    if type(sentence) == str:
        sentence = [sentence]
    tts.read_to_file(text_list=sentence, file_location=filename)
    del tts

# dropdown options
INTERPOLATION_OPTIONS = {
    "No interpolation": ["Standard", "Viennese", "Goisern", "Innervillgraten"],
    "0.5 Standard, 0.5 Viennese": ["Standard", "Viennese"],
    "0.5 Standard, 0.5 Goisern": ["Standard", "Goisern"],
    "0.5 Standard, 0.5 Innervillgraten": ["Standard", "Innervillgraten"],
    "0.5 Viennese, 0.5 Goisern": ["Viennese", "Goisern"],
    "0.5 Viennese, 0.5 Innervillgraten": ["Viennese", "Innervillgraten"],
    "0.5 Goisern, 0.5 Innervillgraten": ["Goisern", "Innervillgraten"],
}

def tts_demo_fn(text, speaker='Wiener Dialekt (m)', lang_emb_avg="Standard", interpolation="No interpolation"):
    if len(text) == 0:
        raise ValueError('Empty text.')
    tts_audio = synthesize(text, speaker, lang_emb_avg, interpolation)
    return tts_audio

def load_models():
    download_hf_models()
    print("load_models (model_dir=" + MODEL_DIR + ")")

def tts_demo():
    interface = gr.Interface(
        fn=tts_demo_fn,
        inputs=[
            gr.Textbox(label='Text'),
            gr.Radio(
                SPEAKERS,
                value=default_spk,
                label='Speaker'
            ),
            gr.Radio(
                DIALECT_EMBEDDINGS,
                value=default_lang_emb,
                label='Standard/Dialect Embedding',
            ),
            gr.Dropdown(
                choices=INTERPOLATION_OPTIONS,
                value=default_interpolation,
                label='Interpolation (overrides Standard/Dialect Embedding!))',
            ),
        ],
        outputs=[
            gr.Audio(label='Audio'),
        ],
        title=TTS_TITLE,
        description=TTS_DESCRIPTION,
        article=ARTICLE,
    )
    return interface

def demo():
    interface = gr.TabbedInterface(
        interface_list=[tts_demo()],
        tab_names=['TTS'],
    )
    return interface

def main():
    """Main entry point of the program."""
    parser = argparse.ArgumentParser(description='Sprechen wir.')
    parser.add_argument(
        '--queue', action='store_true',
        help=(
            'Schreibens was!'
        )
    )
    parser.add_argument(
        '--threads', type=int, default=40,
        help='Allow up to `threads` to be processed in parallel.'
    )
    parser.add_argument(
        '--share', action='store_true',
        help='Whether to create a publicly shareable link for the interface.'
    )
    parser.add_argument(
        '--host',
        help=(
            'Host to listen connections into. '
            'Set to "0.0.0.0" to make it accessible in local network.'
        )
    )
    parser.add_argument(
        '--port',
        help='Start gradio app on this port (if available).'
    )
    parser.add_argument(
        '--debug', action='store_true',
        help='If True, blocks the main thread from running.'
    )
    parser.add_argument(
        '--nolock', action='store_true',
        help=(
            'If True, the interface will block the main thread while the '
            'server is running.'
        )
    )
    args = parser.parse_args()
    load_models()
    interface = demo()
    if args.queue:
        interface = interface.queue()
    interface.launch(   server_name="127.0.0.1",
                        server_port=8080,
                        debug=False,
                        share=False
                    )

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("exit")
        gr.close_all()
        exit()