"""
Moodwave — Text + Audio Emotion Detector (themed demo)
========================================================
Matches the tactical/HUD aesthetic of radioweb.info.

Run with: python app.py
Then open the local URL Gradio prints.
"""

import os
import warnings

os.environ["TOKENIZERS_PARALLELISM"] = "false"
warnings.filterwarnings("ignore")

import gradio as gr
import pandas as pd
from transformers import pipeline as hf_pipeline

from spotify_client import search_tracks_for_mood, tracks_to_html, spotify_configured

# ── Config ──────────────────────────────────────────────────────────────
TARGET_LABELS = ["anger", "fear", "joy", "neutral", "sadness"]
TEXT_ZERO_SHOT_NAME = "SamLowe/roberta-base-go_emotions"

# Public Wav2Vec2 model fine-tuned for speech emotion recognition.
# Used as a fallback since the original university-GPU-trained weights
# are no longer available — mirrors the same "public model fallback"
# pattern already used for text mode above.
AUDIO_MODEL_NAME = "Dpngtm/wav2vec2-emotion-recognition"

# Late-fusion weights from the research project (text=0.2 / audio=0.8 gave
# the best F1 = 0.850 — see README "Multimodal" results). Exposed as sliders
# below so visitors can experiment, but these are the defaults.
DEFAULT_TEXT_WEIGHT = 0.2
DEFAULT_AUDIO_WEIGHT = 0.8

MOOD_MAP = {
    "admiration": "joy", "amusement": "joy",
    "anger": "anger", "annoyance": "anger",
    "approval": "joy", "caring": "joy",
    "confusion": "neutral", "curiosity": "neutral",
    "desire": "joy", "disappointment": "sadness",
    "disapproval": "anger", "disgust": "anger",
    "embarrassment": "sadness", "excitement": "joy",
    "fear": "fear", "gratitude": "joy",
    "grief": "sadness", "joy": "joy",
    "love": "joy", "nervousness": "fear",
    "optimism": "joy", "pride": "joy",
    "realization": "neutral", "relief": "joy",
    "remorse": "sadness", "sadness": "sadness",
    "surprise": "neutral", "neutral": "neutral",
}

# Maps the audio model's RAVDESS-style output labels onto the same
# 5-mood scheme used by the text model, so both tabs feel consistent.
AUDIO_MOOD_MAP = {
    "angry": "anger",
    "anger": "anger",
    "disgust": "anger",
    "fear": "fear",
    "fearful": "fear",
    "happy": "joy",
    "happiness": "joy",
    "joy": "joy",
    "neutral": "neutral",
    "calm": "neutral",
    "sad": "sadness",
    "sadness": "sadness",
    "surprise": "neutral",
    "surprised": "neutral",
}

MOOD_EMOJI = {
    "anger": "🔴",
    "fear": "🟣",
    "joy": "🟢",
    "neutral": "⚪",
    "sadness": "🔵",
}

print("Loading text model…")
_text_pipe = hf_pipeline(
    "text-classification",
    model=TEXT_ZERO_SHOT_NAME,
    top_k=None,
    truncation=True,
    max_length=128,
)
print("Text model loaded.")

print("Loading audio model…")
try:
    _audio_pipe = hf_pipeline(
        "audio-classification",
        model=AUDIO_MODEL_NAME,
    )
    print("Audio model loaded.")
except Exception as e:
    print(f"Audio model failed to load: {e}")
    _audio_pipe = None


def predict_text(text: str):
    text = (text or "").strip()
    if not text:
        return None
    raw = _text_pipe(text)[0]
    mood_scores = {m: 0.0 for m in TARGET_LABELS}
    for item in raw:
        mood = MOOD_MAP.get(item["label"])
        if mood:
            mood_scores[mood] += item["score"]
    total = sum(mood_scores.values()) or 1.0
    scores = {m: v / total for m, v in mood_scores.items()}
    best = max(scores, key=scores.get)
    return {"mood": best, "confidence": scores[best], "scores": scores}


def predict_audio(audio_path):
    if audio_path is None or _audio_pipe is None:
        return None
    raw = _audio_pipe(audio_path, top_k=None)
    mood_scores = {m: 0.0 for m in TARGET_LABELS}
    for item in raw:
        mood = AUDIO_MOOD_MAP.get(item["label"].lower())
        if mood:
            mood_scores[mood] += item["score"]
    total = sum(mood_scores.values()) or 1.0
    scores = {m: v / total for m, v in mood_scores.items()}
    best = max(scores, key=scores.get)
    return {"mood": best, "confidence": scores[best], "scores": scores}


def fuse_scores(text_scores, audio_scores, text_weight=DEFAULT_TEXT_WEIGHT, audio_weight=DEFAULT_AUDIO_WEIGHT):
    """
    Combine text + audio mood-probability vectors into one fused mood.
    Falls back to whichever single modality is available if only one
    was provided.
    """
    if text_scores and audio_scores:
        total_w = (text_weight + audio_weight) or 1.0
        fused = {
            m: (text_scores.get(m, 0.0) * text_weight + audio_scores.get(m, 0.0) * audio_weight) / total_w
            for m in TARGET_LABELS
        }
    elif text_scores:
        fused = dict(text_scores)
    elif audio_scores:
        fused = dict(audio_scores)
    else:
        return None

    best = max(fused, key=fused.get)
    return {"mood": best, "confidence": fused[best], "scores": fused}


def scores_to_bar_data(scores: dict):
    items = sorted(scores.items(), key=lambda x: -x[1])
    labels = [f"{MOOD_EMOJI.get(m, '')} {m.upper()}" for m, _ in items]
    values = [round(v * 100, 1) for _, v in items]
    return labels, values


def run_text_only(text):
    pred = predict_text(text)
    if pred is None:
        return "ENTER TEXT TO ANALYZE", gr.BarPlot()
    labels, values = scores_to_bar_data(pred["scores"])
    mood = pred["mood"]
    emoji = MOOD_EMOJI.get(mood, "")
    result = f"{emoji} **{mood.upper()}** — {pred['confidence']*100:.1f}% confidence"
    df = pd.DataFrame({"mood": labels, "score": values})
    return result, gr.BarPlot(value=df, x="mood", y="score", title="MOOD PROBABILITY (%)", y_lim=[0, 100])


def run_audio_only(audio_path):
    if _audio_pipe is None:
        return "AUDIO MODEL UNAVAILABLE — TRY AGAIN LATER", gr.BarPlot()
    pred = predict_audio(audio_path)
    if pred is None:
        return "RECORD OR UPLOAD AUDIO TO ANALYZE", gr.BarPlot()
    labels, values = scores_to_bar_data(pred["scores"])
    mood = pred["mood"]
    emoji = MOOD_EMOJI.get(mood, "")
    result = f"{emoji} **{mood.upper()}** — {pred['confidence']*100:.1f}% confidence"
    df = pd.DataFrame({"mood": labels, "score": values})
    return result, gr.BarPlot(value=df, x="mood", y="score", title="MOOD PROBABILITY (%)", y_lim=[0, 100])


def run_fusion(text, audio_path, text_weight, audio_weight):
    """
    The missing piece: combine TEXT + AUDIO mood detection into one fused
    signal, then turn that fused mood into an actual track recommendation
    (the "accumulation of these two -> music" step).
    """
    text_pred = predict_text(text)
    audio_pred = predict_audio(audio_path) if _audio_pipe is not None else None

    text_scores = text_pred["scores"] if text_pred else None
    audio_scores = audio_pred["scores"] if audio_pred else None

    fused = fuse_scores(text_scores, audio_scores, text_weight, audio_weight)
    if fused is None:
        return (
            "ENTER TEXT AND/OR AUDIO TO ANALYZE",
            gr.BarPlot(),
            "",
            "",
        )

    labels, values = scores_to_bar_data(fused["scores"])
    mood = fused["mood"]
    emoji = MOOD_EMOJI.get(mood, "")

    sources = []
    if text_scores:
        sources.append("TEXT")
    if audio_scores:
        sources.append("AUDIO")
    source_str = " + ".join(sources)

    result = (
        f"{emoji} **{mood.upper()}** — {fused['confidence']*100:.1f}% confidence "
        f"(fused from {source_str})"
    )
    df = pd.DataFrame({"mood": labels, "score": values})
    chart = gr.BarPlot(value=df, x="mood", y="score", title="FUSED MOOD PROBABILITY (%)", y_lim=[0, 100])

    try:
        tracks = search_tracks_for_mood(mood, n=5)
        tracks_html = tracks_to_html(tracks, mood)
        playlist_label = f"### {emoji} Recommended for **{mood.upper()}**"
    except RuntimeError as e:
        tracks_html = f"<p style='color:#ff4655'>{e}</p>"
        playlist_label = ""

    return result, chart, playlist_label, tracks_html


# ── Tactical theme CSS — matches radioweb.info palette ─────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;700&family=Syne:wght@700;800&display=swap');

:root, .gradio-container {
    --body-background-fill: #050510 !important;
    --background-fill-primary: #0a0a18 !important;
    --background-fill-secondary: #0a0a18 !important;
    --block-background-fill: rgba(255,255,255,0.02) !important;
    --block-border-color: rgba(255,255,255,0.1) !important;
    --block-label-background-fill: transparent !important;
    --block-label-text-color: rgba(255,255,255,0.55) !important;
    --block-label-border-color: rgba(255,255,255,0.1) !important;
    --block-title-text-color: #ffffff !important;
    --body-text-color: #ffffff !important;
    --body-text-color-subdued: rgba(255,255,255,0.45) !important;
    --border-color-primary: rgba(255,255,255,0.1) !important;
    --border-color-accent: #ff4655 !important;
    --border-color-accent-subdued: rgba(255,70,85,0.4) !important;
    --input-background-fill: #0d0d1f !important;
    --input-background-fill-hover: #11112a !important;
    --input-background-fill-focus: #11112a !important;
    --input-border-color: rgba(255,255,255,0.15) !important;
    --input-border-color-hover: #ff4655 !important;
    --input-border-color-focus: #ff4655 !important;
    --button-primary-background-fill: #ff4655 !important;
    --button-primary-background-fill-hover: #ff5e6b !important;
    --button-primary-border-color: #ff4655 !important;
    --button-primary-border-color-hover: #ff5e6b !important;
    --button-secondary-background-fill: rgba(255,255,255,0.06) !important;
    --button-secondary-background-fill-hover: rgba(255,255,255,0.12) !important;
    --button-secondary-border-color: rgba(255,255,255,0.15) !important;
    --panel-background-fill: #0a0a18 !important;
    --panel-border-color: rgba(255,255,255,0.1) !important;
    --table-even-background-fill: #0a0a18 !important;
    --table-odd-background-fill: #0d0d1f !important;
    --table-border-color: rgba(255,255,255,0.1) !important;
    --code-background-fill: #0a0a18 !important;
    --error-background-fill: rgba(255,70,85,0.08) !important;
    --error-border-color: #ff4655 !important;
}

.gradio-container {
    background: #050510 !important;
    font-family: 'JetBrains Mono', monospace !important;
}

h1, h2, h3 {
    font-family: 'Syne', sans-serif !important;
    font-weight: 800 !important;
    letter-spacing: -0.5px;
    color: #ffffff !important;
}

#title-block {
    border-bottom: 1px solid rgba(255,255,255,0.1);
    padding-bottom: 16px;
    margin-bottom: 12px;
}

.accent { color: #ff4655 !important; }

label span, .label-wrap span {
    font-size: 0.7rem !important;
    letter-spacing: 0.15em !important;
    text-transform: uppercase !important;
}

button.primary, button.primary span {
    font-family: 'JetBrains Mono', monospace !important;
    letter-spacing: 0.15em !important;
    text-transform: uppercase !important;
    font-size: 0.75rem !important;
    font-weight: 700 !important;
}

button.primary:hover {
    box-shadow: 0 0 20px rgba(255,70,85,0.4) !important;
}

.result-box {
    font-size: 1.1rem !important;
    padding: 16px !important;
    border-left: 3px solid #ff4655 !important;
}

a { color: #00ccff !important; }

.gr-samples-table td, .gr-sample-textbox, table.gr-dataset td {
    color: #ffffff !important;
    background: rgba(255,255,255,0.04) !important;
    border-color: rgba(255,255,255,0.15) !important;
}

#examples-wrap button {
    color: #ffffff !important;
    background: rgba(255,255,255,0.05) !important;
    border-color: rgba(255,255,255,0.15) !important;
}

#examples-wrap button span {
    color: #ffffff !important;
}

#examples-wrap button:hover {
    background: rgba(255,70,85,0.1) !important;
    border-color: #ff4655 !important;
}

.tab-nav button {
    font-family: 'JetBrains Mono', monospace !important;
    letter-spacing: 0.1em !important;
    text-transform: uppercase !important;
    font-size: 0.75rem !important;
}

footer { display: none !important; }
"""

with gr.Blocks(title="Moodwave — Emotion Detector", css=CSS, theme=gr.themes.Base(
    primary_hue="red", neutral_hue="slate",
)) as demo:
    with gr.Column(elem_id="title-block"):
        gr.Markdown(
            "# MOODWAVE <span class='accent'>//</span> EMOTION DETECTOR\n"
            "Multimodal Emotion Recognition project. Analyze emotion from "
            "**text** or **speech audio**, mapped to 5 moods — then fuse both "
            "signals together to get matching track recommendations."
        )

    with gr.Tabs():
        with gr.Tab("TEXT"):
            txt_input = gr.Textbox(
                label="INPUT TEXT",
                placeholder="Type how you're feeling, or anything at all…",
                lines=4,
            )
            txt_btn = gr.Button("ANALYZE", variant="primary")
            txt_result = gr.Markdown(elem_classes="result-box", value="ENTER TEXT TO ANALYZE")
            txt_chart = gr.BarPlot(
                x="mood", y="score",
                title="MOOD PROBABILITY (%)",
                y_lim=[0, 100],
            )

            txt_btn.click(run_text_only, inputs=txt_input, outputs=[txt_result, txt_chart])

            with gr.Group(elem_id="examples-wrap"):
                gr.Examples(
                    examples=[
                        ["I finally got the internship offer, I'm over the moon!"],
                        ["I don't really care either way, it is what it is."],
                        ["This keeps happening and I'm so done with it."],
                        ["I keep thinking something terrible is about to happen."],
                        ["I miss how things used to be."],
                    ],
                    inputs=txt_input,
                )

        with gr.Tab("AUDIO"):
            gr.Markdown(
                "Record or upload a short voice clip. Powered by a public "
                "Wav2Vec2 speech-emotion model (fallback, since the original "
                "fine-tuned weights trained on RAVDESS are not hosted here).\n\n"
                "**Note:** this model was trained on actors performing exaggerated "
                "emotions, so calm/normal speaking voices often read as NEUTRAL "
                "or SADNESS even when you're not sad — that's a dataset bias, not "
                "a bug. For best results, try adding clear vocal expression."
            )
            audio_input = gr.Audio(
                label="VOICE INPUT",
                sources=["microphone", "upload"],
                type="filepath",
            )
            audio_btn = gr.Button("ANALYZE", variant="primary")
            audio_result = gr.Markdown(elem_classes="result-box", value="RECORD OR UPLOAD AUDIO TO ANALYZE")
            audio_chart = gr.BarPlot(
                x="mood", y="score",
                title="MOOD PROBABILITY (%)",
                y_lim=[0, 100],
            )

            audio_btn.click(run_audio_only, inputs=audio_input, outputs=[audio_result, audio_chart])

        with gr.Tab("🎵 MOOD ➜ MUSIC"):
            gr.Markdown(
                "This is the **accumulation** step: combine text + voice into one "
                "fused mood, then get matching track recommendations (powered by "
                "Apple's free iTunes Search API — 30-second previews, no login "
                "needed). Fill in either or both — when both are given, they're "
                "blended using the project's research fusion weights (text 20% / "
                "audio 80%, F1 = 0.850); adjust the sliders to experiment."
            )

            with gr.Row():
                fusion_text = gr.Textbox(
                    label="TEXT (OPTIONAL)",
                    placeholder="Type how you're feeling…",
                    lines=3,
                )
                fusion_audio = gr.Audio(
                    label="VOICE (OPTIONAL)",
                    sources=["microphone", "upload"],
                    type="filepath",
                )
            with gr.Row():
                text_weight_slider = gr.Slider(0, 1, value=DEFAULT_TEXT_WEIGHT, step=0.05, label="TEXT WEIGHT")
                audio_weight_slider = gr.Slider(0, 1, value=DEFAULT_AUDIO_WEIGHT, step=0.05, label="AUDIO WEIGHT")

            fusion_btn = gr.Button("ANALYZE + RECOMMEND", variant="primary")
            fusion_result = gr.Markdown(elem_classes="result-box", value="ENTER TEXT AND/OR AUDIO TO ANALYZE")
            fusion_chart = gr.BarPlot(
                x="mood", y="score",
                title="FUSED MOOD PROBABILITY (%)",
                y_lim=[0, 100],
            )
            playlist_label = gr.Markdown("")
            playlist_html = gr.HTML("")

            fusion_btn.click(
                run_fusion,
                inputs=[fusion_text, fusion_audio, text_weight_slider, audio_weight_slider],
                outputs=[fusion_result, fusion_chart, playlist_label, playlist_html],
            )

    gr.Markdown(
        "---\n"
        "**MULTIMODAL EMOTION RECOGNITION** · text + speech fusion research project · "
        "[View full project on GitHub](https://github.com/sradowana-ux/moodwave)"
    )

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