"""
Austrian Dialect Text-to-Speech Interface for Hugging Face Spaces
Based on FastSpeech 2 with dialect embeddings

Models are downloaded from a private Hugging Face model repository at startup.
Set the Space variable FastSpeechHifiGANckpt to your repo ID (e.g. username/repo-name).
"""
import gradio as gr
import os
import spaces
from huggingface_hub import hf_hub_download
from InferenceInterfaces.InferenceFastSpeech2 import InferenceFastSpeech2

# ---------------------------------------------------------------------------
# Model repo configuration
# The Space variable "FastSpeechHifiGANckpt" holds the repo ID.
# Fallback to env var, then to a placeholder so the app can still import.
# ---------------------------------------------------------------------------
MODEL_REPO_ID = os.getenv("FastSpeechHifiGANckpt", os.getenv("MODEL_REPO_ID", ""))
MODEL_DIR = os.path.expanduser("~/.cache/austrian_tts_models")
os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs("audios", exist_ok=True)

# Files to download from the model repo and where they need to end up
# so that InferenceFastSpeech2 can find them under MODEL_DIR.
MODEL_FILES = {
    # repo filename -> local path relative to MODEL_DIR
    "best.pt":         "FastSpeech2_Austrian_vox107_wav2vec_espeak/best.pt",
    "best_HiFiGAN.pt": "HiFiGAN_aridialect/best.pt",
}


def download_models():
    """Download model checkpoints from HF Hub and lay out the Models/ tree."""
    if not MODEL_REPO_ID:
        raise RuntimeError(
            "No model repo ID found. Set the Space variable "
            "FastSpeechHifiGANckpt or the MODEL_REPO_ID env var."
        )

    os.makedirs(os.path.join(MODEL_DIR, "FastSpeech2_Austrian_vox107_wav2vec_espeak"), exist_ok=True)
    os.makedirs(os.path.join(MODEL_DIR, "HiFiGAN_aridialect"), exist_ok=True)

    for repo_file, local_rel in MODEL_FILES.items():
        target = os.path.join(MODEL_DIR, local_rel)
        if os.path.exists(target):
            print(f"Already cached: {target}")
            continue
        print(f"Downloading {repo_file} from {MODEL_REPO_ID} ...")
        cached = hf_hub_download(
            repo_id=MODEL_REPO_ID,
            filename=repo_file,
            repo_type="model",
            token=os.environ["TTS_access_model"],
        )
        os.symlink(cached, target)
        print(f"  -> {target}")

    print("All models ready.")


# ---------------------------------------------------------------------------
# UI constants (identical to set_up_interface_tab_lang.py)
# ---------------------------------------------------------------------------
HEADER = (
    '<center>'
    '<a href="https://www.ofai.at/">'
    '</a>'
    '</center>\n'
)

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_DESCRIPTION = (
    HEADER +
    "<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>: "
    "<a href='https://www.isca-archive.org/sigul_2023/gutscher23_sigul.html'>https://www.isca-archive.org/sigul_2023/gutscher23_sigul.html</a>"
    "<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 temporally 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"

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"],
}


# ---------------------------------------------------------------------------
# Embedding / speaker helpers
# ---------------------------------------------------------------------------
def get_language_embedding(lang_emb_avg: str, interpolation: str) -> str:
    if lang_emb_avg == "Standard" and interpolation == "No interpolation":
        return 'Preprocessing/wav2vec_embeddings/at_emb_vox107_wav2vec.pt'
    elif interpolation == "0.5 Standard, 0.5 Viennese":
        return "Preprocessing/wav2vec_embeddings/at_to_vd_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Standard, 0.5 Goisern":
        return "Preprocessing/wav2vec_embeddings/at_to_goi_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Standard, 0.5 Innervillgraten":
        return "Preprocessing/wav2vec_embeddings/at_to_ivg_emb_vox107_0_50.pt"
    elif lang_emb_avg == "Viennese" and interpolation == "No interpolation":
        return "Preprocessing/wav2vec_embeddings/vd_emb_vox107_wav2vec.pt"
    elif interpolation == "0.5 Viennese, 0.5 Goisern":
        return "Preprocessing/wav2vec_embeddings/vd_to_goi_emb_vox107_0_50.pt"
    elif interpolation == "0.5 Viennese, 0.5 Innervillgraten":
        return "Preprocessing/wav2vec_embeddings/vd_to_ivg_emb_vox107_0_50.pt"
    elif lang_emb_avg == "Goisern" and interpolation == "No interpolation":
        return "Preprocessing/wav2vec_embeddings/goi_emb_vox107_wav2vec.pt"
    elif interpolation == "0.5 Goisern, 0.5 Innervillgraten":
        return "Preprocessing/wav2vec_embeddings/goi_to_ivg_emb_vox107_0_50.pt"
    elif lang_emb_avg == "Innervillgraten" and interpolation == "No interpolation":
        return "Preprocessing/wav2vec_embeddings/ivg_emb_vox107_wav2vec.pt"
    else:
        return "Preprocessing/wav2vec_embeddings/at_emb_vox107_wav2vec.pt"


def get_speaker_embedding(speaker: str, lang_emb_avg: str) -> str:
    speaker_map = {
        ("Austrian Standard (m)", "Standard"): "Utility/example_wavs/spo_at_berlin_001.wav",
        ("Austrian Standard (m)", None): "Utility/example_wavs/spo_vd_vdftw_001278.wav",
        ("Viennese (f)", None): "Utility/example_wavs/hga_vd_berlin_003.wav",
        ("Viennese (m)", "Standard"): "Utility/example_wavs/hpo_at_nordwind_005.wav",
        ("Viennese (m)", None): "Utility/example_wavs/hpo_vd_wean_0002.wav",
        ("young Viennese (f)", None): "Utility/example_wavs/joe_vd_fritz_048.wav",
        ("Goisern (f1)", "Standard"): "Utility/example_wavs/gun_at_berlin_001.wav",
        ("Goisern (f1)", None): "Utility/example_wavs/gun_goi_goi_001.wav",
        ("Goisern (m1)", "Standard"): "Utility/example_wavs/hoi_at_berlin_001.wav",
        ("Goisern (m1)", None): "Utility/example_wavs/hoi_goi_goi_001.wav",
        ("Goisern (f2)", None): "Utility/example_wavs/tfe_goi_goi_001.wav",
        ("Goisern (m2)", None): "Utility/example_wavs/wke_goi_goi_001.wav",
        ("Innervillgraten (f1)", "Standard"): "Utility/example_wavs/bsc_at_berlin_001.wav",
        ("Innervillgraten (f1)", None): "Utility/example_wavs/bsc_ivg_ivg_009.wav",
        ("Innervillgraten (m1)", "Standard"): "Utility/example_wavs/csc_at_berlin_001.wav",
        ("Innervillgraten (m1)", None): "Utility/example_wavs/csc_ivg_ivg_009.wav",
        ("Innervillgraten (m2)", None): "Utility/example_wavs/lsc_ivg_ivg_009.wav",
        ("Innervillgraten (f2)", None): "Utility/example_wavs/psc_ivg_ivg_009.wav",
    }

    key = (speaker, lang_emb_avg if lang_emb_avg == "Standard" else None)
    return speaker_map.get(key, "Utility/example_wavs/hpo_vd_wean_0002.wav")


# ---------------------------------------------------------------------------
# Synthesis
# ---------------------------------------------------------------------------
def read_texts(
    model_id: str,
    sentence: str,
    filename: str,
    speaker: str = None,
    device: str = "cuda",
    language: str = "de",
    input_is_phones: bool = False,
    lang_emb_avg: str = None,
    interpolation: str = "No interpolation",
):
    lang_emb = get_language_embedding(lang_emb_avg, interpolation)

    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)

    speaker_emb = get_speaker_embedding(speaker, lang_emb_avg)
    tts.set_utterance_embedding(speaker_emb)
    tts.set_language_embedding(lang_emb, use_avg=True)

    if isinstance(sentence, str):
        sentence = [sentence]

    tts.read_to_file(text_list=sentence, file_location=filename)
    del tts

@spaces.GPU
def synthesize(input_txt: str, speaker: str, lang_emb_avg: str, interpolation: str) -> str:
    model_id = "Austrian_vox107_wav2vec_espeak"
    filename = "audios/output.wav"

    read_texts(
        model_id=model_id,
        sentence=input_txt,
        filename=filename,
        language="de",
        speaker=speaker,
        lang_emb_avg=lang_emb_avg,
        input_is_phones=False,
        interpolation=interpolation,
    )
    return filename


def tts_demo_fn(text: str, speaker: str = DEFAULT_SPK, lang_emb_avg: str = DEFAULT_LANG_EMB, interpolation: str = DEFAULT_INTERPOLATION) -> str:
    if len(text) == 0:
        raise ValueError('Empty text. Please enter some text to synthesize.')
    tts_audio = synthesize(text, speaker, lang_emb_avg, interpolation)
    return tts_audio


# ---------------------------------------------------------------------------
# Gradio interface
# ---------------------------------------------------------------------------
def tts_demo():
    interface = gr.Interface(
        fn=tts_demo_fn,
        inputs=[
            gr.Textbox(label='Text', placeholder='Enter German text here...'),
            gr.Radio(
                SPEAKERS,
                value=DEFAULT_SPK,
                label='Speaker',
            ),
            gr.Radio(
                DIALECT_EMBEDDINGS,
                value=DEFAULT_LANG_EMB,
                label='Standard/Dialect Embedding',
            ),
            gr.Dropdown(
                choices=list(INTERPOLATION_OPTIONS.keys()),
                value=DEFAULT_INTERPOLATION,
                label='Interpolation (overrides Standard/Dialect Embedding)',
            ),
        ],
        outputs=[
            gr.Audio(label='Audio Output'),
        ],
        title=TTS_TITLE,
        description=TTS_DESCRIPTION,
        article=ARTICLE,
        examples=[
            ["Der Nordwind und die Sonne stritten sich, wer von ihnen wohl der Stärkere wäre.", DEFAULT_SPK, DEFAULT_LANG_EMB, DEFAULT_INTERPOLATION],
            ["Grüß Gott, wie geht es dir?", "Viennese (f)", "Viennese", DEFAULT_INTERPOLATION],
        ],
    )
    return interface


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


if __name__ == '__main__':
    download_models()
    interface = demo()
    interface.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )