import warnings
import os

warnings.filterwarnings("ignore", category=UserWarning, module="google.protobuf")

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import argparse
import math
import cv2
from mediapipe.python.solutions import pose
import logging


class Landmarker:

    resized_height = 256
    resized_width = 300

    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)

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

        self.person_height = args.person_height
        self.pixel_height = args.pixel_height

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

        self.landmarks_indices = [
            pose.PoseLandmark.LEFT_SHOULDER.value,
            pose.PoseLandmark.RIGHT_SHOULDER.value,
            pose.PoseLandmark.LEFT_ELBOW.value,
            pose.PoseLandmark.RIGHT_ELBOW.value,
            pose.PoseLandmark.LEFT_WRIST.value,
            pose.PoseLandmark.RIGHT_WRIST.value,
            pose.PoseLandmark.LEFT_HIP.value,
            pose.PoseLandmark.RIGHT_HIP.value,
            pose.PoseLandmark.LEFT_KNEE.value,
            pose.PoseLandmark.RIGHT_KNEE.value,
            pose.PoseLandmark.LEFT_ANKLE.value,
            pose.PoseLandmark.RIGHT_ANKLE.value,
        ]

    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",
            dest="pose_detection_confidence",
            default=0.5,
            type=float,
            help="Confidence score for pose detection",
        )
        parser.add_argument(
            "--pose_tracking_confidence",
            dest="pose_tracking_confidence",
            default=0.5,
            type=float,
            help="Confidence score for pose tracking",
        )
        parser.add_argument(
            "--person_height",
            default=153,
            dest="person_height",
            type=int,
            help="person height of person",
        )
        parser.add_argument(
            "--pixel_height",
            #    default=216,
            dest="pixel_height",
            type=int,
            help="pixel height of person",
        )
        return parser.parse_args()

    def run(self):

        logging.warning("person's height: %s", self.person_height)

        logging.warning("person's pixel height: %s", self.pixel_height)

        front_results = self.process_images()

        self.get_center_top_point(front_results)

        self.calculate_distance_betn_landmarks(front_results)

        self.print_distance()

        self.display_images()

        self.pose.close()

    def process_images(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:  # type: ignore
            self.draw_landmarks(
                self.front_image_keypoints,
                front_results.pose_landmarks,  # type: ignore
                self.landmarks_indices,
            )
        if side_results.pose_landmarks:  # type: ignore
            self.draw_landmarks(
                self.side_image_keypoints,
                side_results.pose_landmarks,  # type: ignore# type: ignore
                self.landmarks_indices,
            )
        return front_results

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

    def pixel_to_metric_ratio(self):
        self.pixel_height = self.pixel_distance * 2
        pixel_to_metric_ratio = self.person_height / self.pixel_height
        logging.warning("pixel_to_metric_ratio %s", 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)
            self.circle(image, cx, cy)

    def circle(self, image, cx, cy):
        return cv2.circle(image, (cx, cy), 2, (255, 0, 0), -1)

    def print_distance(self):
        print(
            "Distance between left shoulder and left elbow:",
            self.distance_left_hand_up,
        )

        print(
            "Distance between left elbow and left wrist:",
            self.distance_left_hand_down,
        )

        print("Distance between left hip and left knee:",
              self.distance_left_leg_up)

        print("Distance between left knee and left ankle:",
              self.distance_left_leg_down)

        print("Distance between center and top point:", self.distance)

        print("Height of person:", self.distance * 2)

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

        landmarks = front_results.pose_landmarks.landmark

        shoulder_left = landmarks[pose.PoseLandmark.LEFT_SHOULDER.value]
        elbow_left = landmarks[pose.PoseLandmark.LEFT_ELBOW.value]
        wrist_left = landmarks[pose.PoseLandmark.LEFT_WRIST.value]
        hip_left = landmarks[pose.PoseLandmark.LEFT_HIP.value]
        knee_left = landmarks[pose.PoseLandmark.LEFT_KNEE.value]
        ankle_left = landmarks[pose.PoseLandmark.LEFT_ANKLE.value]

        self.distance_left_hand_up = self.calculate_distance(shoulder_left, elbow_left)

        self.distance_left_hand_down = self.calculate_distance(elbow_left, wrist_left)

        self.distance_left_leg_up = self.calculate_distance(hip_left, knee_left)

        self.distance_left_leg_down = self.calculate_distance(knee_left, ankle_left)

    def euclidean_distance(self, x1, x2, y1, y2):
        distance = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
        return distance

    def destroy(self):
        cv2.destroyAllWindows()

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

    def get_center_top_point(self, side_results):
        gray_image = cv2.cvtColor(self.side_image_keypoints, cv2.COLOR_BGR2GRAY)
        blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
        roi = blurred_image[0 : int(self.side_image_resized.shape[0] / 2), :]
        self.edges = cv2.Canny(roi, 50, 150)
        contours, _ = cv2.findContours(
            self.edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        xt, yt = None, None
        self.topmost_point = None

        if contours:
            largest_contour = max(contours, key=cv2.contourArea)
            self.topmost_point = tuple(
                largest_contour[largest_contour[:, :, 1].argmin()][0]
            )
            xt, yt = self.topmost_point
            self.circle(self.side_image_keypoints, xt, yt)

            logging.warning("xt: %s", xt)
            logging.warning("yt: %s", yt)
        xc, yc = None, None
        landmarks = side_results.pose_landmarks.landmark

        if side_results.pose_landmarks:  # type: ignore
            left_hip = landmarks[pose.PoseLandmark.LEFT_HIP.value]
            right_hip = landmarks[pose.PoseLandmark.RIGHT_HIP.value]
            center_point = (
                (left_hip.x + right_hip.x) / 2,
                (left_hip.y + right_hip.y) / 2,
            )
            center_point = (
                int(center_point[0] * self.side_image_resized.shape[1]),
                int(center_point[1] * self.side_image_resized.shape[0]),
            )
            xc, yc = center_point
            logging.warning("xc: %s", xc)
            logging.warning("yc: %s", yc)
            self.circle(self.side_image_keypoints, xc, yc)

            self.pixel_distance = self.euclidean_distance(xc, xt, yc, yt)
            self.distance = (
                self.euclidean_distance(xc, xt, yc, yt) * self.pixel_to_metric_ratio()
            )
        return self.distance


l = Landmarker()
try:
    l.run()
except:
    print("error")
finally:
    l.destroy()
