import gradio as gr
from transformers import pipeline
import pandas as pd
import time

# Pesan saat startup untuk memastikan pipeline mulai dimuat.
print("Memuat model pipeline...")

# -----------------------------------------------------------------------------
# 1. Muat semua model pipeline saat aplikasi dimulai.
# -----------------------------------------------------------------------------
try:
    pipe_distilbert = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
    pipe_bert = pipeline("text-classification", model="gchhablani/bert-base-cased-finetuned-sst2")
    pipe_roberta = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment")
    print("Semua model berhasil dimuat.")
except Exception as e:
    print(f"Error saat memuat model: {e}")

# Mapping label untuk model RoBERTa agar lebih mudah dibaca.
roberta_label_map = {
    "LABEL_0": "NEGATIVE",
    "LABEL_1": "NEUTRAL",
    "LABEL_2": "POSITIVE"
}

# -----------------------------------------------------------------------------
# 2. Fungsi utama untuk prediksi dan pengukuran performa.
# -----------------------------------------------------------------------------
def get_performance_data(text):
    """
    Menerima input teks, melakukan prediksi dengan tiga model, mengukur waktu,
    dan mengembalikan hasilnya dalam format DataFrame untuk dianalisis.
    """
    if not text.strip():
        return pd.DataFrame(columns=["Model", "Label", "Confidence Score", "Waktu Pemrosesan (ms)"])

    results = []

    # --- Prediksi dengan DistilBERT ---
    start_time = time.time()
    pred_distilbert = pipe_distilbert(text)[0]
    processing_time = (time.time() - start_time) * 1000
    results.append(["DistilBERT", pred_distilbert['label'], f"{pred_distilbert['score']:.4f}", f"{processing_time:.0f}"])

    # --- Prediksi dengan BERT ---
    start_time = time.time()
    pred_bert = pipe_bert(text)[0]
    processing_time = (time.time() - start_time) * 1000
    results.append(["BERT", pred_bert['label'], f"{pred_bert['score']:.4f}", f"{processing_time:.0f}"])

    # --- Prediksi dengan RoBERTa ---
    start_time = time.time()
    pred_roberta = pipe_roberta(text)[0]
    processing_time = (time.time() - start_time) * 1000
    label_roberta = roberta_label_map.get(pred_roberta['label'], pred_roberta['label'])
    results.append(["RoBERTa", label_roberta, f"{pred_roberta['score']:.4f}", f"{processing_time:.0f}"])

    df = pd.DataFrame(results, columns=["Model", "Label", "Confidence Score", "Waktu Pemrosesan (ms)"])
    return df

# -----------------------------------------------------------------------------
# 3. Buat antarmuka menggunakan gr.Blocks untuk tata letak kustom.
# -----------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # --- Blok Identitas di Bagian Atas ---
    gr.Markdown(
        """
        <div style="text-align: left;">
            <p><strong>Nama:</strong> Ravinasa Deo</p>
            <p><strong>NIM:</strong> 2304130048</p>
            <p><strong>Prodi:</strong> Teknik Informatika</p>
        </div>
        """
    )
    
    # --- Judul dan Deskripsi Aplikasi ---
    gr.Markdown(
        """
        <h1 style="text-align: center; font-size: 2em;">🧪 Eksperimen Performa Model Sentimen</h1>
        <p style="text-align: center;">Masukkan teks untuk mendapatkan data hasil prediksi dan waktu pemrosesan (latensi) dari tiga model berbeda. Data ini dapat digunakan untuk analisis performa.</p>
        """
    )

    # --- Komponen Input dan Output ---
    with gr.Row():
        input_textbox = gr.Textbox(
            lines=8,
            placeholder="Masukkan teks di sini untuk eksperimen...",
            label="Input Teks",
            scale=1  
        )
        output_dataframe = gr.Dataframe(
            headers=["Model", "Label", "Confidence Score", "Waktu Pemrosesan (ms)"],
            datatype=["str", "str", "str", "number"],
            label="Data Hasil Eksperimen",
            wrap=True,  
            scale=2
        )
    
    submit_button = gr.Button("Analisis Sekarang", variant="primary")
    
    # --- Contoh Input ---
    gr.Examples(
        examples=[
            ["The new design is absolutely gorgeous and the user experience is top-notch. I'm very impressed!"],
            ["I've been waiting for this feature for a long time, but the implementation is buggy and unreliable."],
            ["It's a decent product. Nothing special, but it gets the job done without any major issues."]
        ],
        inputs=input_textbox
    )
    # --- Hubungkan Aksi Tombol ke Fungsi ---
    submit_button.click(
        fn=get_performance_data,
        inputs=input_textbox,
        outputs=output_dataframe
    )

    gr.Markdown(
        """
        ---
        ### Kredit dan Sumber Daya
        Aplikasi ini dibangun menggunakan sumber daya berikut:
        * **Model:**
            * [gchhablani/bert-base-cased-finetuned-sst2](https://huggingface.co/gchhablani/bert-base-cased-finetuned-sst2)
            * [cardiffnlp/twitter-roberta-base-sentiment](https://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment)
            * [distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english)
        * **Platform & Tools:**
            * Hosting: [Hugging Face Spaces](https://huggingface.co/spaces)
            * UI Framework: [Gradio](https://www.gradio.app/)
        * **Bantuan Pengembangan:**
            * [ChatGPT](https://chat.openai.com/)
            * [Google Gemini](https://gemini.google.com/)
        """
    )

# Jalankan aplikasi
if __name__ == "__main__":
    demo.launch()