import gradio as gr
import spaces
import cv2
import pandas as pd
import plotly.express as px
from ultralytics import YOLO

# 1. INISIALISASI MODEL PYTORCH (.pt) DI LUAR FUNGSI
# Tujuannya agar model di-load sekali ke memori saat server berjalan, bukan setiap kali tombol ditekan.
face_model = YOLO('yolov11n-face.pt') 
emotion_model = YOLO('best.pt')

# 2. DEKORATOR ZEROGPU
# Meminjam GPU NVIDIA saat fungsi ini dieksekusi dengan batas waktu maksimal 120 detik.
@spaces.GPU(duration=120)
def analyze_video(video_path):
    if not video_path:
        return "❌ Silakan unggah video terlebih dahulu.", None, None
        
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps == 0 or fps != fps:
        fps = 30.0
        
    emotion_counts = {
        "Happy": 0, "Neutral": 0, "Angry": 0, "Contempt": 0,
        "Sad": 0, "Surprised": 0, "Fear": 0, "Disgust": 0
    }
    total_detections = 0
    unique_face_ids = set()
    saved_frames_pool = []
    
    FRAME_SKIP = 5
    TARGET_WIDTH = 640
    frame_idx = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        frame_idx += 1
        
        # Frame Skipping untuk efisiensi
        if frame_idx % FRAME_SKIP != 0:
            continue
            
        # Resize ukuran video
        h, w = frame.shape[:2]
        aspect_ratio = h / w
        target_height = int(TARGET_WIDTH * aspect_ratio)
        frame = cv2.resize(frame, (TARGET_WIDTH, target_height))
        
        annotated_frame = frame.copy()
        face_detected_in_this_frame = False
        
        # TAHAP 1: Deteksi Wajah (Otomatis berjalan di GPU)
        face_results = face_model.track(frame, conf=0.4, persist=True, verbose=False)
        
        for r_face in face_results:
            if r_face.boxes is not None and r_face.boxes.id is not None:
                boxes = r_face.boxes.xyxy
                track_ids = r_face.boxes.id.int().tolist()
                
                for box, track_id in zip(boxes, track_ids):
                    unique_face_ids.add(track_id)
                    
                    x1, y1, x2, y2 = map(int, box)
                    x1, y1 = max(0, x1), max(0, y1)
                    x2, y2 = min(TARGET_WIDTH, x2), min(target_height, y2)
                    
                    face_crop = frame[y1:y2, x1:x2]
                    label_name = "Unknown"
                    
                    if face_crop.size > 0:
                        # TAHAP 2: Klasifikasi Emosi
                        emotion_results = emotion_model(face_crop, conf=0.4, verbose=False)
                        for r_emotion in emotion_results:
                            if len(r_emotion.boxes) > 0:
                                top_box = r_emotion.boxes[0]
                                cls_id = int(top_box.cls)
                                label_name = emotion_model.names[cls_id].capitalize()
                                
                                if label_name in emotion_counts:
                                    emotion_counts[label_name] += 1
                                    total_detections += 1
                                    face_detected_in_this_frame = True
                                    
                    # Gambar bounding box dan teks
                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(annotated_frame, f"ID:{track_id} {label_name}", (x1, y1 - 7),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1, cv2.LINE_AA)
                                
        # Simpan ke pool gambar jika ada wajah dengan timestamp waktu
        if face_detected_in_this_frame and len(saved_frames_pool) < 60:
            time_in_seconds = frame_idx / fps
            minutes = int(time_in_seconds // 60)
            seconds = int(time_in_seconds % 60)
            timestamp_text = f"{minutes:02d}:{seconds:02d}"
            
            cv2.rectangle(annotated_frame, (10, 10), (100, 40), (0, 0, 0), -1)
            cv2.putText(annotated_frame, timestamp_text, (15, 33),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
                        
            rgb_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
            saved_frames_pool.append(rgb_frame)
            
    cap.release()
    
    # --- MENYIAPKAN OUTPUT UNTUK UI GRADIO ---
    
    if total_detections == 0:
        return "❌ Tidak ada wajah siswa yang terdeteksi di dalam video. Pastikan video cukup jelas.", None, None
        
    # 1. Output Gallery (4 Gambar Sampling)
    pool_size = len(saved_frames_pool)
    if pool_size >= 4:
        indices = [0, pool_size // 3, (pool_size * 2) // 3, pool_size - 1]
        selected_frames = [saved_frames_pool[i] for i in indices]
    else:
        selected_frames = saved_frames_pool
        
    # 2. Output Grafik Plotly
    labels = []
    percentages = []
    counts = []
    for emotion, count in emotion_counts.items():
        pct = (count / total_detections) * 100
        labels.append(emotion)
        percentages.append(round(pct, 1))
        counts.append(count)
        
    df = pd.DataFrame({
        "Kategori Emosi": labels,
        "Persentase (%)": percentages,
        "Jumlah": counts
    })
    
    fig = px.bar(
        df, x="Kategori Emosi", y="Persentase (%)", text="Persentase (%)",
        color="Kategori Emosi", title="<b>Mood Breakdown (Profil Emosi Kelas)</b>",
        color_discrete_sequence=px.colors.qualitative.Pastel
    )
    fig.update_traces(texttemplate='%{text}%', textposition='outside')
    fig.update_layout(showlegend=False, yaxis_range=[0, 110])
    
    # 3. Output Teks Kesimpulan (Markdown)
    max_emotion = df.loc[df['Persentase (%)'].idxmax()]['Kategori Emosi']
    max_pct = df['Persentase (%)'].max()
    
    summary_md = f"### 📊 Ringkasan Hasil Analisis\n"
    summary_md += f"- **Total ID Sesi Wajah:** {len(unique_face_ids)} ID\n"
    summary_md += f"- **Total Deteksi Emosi:** {total_detections} Kali\n\n"
    summary_md += f"### 💡 Rekomendasi Pengajaran\n"
    
    if max_emotion in ["Happy", "Neutral"]:
        summary_md += f"👉 **Insight Utama:** Kelas didominasi oleh emosi **{max_emotion} ({max_pct}%)**. Menandakan suasana belajar kondusif.\n\n"
        summary_md += f"**Saran Perbaikan:** Pertahankan ritme mengajar Anda. Sesi interaktif sudah berjalan efektif."
    else:
        summary_md += f"👉 **Insight Utama:** Terdeteksi tingkat emosi **{max_emotion} sebesar {max_pct}%** di dalam kelas.\n\n"
        summary_md += f"**Saran Perbaikan:** Angka emosi negatif ({max_emotion}) yang cukup tinggi menandakan siswa mengalami kendala. Disarankan untuk mengevaluasi kembali bagian materi yang rumit, memberikan jeda *ice breaking*, atau memperlambat tempo penjelasan pada pertemuan berikutnya."
        
    # Kembalikan 3 variabel sesuai urutan output pada blok gr.Button.click
    return summary_md, fig, selected_frames


# 3. MEMBANGUN UI (USER INTERFACE) DENGAN GRADIO
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🏫 EduReflect: Analisis Emosi Kelas")
    gr.Markdown("Upload rekaman video kelas untuk memproses ekspresi siswa dengan kecerdasan buatan.")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_video = gr.Video(label="Upload Rekaman Pembelajaran (MP4/MOV)")
            analyze_btn = gr.Button("Mulai Analisis 🚀", variant="primary")
        
        with gr.Column(scale=1):
            output_markdown = gr.Markdown(label="Kesimpulan Analisis")
            
    with gr.Row():
        output_gallery = gr.Gallery(label="Cuplikan Rekaman Analisis (Key Frames)", columns=4, height="auto")
        
    with gr.Row():
        output_plot = gr.Plot(label="Grafik Analisis")
        
    # Menghubungkan Tombol dengan Fungsi analyze_video
    analyze_btn.click(
        fn=analyze_video,
        inputs=[input_video],
        outputs=[output_markdown, output_plot, output_gallery]
    )

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