import cv2
import numpy as np
import matplotlib.pyplot as plt
import mediapipe as mp
from app.face_utils import get_box

mp_face_mesh = mp.solutions.face_mesh

def preprocess_video_and_predict_sleep_quality(video):
    cap = cv2.VideoCapture(video)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = np.round(cap.get(cv2.CAP_PROP_FPS))

    path_save_video_original = 'result_original.mp4'
    path_save_video_face = 'result_face.mp4'
    path_save_video_sleep = 'result_sleep.mp4'
    
    vid_writer_original = cv2.VideoWriter(path_save_video_original, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    vid_writer_face = cv2.VideoWriter(path_save_video_face, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
    vid_writer_sleep = cv2.VideoWriter(path_save_video_sleep, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))

    frames = []
    sleep_quality_scores = []
    eye_bags_images = []

    with mp_face_mesh.FaceMesh(
    max_num_faces=1,
    refine_landmarks=False,
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5) as face_mesh:

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = face_mesh.process(frame_rgb)

            if results.multi_face_landmarks:
                for fl in results.multi_face_landmarks:
                    startX, startY, endX, endY = get_box(fl, w, h)
                    cur_face = frame_rgb[startY:endY, startX:endX]
                    
                    sleep_quality_score, eye_bags_image = analyze_sleep_quality(cur_face)
                    sleep_quality_scores.append(sleep_quality_score)
                    eye_bags_images.append(cv2.resize(eye_bags_image, (224, 224)))

                    sleep_quality_viz = create_sleep_quality_visualization(cur_face, sleep_quality_score)
                    
                    cur_face = cv2.resize(cur_face, (224, 224))
                    
                    vid_writer_face.write(cv2.cvtColor(cur_face, cv2.COLOR_RGB2BGR))
                    vid_writer_sleep.write(sleep_quality_viz)

            vid_writer_original.write(frame)
            frames.append(len(frames) + 1)

    cap.release()
    vid_writer_original.release()
    vid_writer_face.release()
    vid_writer_sleep.release()

    sleep_stat = sleep_quality_statistics_plot(frames, sleep_quality_scores)
    
    if eye_bags_images:
        average_eye_bags_image = np.mean(np.array(eye_bags_images), axis=0).astype(np.uint8)
    else:
        average_eye_bags_image = np.zeros((224, 224, 3), dtype=np.uint8)

    return (path_save_video_original, path_save_video_face, path_save_video_sleep, 
            average_eye_bags_image, sleep_stat)

def analyze_sleep_quality(face_image):
    # Placeholder function - implement your sleep quality analysis here
    sleep_quality_score = np.random.random()
    eye_bags_image = cv2.resize(face_image, (224, 224))
    return sleep_quality_score, eye_bags_image

def create_sleep_quality_visualization(face_image, sleep_quality_score):
    viz = face_image.copy()
    cv2.putText(viz, f"Sleep Quality: {sleep_quality_score:.2f}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    return cv2.cvtColor(viz, cv2.COLOR_RGB2BGR)

def sleep_quality_statistics_plot(frames, sleep_quality_scores):
    fig, ax = plt.subplots()
    ax.plot(frames, sleep_quality_scores)
    ax.set_xlabel('Frame')
    ax.set_ylabel('Sleep Quality Score')
    ax.set_title('Sleep Quality Over Time')
    plt.tight_layout()
    return fig