"""
HuggingFace Space - ESS Variable Classification Demo
Interactive Gradio interface for the XLM-RoBERTa ESS classifier.
"""
import gradio as gr
from transformers import pipeline

# Load the model
MODEL_NAME = "benjaminBeuster/xlm-roberta-base-ess-classification"
classifier = pipeline("text-classification", model=MODEL_NAME)

# Category descriptions
CATEGORY_INFO = {
    "DEMOGRAPHY (POPULATION, VITAL STATISTICS, AND CENSUSES)": "Demographics, population statistics, age, gender",
    "ECONOMICS": "Economic issues, finance, income",
    "EDUCATION": "Education, schooling, qualifications",
    "HEALTH": "Healthcare, medical services, health satisfaction",
    "POLITICS": "Political systems, trust in government, parliament",
    "SOCIETY AND CULTURE": "Social issues, cultural topics, religion",
    "LABOUR AND EMPLOYMENT": "Work, occupation, employment status",
    "PSYCHOLOGY": "Mental health, psychological wellbeing",
    "OTHER": "General or uncategorized topics"
}

def classify_text(text):
    """Classify survey question/variable."""
    if not text.strip():
        return "Please enter some text to classify."

    result = classifier(text)[0]
    label = result['label']
    score = result['score']

    # Format output
    output = f"**Category:** {label}\n\n"
    output += f"**Confidence:** {score:.2%}\n\n"

    if label in CATEGORY_INFO:
        output += f"**Description:** {CATEGORY_INFO[label]}"

    return output

# Example questions
examples = [
    ["What is your age?"],
    ["How satisfied are you with the healthcare system?"],
    ["Trust in country's parliament"],
    ["What is your highest level of education?"],
    ["How often do you pray?"],
    ["What is your occupation?"],
    ["Do you feel safe walking alone at night?"],
]

# Create Gradio interface
demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Enter a survey question or variable description...",
        label="Input Text"
    ),
    outputs=gr.Markdown(label="Classification Result"),
    title="ESS Variable Classification",
    description="""
    Classify European Social Survey (ESS) variables into 19 subject categories.

    This model is fine-tuned from XLM-RoBERTa-Base and achieves 83.8% accuracy on the test set.
    """,
    examples=examples,
    article="""
    ### About

    This classifier helps organize survey variables by automatically categorizing them into subject areas.
    Built on [XLM-RoBERTa-Base](https://huggingface.co/FacebookAI/xlm-roberta-base),
    trained on European Social Survey metadata.

    **Model:** [benjaminBeuster/xlm-roberta-base-ess-classification](https://huggingface.co/benjaminBeuster/xlm-roberta-base-ess-classification)

    **Performance:** 83.8% accuracy | F1: 0.796 (weighted)
    """,
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

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