import gradio as gr
import spaces
from transformers import pipeline

# Available models shown in the dropdown.
MODEL_OPTIONS = {
    "XLM-RoBERTa X-Stance": "MatteoFasulo/xlm-roberta-xstance",
    "mDeBERTa-v3 X-Stance": "MatteoFasulo/mdeberta-v3-xstance",
}

# Cache one classifier for each model so it is not reloaded on every request.
classifiers = {}


def get_classifier(model_name):
    if model_name not in classifiers:
        classifiers[model_name] = pipeline(
            task="text-classification",
            model=model_name,
            device=0,
        )

    return classifiers[model_name]


@spaces.GPU
def predict_stance(selected_model, question, comment):
    if not question.strip() or not comment.strip():
        return "⚠️ Please provide both a question and a comment.", None

    try:
        model_name = MODEL_OPTIONS[selected_model]
        model = get_classifier(model_name)

        result = model(
            {
                "text": question,
                "text_pair": comment,
            }
        )

        if isinstance(result, list) and len(result) > 0:
            prediction = result[0]
            label = prediction["label"]
            score = prediction["score"]
        elif isinstance(result, dict):
            label = result.get("label", "Unknown")
            score = result.get("score", 0.0)
        else:
            return "⚠️ Unexpected model output format.", None

        normalized_label = label.upper()

        if "FAVOR" in normalized_label:
            emoji = "✅"
            color = "green"
            explanation = "The comment **supports** the political question."
        elif "AGAINST" in normalized_label:
            emoji = "❌"
            color = "red"
            explanation = "The comment **opposes** the political question."
        else:
            emoji = "❓"
            color = "orange"
            explanation = "The model returned an unrecognized stance label."

        output = f"""
### {emoji} Prediction: **{label}**

<div style="
    padding: 10px;
    border-left: 4px solid {color};
    background-color: #f5f5f5;
    margin: 10px 0;
">
{explanation}
</div>

**Confidence:** {score:.2%}

**Model:** `{model_name}`

---
*💡 Tip: Try questions in German or French!*
        """

        if "FAVOR" in normalized_label:
            confidence_dist = {
                "FAVOR": score,
                "AGAINST": 1 - score,
            }
        elif "AGAINST" in normalized_label:
            confidence_dist = {
                "AGAINST": score,
                "FAVOR": 1 - score,
            }
        else:
            confidence_dist = {label: score}

        return output, confidence_dist

    except Exception as e:
        error_msg = f"❌ Error during prediction: {str(e)}"
        return error_msg, None


with gr.Blocks(
    title="Multilingual Stance Detection",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown("""
    # 🌍 Multilingual Political Stance Detection

    Select a model and predict whether a comment **supports (FAVOR)**
    or **opposes (AGAINST)** a political question.

    The models support multilingual stance detection, including
    🇩🇪 German and 🇫🇷 French.
    """)

    with gr.Row():
        with gr.Column():
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_OPTIONS.keys()),
                value="XLM-RoBERTa X-Stance",
                label="🤖 Model",
                info="Choose the model used for stance classification.",
                interactive=True,
            )

            question_input = gr.Textbox(
                label="🗳️ Political Question",
                lines=3,
                placeholder=(
                    "e.g. Sollte die Schweiz die Kernenergie verbieten?"
                ),
            )

            comment_input = gr.Textbox(
                label="💬 Comment",
                lines=5,
                placeholder=(
                    "e.g. Erneuerbare Energien sollten Kernenergie "
                    "ersetzen, weil sie sicherer und nachhaltiger sind."
                ),
            )

            submit_btn = gr.Button(
                "🔍 Analyze Stance",
                variant="primary",
            )

        with gr.Column():
            output_text = gr.Markdown(
                label="Analysis Results",
            )

            confidence_plot = gr.Label(
                label="Confidence Distribution",
                num_top_classes=2,
            )

    gr.Examples(
        examples=[
            [
                "Sollte die Schweiz erneuerbare Energien stärker fördern?",
                (
                    "Investitionen in erneuerbare Energien reduzieren "
                    "Emissionen und verbessern die Energieunabhängigkeit."
                ),
            ],
            [
                "Sollte die Schweiz der Europäischen Union beitreten?",
                (
                    "Die Schweiz muss ihre Unabhängigkeit und Neutralität "
                    "um jeden Preis bewahren."
                ),
            ],
            [
                "Sollte die Schweiz die Kernenergie schrittweise abschaffen?",
                (
                    "Kernenergie ist gefährlich und sollte durch sicherere "
                    "Alternativen ersetzt werden."
                ),
            ],
            [
                "Sollte die Schweiz die Einwanderung begrenzen?",
                (
                    "Eine Begrenzung der Einwanderung schadet der Wirtschaft "
                    "und dem kulturellen Austausch."
                ),
            ],
            [
                (
                    "Sollte die Schweiz ein bedingungsloses "
                    "Grundeinkommen einführen?"
                ),
                (
                    "Ein bedingungsloses Grundeinkommen würde die soziale "
                    "Sicherheit stärken und Armut reduzieren."
                ),
            ],
            [
                (
                    "La Suisse devrait-elle promouvoir davantage "
                    "les énergies renouvelables?"
                ),
                (
                    "Les investissements dans les énergies renouvelables "
                    "réduisent les émissions et améliorent "
                    "l'indépendance énergétique."
                ),
            ],
            [
                "La Suisse devrait-elle adhérer à l'Union européenne?",
                (
                    "La Suisse doit préserver son indépendance et sa "
                    "neutralité à tout prix."
                ),
            ],
            [
                (
                    "La Suisse devrait-elle éliminer progressivement "
                    "l'énergie nucléaire?"
                ),
                (
                    "L'énergie nucléaire est dangereuse et devrait être "
                    "remplacée par des alternatives plus sûres."
                ),
            ],
            [
                "La Suisse devrait-elle limiter l'immigration?",
                (
                    "La limitation de l'immigration nuit à l'économie "
                    "et aux échanges culturels."
                ),
            ],
            [
                (
                    "La Suisse devrait-elle introduire un revenu "
                    "de base inconditionnel?"
                ),
                (
                    "Un revenu de base inconditionnel renforcerait la "
                    "sécurité sociale et réduirait la pauvreté."
                ),
            ],
        ],
        inputs=[question_input, comment_input],
        label="📝 Try these examples",
    )

    prediction_inputs = [
        model_dropdown,
        question_input,
        comment_input,
    ]

    prediction_outputs = [
        output_text,
        confidence_plot,
    ]

    submit_btn.click(
        fn=predict_stance,
        inputs=prediction_inputs,
        outputs=prediction_outputs,
    )

    comment_input.submit(
        fn=predict_stance,
        inputs=prediction_inputs,
        outputs=prediction_outputs,
    )


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