import cv2
import mediapipe as mp
import math


class poseDetector():

    def __init__(self, mode=False, complexity=1, smooth_landmarks=True,
                 enable_segmentation=False, smooth_segmentation=True,
                 detectionCon=0.5, trackCon=0.5):

        self.mode = mode
        self.complexity = complexity
        self.smooth_landmarks = smooth_landmarks
        self.enable_segmentation = enable_segmentation
        self.smooth_segmentation = smooth_segmentation
        self.detectionCon = detectionCon
        self.trackCon = trackCon

        self.mpDraw = mp.solutions.drawing_utils
        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(self.mode, self.complexity, self.smooth_landmarks,
                                     self.enable_segmentation, self.smooth_segmentation,
                                     self.detectionCon, self.trackCon)

    def findPose(self, img, draw=True):
        imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        self.results = self.pose.process(imgRGB)

        if self.results.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(img, self.results.pose_landmarks,
                                           self.mpPose.POSE_CONNECTIONS)

        return img

    def drawPoint(self, img, p1, p2, p3):
        x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
        x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
        x3, y3 = self.lmList[p3][1], self.lmList[p3][2]
        cv2.circle(img, (x1, y1), 5, (102, 106, 233), cv2.FILLED)
        cv2.circle(img, (x2, y2), 5, (50, 80, 4), cv2.FILLED)
        cv2.circle(img, (x3, y3), 5, (20, 95, 104), cv2.FILLED)

    def Screen_occupation_left(self):
        if self.results.pose_landmarks :
            for id, lm in enumerate(self.results.pose_landmarks.landmark):
                if id == 23:
                    x1 = lm.x
                if id == 29:
                    x2 = lm.x
            ocp = x2 - x1
        return ocp

    def Screen_occupation_right(self):
        if self.results.pose_landmarks:
            for id, lm in enumerate(self.results.pose_landmarks.landmark):
                if id == 24:
                    x1 = lm.x
                if id == 30:
                    x2 = lm.x
            ocp = x2 - x1
        return ocp

    def findPosition(self, img, draw=True):
        self.lmList = []
        if self.results.pose_landmarks:
            for id, lm in enumerate(self.results.pose_landmarks.landmark):
                # finding height, width of the image printed
                h, w, c = img.shape
                # Determining the pixels of the landmarks
                cx, cy, cz = int(lm.x * w), int(lm.y * h), lm.z
                vis = lm.visibility
                self.lmList.append([id, cx, cy, cz, vis])
                if draw:
                    cv2.circle(img, (cx, cy), 5, (255, 0, 0), cv2.FILLED)
        return self.lmList

    def findConcave(self, p1, p2, p3):
        x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
        x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
        x3, y3 = self.lmList[p3][1], self.lmList[p3][2]

        # Calculate Angle
        Concave_angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
        if Concave_angle < 0:
            Concave_angle += 360

        return Concave_angle

    def findAngle(self, img, p1, p2, p3, draw=True):
        # Get the landmarks
        x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
        x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
        x3, y3 = self.lmList[p3][1], self.lmList[p3][2]

        # Calculate Angle
        angle = math.degrees(math.atan2(y3 - y2, x3 - x2) -
                             math.atan2(y1 - y2, x1 - x2))
        if angle < 0:
            angle += 360
            if angle > 180:
                angle = 360 - angle
        elif angle > 180:
            angle = 360 - angle
        # print(angle)

        # Draw
        if draw:
            cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 2)
            cv2.line(img, (x3, y3), (x2, y2), (255, 255, 255), 2)

            cv2.circle(img, (x1, y1), 3, (0, 0, 255), cv2.FILLED)
            cv2.circle(img, (x1, y1), 10, (0, 0, 255), 2)
            cv2.circle(img, (x2, y2), 3, (0, 0, 255), cv2.FILLED)
            cv2.circle(img, (x2, y2), 10, (0, 0, 255), 2)
            cv2.circle(img, (x3, y3), 3, (0, 0, 255), cv2.FILLED)
            cv2.circle(img, (x3, y3), 10, (0, 0, 255), 2)

            cv2.putText(img, str(int(angle)), (x2 - 20, y2 + 20),
                        cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
        return angle

    def findIncludedAngle(self, x1, y1, x2, y2):
        IncludedAngle = math.asin((y1 - y2) / math.hypot((y1 - y2), (x1 - x2))) * 180 / math.pi
        return IncludedAngle

    def findRange(self, x1, y1, x2, y2):
        Range = math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2))
        return Range


def main():
    detector = poseDetector()
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, img = cap.read()  # ret is just the return variable, not much in there that we will use.
        if ret:
            img = detector.findPose(img)
            cv2.imshow('Pose Detection', img)
        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()