import argparse
import math
import cv2
import mediapipe as mp


class Landmarker:

    def __init__(self) -> None:
        args = self.parse_args()
        if args.front_image == None:
            raise Exception("front image needs to be passed")
        if args.side_image == None:
            raise Exception("side image needs to be passed")

        self.front_image = cv2.imread(args.front_image)
        self.side_image = cv2.imread(args.side_image)

        mp_pose = mp.solutions.pose
        self.pose = mp_pose.Pose(
            static_image_mode=True,
            min_detection_confidence=args.pose_detection_confidence,
            min_tracking_confidence=args.pose_tracking_confidence,
        )

        self.resized_height = 256
        self.resized_width = int(
            self.front_image.shape[1] *
            (self.resized_height / self.front_image.shape[0]))

        self.front_image_resized = cv2.resize(
            self.front_image, (self.resized_width, self.resized_height))
        self.side_image_resized = cv2.resize(
            self.side_image, (self.resized_width, self.resized_height))

        self.landmarks_indices = [
            11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28
        ]

    def parse_args(self):
        parser = argparse.ArgumentParser()
        parser.add_argument("--front",
                            dest="front_image",
                            type=str,
                            help="Front image")
        parser.add_argument("--side",
                            dest="side_image",
                            type=str,
                            help="Side image")
        parser.add_argument("--pose_detection_confidence",
                            deafult=0.5,
                            dest="Pose detection confidence",
                            type=float,
                            help="Confidence score for pose detection")
        parser.add_argument("--pose_tracking_confidence",
                            default=0.5,
                            dest="Pose detection tracking",
                            type=float,
                            help="Confidence score for pose tracking")
        return parser.parse_args()

    def run(self):
        front_results = self.pose.process(
            cv2.cvtColor(self.front_image_resized, cv2.COLOR_BGR2RGB))
        side_results = self.pose.process(
            cv2.cvtColor(self.side_image_resized, cv2.COLOR_BGR2RGB))

        self.side_image_keypoints = self.side_image_resized.copy()
        self.front_image_keypoints = self.front_image_resized.copy()

        if front_results.pose_landmarks:
            self.draw_landmarks(
                self.front_image_keypoints,
                front_results.pose_landmarks,
                self.landmarks_indices,
            )
        if side_results.pose_landmarks:
            self.draw_landmarks(
                self.side_image_keypoints,
                side_results.pose_landmarks,
                self.landmarks_indices,
            )

        self.get_center_top_point(front_results)

        self.print_distance(front_results)

        self.display_images()

        self.pose.close()

    def calculate_distance(self, landmark1, landmark2):
        l1 = landmark1.x, landmark1.y
        l1 = int(l1[0] * self.resized_width), int(l1[1] * self.resized_height)
        x1, y1 = l1
        l2 = landmark2.x, landmark2.y
        l2 = int(l2[0] * self.resized_width), int(l2[1] * self.resized_height)
        x2, y2 = l2
        pixel_distance = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
        real_distance = pixel_distance * self.pixel_to_metric_ratio()
        return real_distance

    def pixel_to_metric_ratio(self):
        height_person = 153
        pixel_distance_height = 255
        pixel_to_metric_ratio = height_person / pixel_distance_height
        pixel_to_metric_ratio = round(pixel_to_metric_ratio)
        return pixel_to_metric_ratio

    def draw_landmarks(self, image, landmarks, indices):
        for idx in indices:
            landmark = landmarks.landmark[idx]
            h, w, _ = image.shape
            cx, cy = int(landmark.x * w), int(landmark.y * h)
            cv2.circle(image, (cx, cy), 2, (255, 0, 0), -1)

    def print_distance(self, front_results):
        if not front_results.pose_landmarks:
            return

        landmarks = front_results.pose_landmarks.landmark

        shoulder_left = landmarks[11]
        elbow_left = landmarks[13]
        wrist_left = landmarks[15]
        hip_left = landmarks[23]
        knee_left = landmarks[25]
        ankle_left = landmarks[27]

        distance_left_hand_up = self.calculate_distance(
            shoulder_left, elbow_left)
        print("Distance between left shoulder and left elbow:",
              distance_left_hand_up)

        distance_left_hand_down = self.calculate_distance(
            elbow_left, wrist_left)
        print("Distance between left elbow and left wrist:",
              distance_left_hand_down)

        distance_left_leg_up = self.calculate_distance(hip_left, knee_left)
        print("Distance between left hip and left knee:", distance_left_leg_up)

        distance_left_leg_down = self.calculate_distance(knee_left, ankle_left)
        print("Distance between left knee and left ankle:",
              distance_left_leg_down)

    def destroy():
        cv2.destroyAllWindows()

    def display_images(self):
        cv2.imshow("front_image_keypoints", self.front_image_keypoints)
        cv2.imshow("side_image_keypoints", self.side_image_keypoints)

    def get_center_top_point(self, front_results):
        gray_image = cv2.cvtColor(self.front_image_keypoints,
                                  cv2.COLOR_BGR2GRAY)
        blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
        roi = blurred_image[0:int(self.resized_height / 2), :]
        edges = cv2.Canny(roi, 50, 150)
        contours, hierarchy = cv2.findContours(edges, cv2.RETR_TREE,
                                               cv2.CHAIN_APPROX_SIMPLE)

        self.topmost_point = None

        outermost_contours = [
            contours[i] for i in range(len(contours))
            if hierarchy[0][i][3] == -1
        ]
        if outermost_contours:
            largest_contour = max(outermost_contours, key=cv2.contourArea)
            self.topmost_point = tuple(
                largest_contour[largest_contour[:, :, 1].argmin()][0])
            # topmost_point = (int(topmost_point[0]*self.resized_width), int(topmost_point[1] * self.resized_height))
            topmost_point = (self.topmost_point[0], self.topmost_point[1])
            cv2.circle(self.front_image_keypoints, topmost_point, 2,
                       (255, 0, 0), -1)

        self.left_hip_idx = 23
        self.right_hip_idx = 24

        self.center_point = None

        if front_results.pose_landmarks:
            left_hip = front_results.pose_landmarks.landmark[self.left_hip_idx]
            right_hip = front_results.pose_landmarks.landmark[
                self.right_hip_idx]
            center_point = (
                (left_hip.x + right_hip.x) / 2,
                (left_hip.y + right_hip.y) / 2,
            )
            center_point = (
                int(center_point[0] * self.resized_width),
                int(center_point[1] * self.resized_height),
            )
            cv2.circle(self.front_image_keypoints, center_point, 2,
                       (255, 0, 0), -1)
            xc, yc = center_point
            xt, yt = topmost_point
            distance = math.sqrt((xc - xt)**2 + (yc - yt)**2)
            #return distance
            return center_point


l = Landmarker()
try:
    l.run()
except:
    pass
finally:
    l.destroy()
