from Imports import *
from Configuration import *
from Model_Loading import *
from Weights_Loading import model_state, vocoder_state
from Inference import inference

import gradio as gr

def tts_generate(
    text: str,
    eos_attn_threshold: float,
    max_steps: int,
    progress=gr.Progress(),
):
    text = (text or "")
    if not text:
        return None

    if vocoder_state is None or model_state is None:
        raise RuntimeError(
            "Models are not loaded. Make sure 'tacotron2_inference_ema.pkl' and "
            "'vocoder_inference.pkl' are present in the working directory."
        )

    _, wav_np = inference(
        text,
        model_state,
        vocoder_state,
        max_steps=int(max_steps),
        EOS_ATTN_THRESHOLD=float(eos_attn_threshold),
        progress=progress,
    )

    wav_np = np.asarray(wav_np, dtype=np.float32)
    max_val = np.max(np.abs(wav_np)) + 1e-8
    wav_np = wav_np / max_val

    return SAMPLE_RATE, wav_np


custom_css = """
body {
  background: radial-gradient(circle at top, #1e1b4b 0, #020617 55%, #020617 100%);
  color: #e5e7eb;
  font-family: system-ui, -apple-system, BlinkMacSystemFont, "SF Pro Text", sans-serif;
}
.tts-container {
  max-width: 520px;
  margin: 0 auto;
  padding: 32px 20px 40px;
}
.tts-card {
  background: radial-gradient(circle at top left, #1f2937 0, #020617 55%);
  border-radius: 24px;
  padding: 28px 24px 32px;
  box-shadow:
    0 30px 80px rgba(15, 23, 42, 0.9),
    0 0 0 1px rgba(148, 163, 184, 0.15);
  backdrop-filter: blur(24px);
}
.tts-title {
  text-align: center;
  font-weight: 800;
  font-size: 32px;
  letter-spacing: 0.16em;
  background: linear-gradient(90deg, #38bdf8, #a855f7, #fb7185);
  -webkit-background-clip: text;
  color: transparent;
  margin-bottom: 4px;
}
.tts-subtitle {
  text-align: center;
  font-size: 14px;
  color: #9ca3af;
  margin-bottom: 24px;
}
.tts-textbox .wrap.svelte-sr6c8e textarea {
  background: rgba(15, 23, 42, 0.9);
  border-radius: 18px;
  border: 1px solid rgba(148, 163, 184, 0.4);
  color: #e5e7eb;
}
.tts-generate-btn button {
  background: linear-gradient(90deg, #38bdf8, #a855f7, #fb7185);
  border-radius: 999px;
  border: none;
  color: white;
  font-weight: 600;
  letter-spacing: 0.04em;
  padding: 10px 24px;
  box-shadow: 0 14px 40px rgba(56, 189, 248, 0.35);
}
.tts-generate-btn button:hover {
  filter: brightness(1.05);
}
.tts-footer {
  margin-top: 18px;
  padding: 14px 16px;
  border-radius: 16px;
  background: linear-gradient(to right, rgba(15, 23, 42, 0.95), rgba(15, 23, 42, 0.85));
  border: 1px solid rgba(148, 163, 184, 0.25);
  font-size: 11px;
  color: #9ca3af;
  text-align: center;
}
.tts-footer strong {
  display: block;
  color: #e5e7eb;
  margin-bottom: 4px;
  text-transform: uppercase;
  letter-spacing: 0.15em;
  font-size: 10px;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_classes="tts-container"):
        with gr.Column(elem_classes="tts-card"):
            gr.HTML("<div class='tts-title'>TTSIE</div>")
            gr.HTML("<div class='tts-subtitle'>Text to Speech Synthesis</div>")

            text_in = gr.Textbox(
                placeholder="Enter your text here...",
                lines=4,
                label="",
                elem_classes="tts-textbox",
            )

            with gr.Accordion("Advanced", open=False):
                eos_attn_threshold = gr.Slider(
                    minimum=0.05,
                    maximum=0.95,
                    value=0.2,
                    step=0.01,
                    label="EOS attention threshold (higher = longer speech)",
                )
                max_steps = gr.Slider(
                    minimum=600,
                    maximum=20000,
                    value=1200,
                    step=100,
                    label="Max decoding steps (higher = allows longer audio)",
                )

            audio_out = gr.Audio(
                label="",
                type="numpy",
                interactive=False,
            )

            generate_btn = gr.Button("Generate", elem_classes="tts-generate-btn")

            generate_btn.click(
                fn=tts_generate,
                inputs=[text_in, eos_attn_threshold, max_steps],
                outputs=[audio_out],
            )

            gr.HTML(
                """
                <div class="tts-footer">
                  <strong>Powered by Tacotron 2 &amp; HiFi-GAN</strong>
                  Trained from scratch for educational purposes, focused on learning and experimentation.
                </div>
                """
            )

if __name__ == "__main__":
    demo.launch()