import cv2
import gradio as gr
from ultralytics import YOLO

print("Đang tải mô hình Pose...")
pose_model = YOLO('yolov8n-pose.pt')

def analyze_posture(frame, do_calibrate, baseline_data):
    if frame is None:
        return None, "Đang chờ Camera...", do_calibrate, baseline_data
    
    frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    results = pose_model(frame_bgr, verbose=False, conf=0.5)
    
    annotated_frame = frame_bgr.copy()
    status_text = "Đang theo dõi..."
    
    for r in results:
        annotated_frame = r.plot()
        
        if r.keypoints is not None and len(r.keypoints.data) > 0:
            kpts = r.keypoints.data[0]
            
            if len(kpts) >= 7:
                nose = kpts[0][:2].tolist()
                left_shoulder = kpts[5][:2].tolist()
                right_shoulder = kpts[6][:2].tolist()
                
                if nose[0] != 0 and left_shoulder[0] != 0 and right_shoulder[0] != 0:
                    shoulder_mid_y = (left_shoulder[1] + right_shoulder[1]) / 2
                    current_neck_dist = shoulder_mid_y - nose[1]
                    current_sh_diff = abs(left_shoulder[1] - right_shoulder[1])
                    
                    if current_neck_dist <= 0:
                        continue
                        
                    # BƯỚC 1: Xử lý khi người dùng bấm nút "Hiệu chỉnh"
                    if do_calibrate:
                        baseline_data = {
                            "neck_dist": current_neck_dist,
                            "sh_diff": current_sh_diff
                        }
                        do_calibrate = False # Tắt cờ hiệu chỉnh
                        status_text = "Đã lưu tư thế chuẩn! Bắt đầu giám sát ✅"
                        break # Lưu xong thì thoát vòng lặp khung hình này
                        
                    # BƯỚC 2: Cảnh báo theo Tỷ lệ (Nếu đã có baseline)
                    if baseline_data is not None:
                        base_neck = baseline_data["neck_dist"]
                        base_sh_diff = baseline_data["sh_diff"]
                        
                        status_text = "Tư thế TỐT ✅"
                        
                        # Cảnh báo 1: Gù lưng / Rướn cổ (Rút ngắn > 25% so với chuẩn)
                        if current_neck_dist < 0.75 * base_neck:
                            status_text = "⚠️ CẢNH BÁO: Đang rướn cổ / Gù lưng!"
                            
                        # Cảnh báo 2: Lệch vai
                        # Mẹo: Dùng độ dài cổ (base_neck) làm thước đo. Nếu vai lệch > 20% chiều dài cổ là có vấn đề
                        tilt_threshold = 0.20 * base_neck
                        if abs(current_sh_diff - base_sh_diff) > tilt_threshold:
                            status_text = "⚠️ CẢNH BÁO: Ngồi lệch vai!"
                    else:
                        status_text = "⏳ Vui lòng ngồi thẳng và bấm 'Hiệu chỉnh tư thế chuẩn'."

    annotated_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
    return annotated_rgb, status_text, do_calibrate, baseline_data

# ==========================================
# GIAO DIỆN CÓ THÊM NÚT CALIBRATION
# ==========================================
with gr.Blocks(title="SmartErgo - Dynamic Pose Tracking") as demo:
    gr.Markdown("## 🎥 Giám Sát Tư Thế (Có Hiệu Chỉnh Động)")
    
    with gr.Row():
        with gr.Column(scale=2):
            live_input = gr.Image(sources=["webcam"], streaming=True, label="Camera của bạn")
        
        with gr.Column(scale=1):
            live_output = gr.Image(label="AI Tracking")
            status_output = gr.Textbox(label="Trạng thái", text_align="center", value="Chưa hiệu chỉnh")
            
            gr.Markdown("**Hướng dẫn:** Ngồi thẳng lưng, mắt nhìn ngang màn hình rồi bấm nút bên dưới.")
            calibrate_btn = gr.Button("🎯 Hiệu chỉnh tư thế chuẩn", variant="primary")
            
    # Các biến State để lưu trữ dữ liệu ngầm
    do_calibrate = gr.State(False)
    baseline_data = gr.State(None)
    
    # Khi bấm nút, bật cờ do_calibrate lên True
    calibrate_btn.click(fn=lambda: True, outputs=[do_calibrate])
    
    # Xử lý stream
    live_input.stream(
        fn=analyze_posture,
        inputs=[live_input, do_calibrate, baseline_data],
        outputs=[live_output, status_output, do_calibrate, baseline_data],
        stream_every=0.1
    )

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