import gradio as gr
import cv2
import numpy as np
import torch
from torchvision import models, transforms
from PIL import Image
import mediapipe as mp
from fer import FER  # Facial emotion recognition

# -----------------------------
# Initialize Models and Helpers
# -----------------------------

# MediaPipe Pose for posture analysis
mp_pose = mp.solutions.pose
pose = mp_pose.Pose()
mp_drawing = mp.solutions.drawing_utils

# MediaPipe Face Detection for face detection
mp_face_detection = mp.solutions.face_detection
face_detection = mp_face_detection.FaceDetection(min_detection_confidence=0.5)

# Object Detection Model: Faster R-CNN (pretrained on COCO)
object_detection_model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
object_detection_model.eval()
obj_transform = transforms.Compose([transforms.ToTensor()])

# Facial Emotion Detection using FER (this model will detect emotions from a face)
emotion_detector = FER(mtcnn=True)

# -----------------------------
# Define Analysis Functions
# -----------------------------

def analyze_posture(frame_rgb, output_frame):
    """Runs pose estimation and draws landmarks on the frame."""
    pose_results = pose.process(frame_rgb)
    posture_text = "No posture detected"
    if pose_results.pose_landmarks:
        posture_text = "Posture detected"
        # Draw the pose landmarks on the output image (convert back to BGR for OpenCV)
        mp_drawing.draw_landmarks(
            output_frame, pose_results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
            mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
            mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2)
        )
    return posture_text

def analyze_emotion(frame):
    """Detects emotion from faces using FER. Returns the dominant emotion."""
    # FER expects RGB images
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    emotions = emotion_detector.detect_emotions(frame_rgb)
    if emotions:
        # Use the first detected face and its top emotion
        top_emotion, score = max(emotions[0]["emotions"].items(), key=lambda x: x[1])
        emotion_text = f"{top_emotion} ({score:.2f})"
    else:
        emotion_text = "No face detected for emotion analysis"
    return emotion_text

def analyze_objects(frame_rgb, output_frame):
    """Performs object detection and draws bounding boxes for detections above a threshold."""
    image_pil = Image.fromarray(frame_rgb)
    img_tensor = obj_transform(image_pil)
    with torch.no_grad():
        detections = object_detection_model([img_tensor])[0]

    threshold = 0.8
    detected_boxes = detections["boxes"][detections["scores"] > threshold]
    for box in detected_boxes:
        box = box.int().cpu().numpy()
        cv2.rectangle(output_frame, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 2)
    object_text = f"Detected {len(detected_boxes)} object(s)" if len(detected_boxes) else "No objects detected"
    return object_text

def analyze_faces(frame_rgb, output_frame):
    """Detects faces using MediaPipe and draws bounding boxes."""
    face_results = face_detection.process(frame_rgb)
    face_text = "No faces detected"
    if face_results.detections:
        face_text = f"Detected {len(face_results.detections)} face(s)"
        h, w, _ = output_frame.shape
        for detection in face_results.detections:
            bbox = detection.location_data.relative_bounding_box
            x = int(bbox.xmin * w)
            y = int(bbox.ymin * h)
            box_w = int(bbox.width * w)
            box_h = int(bbox.height * h)
            cv2.rectangle(output_frame, (x, y), (x + box_w, y + box_h), (0, 0, 255), 2)
    return face_text

# -----------------------------
# Main Analysis Function
# -----------------------------

def analyze_webcam(frame):
    """
    Runs posture analysis, facial emotion analysis, object detection, and face detection
    on the given webcam frame. Returns an annotated image and a textual summary.
    """
    if frame is None:
        return None, "No frame provided."

    # The input frame is in BGR (as from OpenCV). Create a copy for drawing.
    output_frame = frame.copy()

    # Convert frame to RGB for analysis
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Run analyses
    posture_result = analyze_posture(frame_rgb, output_frame)
    emotion_result = analyze_emotion(frame)
    object_result = analyze_objects(frame_rgb, output_frame)
    face_result = analyze_faces(frame_rgb, output_frame)

    # Compose the result summary text
    summary = (
        f"Posture Analysis: {posture_result}\n"
        f"Emotion Analysis: {emotion_result}\n"
        f"Object Detection: {object_result}\n"
        f"Face Detection: {face_result}"
    )

    # Optionally, overlay some of the summary text on the image
    cv2.putText(output_frame, f"Emotion: {emotion_result}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    cv2.putText(output_frame, f"Objects: {object_result}", (10, 70),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
    cv2.putText(output_frame, f"Faces: {face_result}", (10, 110),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    return output_frame, summary

# -----------------------------
# Gradio Interface Setup
# -----------------------------

# We output both an image (with drawn annotations) and a text summary.
interface = gr.Interface(
    fn=analyze_webcam,
    inputs=gr.Image(source="webcam", streaming=True, label="Webcam Feed"),
    outputs=[
        gr.Image(type="numpy", label="Annotated Output"),
        gr.Textbox(label="Analysis Summary")
    ],
    title="Real-Time Multi-Analysis App",
    description=(
        "This app performs real-time posture analysis, facial emotion detection, "
        "object detection, and face detection using your webcam."
    ),
    live=True
)

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