import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_PATH = "bert_final_model_v1"
LABELS = {
    0: "Normal",
    1: "Distressed",
    2: "Suicidal"
}

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)


def predict_mental_health(text):
    if not text or not text.strip():
        return "Please enter some text to analyze.", {}, ""

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)[0]
        predicted_class = torch.argmax(probabilities).item()

    prediction = LABELS[predicted_class]
    confidence = probabilities[predicted_class].item()

    prob_dict = {
        LABELS[i]: float(probabilities[i].item())
        for i in range(len(LABELS))
    }

    result_text = (
        f"**Prediction:** {prediction}\n\n"
        f"**Confidence:** {confidence * 100:.1f}%"
    )

    return result_text, prob_dict, text


custom_css = """
.gradio-container {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
}
.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
}
"""


with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Mental Health Text Analyzer (BERT Final v1)
        AI-powered mental health status detection using a fine-tuned BERT model.

        This model classifies text into three categories: **Normal**, **Distressed**, or **Suicidal**.
        """
    )

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Enter text to analyze",
                placeholder="Type or paste text here...",
                lines=5
            )

            submit_btn = gr.Button("Analyze Text", variant="primary", elem_classes="primary-btn")

            gr.Markdown(
                """
                ### Examples
                Try these sample texts to see how the model works.
                """
            )

            gr.Examples(
                examples=[
                    ["I had a wonderful day at the park with my family!"],
                    ["I'm feeling really anxious about my upcoming exam."],
                    ["I feel so hopeless, like nothing will ever get better."],
                    ["Just finished a great workout session, feeling energized!"],
                    ["I can't stop these dark thoughts, everything feels pointless."]
                ],
                inputs=text_input
            )

        with gr.Column():
            result_output = gr.Markdown(label="Result")
            probabilities_output = gr.Label(label="Detailed Probabilities", num_top_classes=3)

    submit_btn.click(
        fn=predict_mental_health,
        inputs=text_input,
        outputs=[result_output, probabilities_output, text_input]
    )

    text_input.submit(
        fn=predict_mental_health,
        inputs=text_input,
        outputs=[result_output, probabilities_output, text_input]
    )

    gr.Markdown(
        """
        ---
        ### Important Disclaimer

        This tool is for research and educational purposes only. It should not be used as a substitute
        for professional mental health care. If you or someone you know is experiencing a mental health crisis,
        please contact a mental health professional or crisis helpline immediately.

        ### Crisis Resources

        - National Suicide Prevention Lifeline: 1-800-273-8255
        - Crisis Text Line: Text HOME to 741741
        - International Association for Suicide Prevention: https://www.iasp.info/resources/Crisis_Centres/
        """
    )


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