import gradio as gr
import pandas as pd
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
import os # Untuk memeriksa environment variabel

# --- 1. Konfigurasi Model ---
# Pastikan ID repositori model Anda sudah benar dan model (config.json, pytorch_model.bin, dll.)
# sudah terunggah sepenuhnya ke repositori ini.
MODEL_ID = "Dawgggggg/modelNER"

# --- 2. Inisialisasi Model dan Pipeline ---
# Gunakan try-except block untuk penanganan error saat memuat model
tokenizer = None
model = None
ner_pipeline = None

try:
    print(f"Loading tokenizer from {MODEL_ID}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    print("Tokenizer loaded.")

    print(f"Loading model from {MODEL_ID}...")
    model = AutoModelForTokenClassification.from_pretrained(MODEL_ID)
    print("Model loaded.")

    print("Creating NER pipeline...")
    ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
    print("NER pipeline created successfully.")

except Exception as e:
    print(f"Error loading model or tokenizer: {e}")
    # Jika ada error saat startup, kita bisa memberikan fungsi dummy
    # agar UI tetap muncul dengan pesan error
    def dummy_pipeline(text):
        print(f"Model not loaded. Error: {e}")
        return []
    ner_pipeline = dummy_pipeline

# --- 3. Fungsi Prediksi NER ---
def predict_ner_all(text):
    # Pastikan pipeline sudah terinisialisasi
    if ner_pipeline is None or isinstance(ner_pipeline, type(lambda:0)): # Cek jika dummy function
        return (
            pd.DataFrame(columns=["Entitas", "Kata", "Skor", "Mulai", "Akhir"]),
            "Error: Model tidak dapat dimuat. Cek log.",
            []
        )

    if not text:
        return (
            pd.DataFrame(columns=["Entitas", "Kata", "Skor", "Mulai", "Akhir"]),
            "", # Teks kosong
            [] # Highlight kosong
        )
    
    try:
        results = ner_pipeline(text)
    except Exception as e:
        print(f"Error during NER pipeline execution: {e}")
        return (
            pd.DataFrame(columns=["Entitas", "Kata", "Skor", "Mulai", "Akhir"]),
            f"Error saat menjalankan NER: {e}",
            []
        )

    # Untuk DataFrame
    df_results = []
    for res in results:
        df_results.append(
            {
                "Entitas": res["entity_group"],
                "Kata": res["word"],
                "Skor": f"{res['score']:.2f}",
                "Mulai": res["start"],
                "Akhir": res["end"]
            }
        )
    df_output = pd.DataFrame(df_results)

    # Untuk HighlightedText
    highlighted_chunks = []
    current_index = 0
    # Pastikan hasil dari pipeline tidak kosong sebelum iterasi
    if results:
        for res in results:
            # Tambahkan teks non-entitas sebelum entitas saat ini
            if res["start"] > current_index:
                highlighted_chunks.append((text[current_index:res["start"]], None))
            
            # Tambahkan entitas yang terdeteksi
            highlighted_chunks.append((res["word"], res["entity_group"], res["score"]))
            
            # Update current_index
            current_index = res["end"]
        
        # Tambahkan sisa teks setelah entitas terakhir
        if current_index < len(text):
            highlighted_chunks.append((text[current_index:], None))
    else:
        # Jika tidak ada entitas yang terdeteksi, tampilkan seluruh teks tanpa highlight
        highlighted_chunks.append((text, None))

    return df_output, text, highlighted_chunks

# --- 4. Antarmuka Gradio ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Demo NER Informasi Obat dan Dosis")
    gr.Markdown("Model ini dirancang untuk mengekstraksi informasi obat, dosis, dan unit dari teks medis berbahasa Indonesia.")

    with gr.Row():
        text_input = gr.Textbox(
            lines=2, 
            placeholder="Masukkan kalimat di sini, contoh: Diberikan parasetamol 5.6 mg/dL setiap 8 jam.", 
            label="Kalimat Input"
        )
        submit_button = gr.Button("Deteksi Entitas")

    # Output dalam bentuk tabel
    with gr.Row():
        df_output = gr.Dataframe(
            headers=["Entitas", "Kata", "Skor", "Mulai", "Akhir"],
            wrap=True,
            label="Hasil Deteksi (Tabel)"
        )
    
    # Output dalam bentuk teks yang disorot
    with gr.Row():
        # Karena gr.HighlightedText memerlukan input teks asli dan list highlight,
        # kita perlu "proxy" teks asli dari fungsi predict_ner_all
        # Output kedua dari predict_ner_all adalah teks asli
        # Output ketiga adalah list highlight
        highlighted_output = gr.HighlightedText(
            label="Deteksi Entitas (Teks Disorot)",
            color_map={
                "DRUG": "red",
                "RESULT": "blue",
                "UNIT": "green",
                # Tambahkan warna lain jika ada entitas lain yang Anda prediksi
                # Contoh: "DOSAGE": "purple"
            },
            show_legend=True # Menampilkan legenda warna
        )
    
    # Menghubungkan tombol dengan fungsi prediksi
    submit_button.click(
        fn=predict_ner_all, 
        inputs=text_input, 
        # Output harus sesuai urutan: df_output, highlighted_output (bagian text), highlighted_output (bagian highlight)
        outputs=[df_output, text_input, highlighted_output] # Output kedua dari predict_ner_all adalah text, yang diumpan kembali ke text_input sebagai "state"
    )

    # Contoh-contoh input untuk mempermudah pengujian
    gr.Examples(
        examples=[
            ["Diberikan parasetamol 5.6 mg/dL setiap 8 jam."],
            ["Pasien diberi insulin 10 unit sebelum makan."],
            ["Resepnya adalah amoxicillin 500 mg tiga kali sehari."]
        ],
        inputs=text_input
    )

# --- 5. Jalankan Aplikasi Gradio ---
# Set share=True untuk mendapatkan tautan publik saat berjalan secara lokal
# Saat di Hugging Face Space, ini otomatis menjadi True
demo.launch()