import cv2
import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig

# Load the ronka/postureDetection-Mediapipe model
import mediapipe as mp
mp_pose = mp.solutions.pose
pose = mp_pose.Pose()
lmPose = mp_pose.PoseLandmark

class PostureDetectionConfig(PretrainedConfig):
    model_type = "posture_detection"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

class PostureDetectionModel(PreTrainedModel):
    config_class = PostureDetectionConfig

    def __init__(self, config):
        super().__init__(config)
        self.pose = pose

    def forward(self, pixel_values):
        images = [cv2.cvtColor(pixel_value.numpy(), cv2.COLOR_RGB2BGR) for pixel_value in pixel_values]
        keypoints = [self.pose.process(image) for image in images]

        results = []
        for image, keypoint in zip(images, keypoints):
            lm = keypoint.pose_landmarks
            if lm is None:
                results.append(None)
                continue

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

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

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

            # Left hip.
            l_hip_x = int(lm.landmark[lmPose.LEFT_HIP].x * image.shape[1])
            l_hip_y = int(lm.landmark[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)

            output = {
                "offset": offset,
                "neck_inclination": neck_inclination,
                "torso_inclination": torso_inclination
            }

            results.append(output)

        return results

    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

# Register the model and configuration
MODEL_MAPPING = {
    PostureDetectionConfig.model_type: PostureDetectionModel
}

CONFIG_MAPPING = {
    PostureDetectionConfig.model_type: PostureDetectionConfig
}