import cv2
import math as m
import mediapipe as mp

def findDistance(x1, y1, x2, y2):
    dist = m.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
    return dist

def findAngle(x1, y1, x2, y2):
    theta = m.acos((y2 - y1) * (-y1) / (m.sqrt(
        (x2 - x1) ** 2 + (y2 - y1) ** 2) * y1))
    degree = int(180 / m.pi) * theta
    return degree

mp_pose = mp.solutions.pose
pose = mp_pose.Pose()

cap = cv2.VideoCapture(0)

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

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    keypoints = pose.process(frame)

    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    if keypoints.pose_landmarks:
        lm = keypoints.pose_landmarks
        lmPose = mp_pose.PoseLandmark

        # Acquire the landmark coordinates.
        l_shldr_x = int(lm.landmark[lmPose.LEFT_SHOULDER].x * frame.shape[1])
        l_shldr_y = int(lm.landmark[lmPose.LEFT_SHOULDER].y * frame.shape[0])
        r_shldr_x = int(lm.landmark[lmPose.RIGHT_SHOULDER].x * frame.shape[1])
        r_shldr_y = int(lm.landmark[lmPose.RIGHT_SHOULDER].y * frame.shape[0])
        l_ear_x = int(lm.landmark[lmPose.LEFT_EAR].x * frame.shape[1])
        l_ear_y = int(lm.landmark[lmPose.LEFT_EAR].y * frame.shape[0])
        l_hip_x = int(lm.landmark[lmPose.LEFT_HIP].x * frame.shape[1])
        l_hip_y = int(lm.landmark[lmPose.LEFT_HIP].y * frame.shape[0])

        # Calculate distance between left shoulder and right shoulder points.
        offset = findDistance(l_shldr_x, l_shldr_y, r_shldr_x, r_shldr_y)

        # Calculate angles.
        neck_inclination = findAngle(l_shldr_x, l_shldr_y, l_ear_x, l_ear_y)
        torso_inclination = findAngle(l_hip_x, l_hip_y, l_shldr_x, l_shldr_y)

        # Draw landmarks and lines.
        cv2.circle(frame, (l_shldr_x, l_shldr_y), 7, (255, 255, 0), -1)
        cv2.circle(frame, (l_ear_x, l_ear_y), 7, (0, 255, 255), -1)
        cv2.circle(frame, (r_shldr_x, r_shldr_y), 7, (255, 0, 255), -1)
        cv2.circle(frame, (l_hip_x, l_hip_y), 7, (0, 255, 0), -1)

        cv2.line(frame, (l_shldr_x, l_shldr_y), (r_shldr_x, r_shldr_y), (255, 255, 0), 2)
        cv2.line(frame, (l_shldr_x, l_shldr_y), (l_ear_x, l_ear_y), (0, 255, 255), 2)
        cv2.line(frame, (l_ear_x, l_ear_y), (r_shldr_x, r_shldr_y), (255, 0, 255), 2)
        cv2.line(frame, (l_shldr_x, l_shldr_y), (l_hip_x, l_hip_y), (0, 255, 0), 2)
        cv2.line(frame, (l_hip_x, l_hip_y), (r_shldr_x, r_shldr_y), (255, 255, 0), 2)

        angle_text_string = 'Neck : ' + str(int(neck_inclination)) + 'degrees. Torso : ' + str(int(torso_inclination)) + 'degrees.'
        cv2.putText(frame, angle_text_string, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        # Determine whether good posture or bad posture.
        if neck_inclination < 40 and torso_inclination < 10:
            cv2.putText(frame, 'Good Posture', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        else:
            cv2.putText(frame, 'Bad Posture', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)

    cv2.imshow('Posture Analysis', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()