# -*- coding: utf-8 -*-
"""GradioWebsite.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/14D-sxTs35_Vc9__q6maqZ6cb8OqfhTqt
"""

import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
from safetensors.torch import load_file
import matplotlib.pyplot as plt
import numpy as np
import os


from transformers import AutoModelForSequenceClassification, AutoTokenizer
from safetensors.torch import load_file
import torch

# Define your model path
model_path = "/content/drive/MyDrive/model_3500_data_pred_label_fold_5_dari_5(indoBERT_lite_large)/model.safetensors"

# Load the model architecture (BERT-like, IndoBERT, etc.)
model_name = "indobert-lite-large-p2"  # Replace with the correct model name for your case
model = AutoModelForSequenceClassification.from_pretrained('indobenchmark/indobert-lite-large-p2', num_labels=3)

# Load the model weights using safetensors
weights = load_file(model_path)

# Load the weights into the model
model.load_state_dict(weights)

# Set the model to evaluation mode
model.eval()

# Verify the model loaded correctly
print(model)

from transformers import AutoTokenizer
from transformers import BertForSequenceClassification, BertTokenizer, AutoModelForSequenceClassification

# Load the tokenizer
tokenizer = BertTokenizer.from_pretrained('indobenchmark/indobert-lite-large-p2')

import torch
import torch.nn.functional as F
import gradio as gr

def predict_stress_with_accuracy(text_input):
    if not text_input.strip():
        return None, None, None, None

    # Tokenize the input text
    inputs = tokenizer(text_input, return_tensors="pt", padding=True, truncation=True, max_length=512)

    # Get the model's output
    with torch.no_grad():
        output = model(**inputs)

    # Apply softmax to get probabilities
    probabilities = F.softmax(output.logits, dim=1)

    # Get predictions for all classes
    probs = probabilities[0].tolist()
    confidence_scores = [round(p * 100, 1) for p in probs]

    # Get main prediction
    predicted_class = torch.argmax(probabilities, dim=1).item()
    main_confidence = confidence_scores[predicted_class]

    # Map the predicted class to stress level
    stress_levels = {0: "Neutral", 1: "Mild Stress", 2: "Very Stress"}
    prediction = stress_levels[predicted_class]

    # Generate HTML for the main result
    result_html = f"""
        <div class="result-card">
            <div class="prediction-text">{prediction}</div>
            <div class="confidence-bar-container">
                <div class="confidence-bar" style="width: {main_confidence}%"></div>
                <span class="confidence-text">{main_confidence}% Confident</span>
            </div>
        </div>
    """

    # Generate HTML for detailed analysis
    detailed_html = f"""
        <div class="detailed-analysis">
            <div class="analysis-title">Detailed Analysis</div>
            <div class="analysis-bars">
                <div class="analysis-bar">
                    <div class="bar-label">Neutral</div>
                    <div class="bar-container">
                        <div class="bar neutral" style="width: {confidence_scores[0]}%"></div>
                        <span class="bar-value">{confidence_scores[0]}%</span>
                    </div>
                </div>
                <div class="analysis-bar">
                    <div class="bar-label">Mild Stress</div>
                    <div class="bar-container">
                        <div class="bar mild" style="width: {confidence_scores[1]}%"></div>
                        <span class="bar-value">{confidence_scores[1]}%</span>
                    </div>
                </div>
                <div class="analysis-bar">
                    <div class="bar-label">Very Stress</div>
                    <div class="bar-container">
                        <div class="bar very" style="width: {confidence_scores[2]}%"></div>
                        <span class="bar-value">{confidence_scores[2]}%</span>
                    </div>
                </div>
            </div>
        </div>
    """

    return result_html, detailed_html, prediction, main_confidence

# Create the interface
with gr.Blocks(css="""
    #component-0 {
        max-width: 900px;
        margin: auto;
        padding: 0 20px;
    }

    .container {
        background: linear-gradient(135deg, #1a1c29, #2d3748);
        border-radius: 20px;
        padding: 2rem;
        box-shadow: 0 10px 30px rgba(0,0,0,0.2);
    }

    .header {
        text-align: center;
        margin-bottom: 2rem;
    }

    .title {
        font-size: 2.5rem;
        font-weight: bold;
        background: linear-gradient(45deg, #00c6ff, #0072ff);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        margin-bottom: 0.5rem;
    }

    .subtitle {
        color: #a0aec0;
        font-size: 1.1rem;
    }

    /* Input styles */
    .input-container {
        background: rgba(255,255,255,0.05);
        border-radius: 15px;
        padding: 1.5rem;
        margin-bottom: 2rem;
    }

    textarea {
        background: rgba(255,255,255,0.07) !important;
        border: 2px solid rgba(255,255,255,0.1) !important;
        border-radius: 12px !important;
        color: white !important;
        font-size: 1.1rem !important;
        transition: all 0.3s ease !important;
    }

    textarea:focus {
        border-color: #00c6ff !important;
        box-shadow: 0 0 20px rgba(0,198,255,0.2) !important;
    }

    /* Result card styles */
    .result-card {
        background: rgba(255,255,255,0.07);
        border-radius: 15px;
        padding: 1.5rem;
        margin-bottom: 1.5rem;
        animation: fadeIn 0.5s ease-out;
    }

    .prediction-text {
        font-size: 1.8rem;
        font-weight: bold;
        color: white;
        text-align: center;
        margin-bottom: 1rem;
    }

    .confidence-bar-container {
        background: rgba(255,255,255,0.1);
        border-radius: 10px;
        height: 20px;
        position: relative;
        overflow: hidden;
    }

    .confidence-bar {
        background: linear-gradient(90deg, #00c6ff, #0072ff);
        height: 100%;
        border-radius: 10px;
        transition: width 0.5s ease-out;
    }

    .confidence-text {
        position: absolute;
        top: 50%;
        left: 50%;
        transform: translate(-50%, -50%);
        color: white;
        font-weight: bold;
        text-shadow: 0 0 10px rgba(0,0,0,0.5);
    }

    /* Detailed analysis styles */
    .detailed-analysis {
        background: rgba(255,255,255,0.07);
        border-radius: 15px;
        padding: 1.5rem;
        animation: fadeIn 0.5s ease-out;
    }

    .analysis-title {
        color: white;
        font-size: 1.3rem;
        font-weight: bold;
        margin-bottom: 1rem;
        text-align: center;
    }

    .analysis-bar {
        margin-bottom: 1rem;
    }

    .bar-label {
        color: #a0aec0;
        margin-bottom: 0.5rem;
    }

    .bar-container {
        background: rgba(255,255,255,0.1);
        border-radius: 8px;
        height: 15px;
        position: relative;
        overflow: hidden;
    }

    .bar {
        height: 100%;
        transition: width 0.5s ease-out;
    }

    .bar.neutral { background: linear-gradient(90deg, #00f2c3, #0098f0); }
    .bar.mild { background: linear-gradient(90deg, #ffd600, #ff9100); }
    .bar.very { background: linear-gradient(90deg, #ff5724, #ff2d55); }

    .bar-value {
        position: absolute;
        right: 10px;
        top: 50%;
        transform: translateY(-50%);
        color: white;
        font-size: 0.9rem;
        font-weight: bold;
    }

    @keyframes fadeIn {
        from { opacity: 0; transform: translateY(10px); }
        to { opacity: 1; transform: translateY(0); }
    }

    @media (max-width: 768px) {
        .container {
            padding: 1rem;
        }

        .title {
            font-size: 2rem;
        }

        .prediction-text {
            font-size: 1.5rem;
        }
    }
""") as iface:
    gr.HTML("""
        <div class="header">
            <div class="title">Klasifikasi Tingkat Stress</div>
            <div class="subtitle">Jelaskan keadaan emosi Anda dan biarkan AI menganalisis tingkat stres Anda</div>
        </div>
    """)

    with gr.Column(elem_classes="container"):
        text_input = gr.Textbox(
            label="Describe Your Emotional State",
            placeholder="Apa kabar hari ini?",
            lines=4,
            elem_classes="input-container"
        )

        result_html = gr.HTML()
        detailed_html = gr.HTML()
        prediction = gr.State()
        confidence = gr.State()

        text_input.change(
            predict_stress_with_accuracy,
            inputs=[text_input],
            outputs=[result_html, detailed_html, prediction, confidence]
        )

iface.launch(share=True)