import gradio as gr
from transformers import pipeline
from collections import defaultdict

# Label mapping
label_mapping = {
    "LABEL_0": "Normal",
    "LABEL_1": "Depression",
    "LABEL_2": "Anxiety"
}

# Load classifier
classifier = pipeline("text-classification", model="coldnasser/mindscape-v2")


def predict(texts):
    try:
        if isinstance(texts, str):
            texts = [texts]

        results = classifier(texts)

        # Initialize score aggregator
        score_sums = defaultdict(float)
        count = len(texts)

        for res in results:
            label = res['label']
            score = res['score']
            score_sums[label] += score

        # Calculate average scores
        avg_scores = {label_mapping.get(label, label): score_sums[label] / count for label in score_sums}

        # Get final predicted label (highest average)
        final_label = max(avg_scores.items(), key=lambda x: x[1])[0]

        return {
            "Predicted Status": final_label,
            "Average Scores": avg_scores
        }

    except Exception as e:
        return {"Error": str(e)}

# Gradio interface
gr.Interface(
     fn=predict,
     inputs=gr.Textbox(
         lines=10,
         placeholder="Enter one or more texts (one per line)",
         label="Input Texts"
     ),
     outputs=gr.JSON(
         label="Predicted Status & Scores"
     ),
     title="Mindscape AI Therapist (Multi-text Support)"
).launch()