import torch
import torch.nn as nn
import mediapipe as mp
import cv2

class PostureAnalysisModel(nn.Module):
    def __init__(self):
        super(PostureAnalysisModel, self).__init__()
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose()
        self.lmPose = self.mp_pose.PoseLandmark

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

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

    def forward(self, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        keypoints = self.pose.process(image)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        lm = keypoints.pose_landmarks

        if lm is None:
            return None

        # Acquire the landmark coordinates.
        # Left shoulder.
        l_shldr_x = int(lm.landmark[self.lmPose.LEFT_SHOULDER].x * image.shape[1])
        l_shldr_y = int(lm.landmark[self.lmPose.LEFT_SHOULDER].y * image.shape[0])

        # Right shoulder
        r_shldr_x = int(lm.landmark[self.lmPose.RIGHT_SHOULDER].x * image.shape[1])
        r_shldr_y = int(lm.landmark[self.lmPose.RIGHT_SHOULDER].y * image.shape[0])

        # Left ear.
        l_ear_x = int(lm.landmark[self.lmPose.LEFT_EAR].x * image.shape[1])
        l_ear_y = int(lm.landmark[self.lmPose.LEFT_EAR].y * image.shape[0])

        # Left hip.
        l_hip_x = int(lm.landmark[self.lmPose.LEFT_HIP].x * image.shape[1])
        l_hip_y = int(lm.landmark[self.lmPose.LEFT_HIP].y * image.shape[0])

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

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

        return offset, neck_inclination, torso_inclination