import joblib
import cv2
import numpy as np
import mediapipe as mp
import gradio as gr
import matplotlib.pyplot as plt
import io
from PIL import Image

# 모델 로드
model = joblib.load("image_model.pkl")
action_labels = {0: "Weightlifting", 1: "Soccer", 2: "Handball"}

# MediaPipe 설정
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5)
mp_drawing = mp.solutions.drawing_utils

def process_frame(frame):
    # 1. None 체크
    if frame is None:
        return None, "Error: No image received from camera", None

    # 2. PIL.Image → numpy 변환
    if isinstance(frame, Image.Image):
        frame = np.array(frame)

    # 3. RGB 변환 예외 처리
    try:
        image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    except Exception as e:
        return frame, f"Error converting image: {e}", None

    results = pose.process(image_rgb)

    # 관절 시각화
    if results.pose_landmarks:
        mp_drawing.draw_landmarks(frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

    try:
        # 모델 입력 전처리
        resized = cv2.resize(image_rgb, (128, 128))
        input_data = resized.flatten().reshape(1, -1)
        prediction_proba = model.predict_proba(input_data)[0]
        predicted_class = np.argmax(prediction_proba)
        prediction = action_labels.get(predicted_class, "Unknown")
        confidence = prediction_proba[predicted_class] * 100
        label = f'{prediction} ({confidence:.2f}%)'
    except Exception as e:
        return frame, f"Prediction error: {e}", None

    # 예측 그래프 생성
    try:
        fig, ax = plt.subplots()
        ax.bar(action_labels.values(), prediction_proba, color='skyblue')
        ax.set_ylim([0, 1])
        ax.set_ylabel("Confidence")
        ax.set_title("Prediction Probabilities")
        buf = io.BytesIO()
        plt.savefig(buf, format="png")
        buf.seek(0)
        graph_img = Image.open(buf).convert("RGB")
        plt.close()
    except Exception as e:
        graph_img = None
        label += f" (Graph error: {e})"

    # 예측 텍스트 이미지에 표시
    try:
        cv2.putText(frame, label, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
    except:
        pass

    return frame, label, graph_img

# Gradio 인터페이스 구성
demo = gr.Interface(
    fn=process_frame,
    inputs=gr.Image(source="webcam", streaming=False),  # ✅ 사진 촬영 버튼 활성화!
    outputs=[
        gr.Image(label="Pose + Prediction"),
        gr.Textbox(label="Predicted Action"),
        gr.Image(label="Confidence Graph")
    ],
    title="Pose Action Classifier",
    description=(
        "📱 On mobile, tap the button to take a photo. "
        "You can switch to the back camera by tapping the camera icon. "
        "Internet Explorer is not supported — use Chrome or Safari."
    )
)

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