import os

import cv2
import mediapipe as mp
from DataStreamModule import DataStreamModule
from IHPEModule import IHPEModule

# initialize mediapipe drawing utilities and pose models
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose


# This class implements Human Pose Estimation (HPE) using the mediapipe library
class HPEModule(IHPEModule):

    # This method starts HPE using a camera specified by its name
    def startHPEwithCamera(self, camera_name):
        out = None
        # check if the camera_name is a file path or not
        if os.path.isfile(camera_name):
            # open the video file using cv2.VideoCapture
            cap = cv2.VideoCapture(camera_name)
            # set the output video file path
            output_path = os.path.splitext(camera_name)[0] + "_output.mp4"
            # get the frame rate and size of the input video
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            # initialize video writer to save the output video
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        else:
            # initialize data stream and camera object
            data_stream = DataStreamModule()
            cap = data_stream.get_camera_stream(camera_name)

        # set the window name using the camera name
        window_name = f"Pose Estimation on Camera {camera_name}"

        # start pose detection
        with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
            while True:
                # get the next frame from the camera
                if out is not None:
                    ret, frame = cap.read()
                    if not ret:
                        break
                else:
                    frame = next(cap)

                # Recolor image to RGB
                image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image.flags.writeable = False

                # Make detection
                results = pose.process(image)

                # Recolor back to BGR
                image.flags.writeable = True
                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

                # Extract landmarks
                try:
                    landmarks = results.pose_landmarks.landmark
                    #print(landmarks)
                except:
                    pass

                # Render detections
                mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                                          mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=2),
                                          mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2)
                                          )

                cv2.imshow(window_name, image)

                if out is not None:
                    out.write(image)

                keyCode = cv2.waitKey(1)
                if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1:
                    break