using Mediapipe;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public class KeypointManagerEmbedding
{

    private int leftShoulderIndex = 11;
    private int rightShoulderIndex = 12;
    private int neckIndex = 33;
    private int noseIndex = 0;
    private int leftEyeIndex = 2;

    private List<int> pose_indices = new List<int> { 0, 33, 5, 2, 8, 7, 12, 11, 14, 13, 16, 15 };
    private List<int> hand_indices = new List<int> { 0, 8, 7, 6, 5, 12, 11, 10, 9, 16, 15, 14, 13, 20, 19, 18, 17, 4, 3, 2, 1 }; 

    private static int BUFFER_SIZE = 10;
    private List<List<List<float>>> keypointsBuffer;

    public KeypointManagerEmbedding()
    {
        keypointsBuffer = new List<List<List<float>>>();
    }

    private (List<float>, List<float>) NormalizeHand(List<float> handX, List<float> handY)
    {
        var xValues = new List<float>();
        var yValues = new List<float>();
        for (int i = 0; i < handX.Count; i++)
        {
            if (handX[i] != 0)
            {
                xValues.Add(handX[i]);
            }
            if (handY[i] != 0)
            {
                yValues.Add(handY[i]);
            }
        }

        if (xValues.Count == 0 || yValues.Count == 0)
        {
            return (handX, handY);
        }

        float width = xValues.Max() - xValues.Min();
        float height = yValues.Max() - yValues.Min();
        float delta_x, delta_y;

        if (width > height)
        {
            delta_x = 0.1f * width;
            delta_y = delta_x + ((width - height) / 2f);
        }
        else
        {
            delta_y = 0.1f * height;
            delta_x = delta_y + ((height - width) / 2f);
        }

        var startingPoint = new Vector2(xValues.Min() - delta_x, yValues.Min() - delta_y);
        var endingPoint = new Vector2(xValues.Max() + delta_x, yValues.Max() + delta_y);

        if (endingPoint.x - startingPoint.x == 0f || endingPoint.y - startingPoint.y == 0f)
        {
            return (handX, handY);
        }

        // normalize keypoints
        for (int i = 0; i < handX.Count; i++)
        {
            handX[i] = (handX[i] - startingPoint.x) / (endingPoint.x - startingPoint.x);
            handY[i] = (handY[i] - startingPoint.y) / (endingPoint.y - startingPoint.y);
        }

        return (handX, handY);
    }

    private (List<float>, List<float>) NormalizePose(List<float> poseX, List<float> poseY)
    {
        var leftShoulder = new Vector2(poseX[leftShoulderIndex], poseY[leftShoulderIndex]);
        var rightShoulder = new Vector2(poseX[rightShoulderIndex], poseY[rightShoulderIndex]);
        var neck = new Vector2(poseX[neckIndex], poseY[neckIndex]);
        var nose = new Vector2(poseX[noseIndex], poseY[noseIndex]);

        // Prevent from even starting the analysis if some necessary elements are not present
        if ((leftShoulder.x == 0 || rightShoulder.x == 0 ||
            (leftShoulder.x == rightShoulder.x && leftShoulder.y == rightShoulder.y)) &&
            (neck.x == 0 || nose.x == 0 || (neck.x == nose.x && neck.y == nose.y)))
        {
            return (poseX, poseY);
        }

        float shoulderDistance, headMetric;
        if (leftShoulder.x != 0 && rightShoulder.x != 0 &&
            (leftShoulder.x != rightShoulder.x || leftShoulder.y != rightShoulder.y))
        {
            shoulderDistance = Mathf.Sqrt(Mathf.Pow(leftShoulder.x - rightShoulder.x, 2) + Mathf.Pow(leftShoulder.y - rightShoulder.y, 2));
            headMetric = shoulderDistance;
        }
        else
        {
            float neckNoseDistance = Mathf.Sqrt(Mathf.Pow(neck.x - nose.x, 2) + Mathf.Pow(neck.y - nose.y, 2));
            headMetric = neckNoseDistance;
        }

        // Set the starting and ending point of the normalization bounding box
        var startingPoint = new Vector2(poseX[neckIndex] - 3 * headMetric, poseY[leftEyeIndex] + headMetric);
        var endingPoint = new Vector2(poseX[neckIndex] + 3 * headMetric, startingPoint.y - 6 * headMetric);

        if (startingPoint.x < 0)
        {
            startingPoint.x = 0;
        }
        if (startingPoint.y < 0)
        {
            startingPoint.y = 0;
        }
        if (endingPoint.x < 0)
        {
            endingPoint.x = 0;
        }
        if (endingPoint.y < 0)
        {
            endingPoint.y = 0;
        }

        // Normalize the keypoints
        for (int i = 0; i < poseX.Count; i++)
        {
            poseX[i] = (poseX[i] - startingPoint.x) / (endingPoint.x - startingPoint.x);
            poseY[i] = (poseY[i] - endingPoint.y) / (startingPoint.y - endingPoint.y);
        }

        return (poseX, poseY);
    }

    private (List<float>, List<float>) CalculateNeck(List<float> keypointsX, List<float> keypointsY)
    {
        var leftShoulder = new Vector2(keypointsX[11], keypointsY[11]);
        var rightShoulder = new Vector2(keypointsX[12], keypointsY[12]);

        var neck = new Vector2((leftShoulder.x + rightShoulder.x) / 2, (leftShoulder.y + rightShoulder.y) / 2);

        // add neck to keypoints
        keypointsX.Add(neck.x);
        keypointsY.Add(neck.y);

        return (keypointsX, keypointsY);
    }


    public void AddLandmarks(NormalizedLandmarkList poseLandmarks, NormalizedLandmarkList leftHandLandmarks, NormalizedLandmarkList rightHandLandmarks)
    {
        List<float> pose_x = new List<float>();
        List<float> pose_y = new List<float>();
        List<float> left_hand_x = new List<float>();
        List<float> left_hand_y = new List<float>();
        List<float> right_hand_x = new List<float>();
        List<float> right_hand_y = new List<float>();

        if (poseLandmarks == null || (leftHandLandmarks == null && rightHandLandmarks == null))
        {
            return;
        }

        if (poseLandmarks != null)
        {
            foreach (NormalizedLandmark landmark in poseLandmarks.Landmark)
            {
                pose_x.Add(landmark.X);
                pose_y.Add(landmark.Y);
            }
        }else{
            for (int i = 0; i < 33; i++)
            {
                pose_x.Add(0);
                pose_y.Add(0);
            }
        }

        // Add neck to pose
        (pose_x, pose_y) = CalculateNeck(pose_x, pose_y);

        // normalize pose
        (pose_x, pose_y) = NormalizePose(pose_x, pose_y);

        // now filter the pose keypoints based on the pose indeces
        List<List<float>> filtered_pose = new List<List<float>>();
        foreach (int index in pose_indices)
        {
            filtered_pose.Add(new List<float> { pose_x[index] - 0.5f, pose_y[index]  - 0.5f });
        }

        // add hand landmarks
        if (leftHandLandmarks != null)
        {
            foreach (NormalizedLandmark landmark in leftHandLandmarks.Landmark)
            {
                left_hand_x.Add(landmark.X);
                left_hand_y.Add(landmark.Y);
            }
        }else{
            for (int i = 0; i < 21; i++)
            {
                left_hand_x.Add(0);
                left_hand_y.Add(0);
            }
        }

        if (rightHandLandmarks != null)
        {
            foreach (NormalizedLandmark landmark in rightHandLandmarks.Landmark)
            {
                right_hand_x.Add(landmark.X);
                right_hand_y.Add(landmark.Y);
            }
        }else{
            for (int i = 0; i < 21; i++)
            {
                right_hand_x.Add(0);
                right_hand_y.Add(0);
            }
        }

        // normalize the hands
        (left_hand_x, left_hand_y) = NormalizeHand(left_hand_x, left_hand_y);
        (right_hand_x, right_hand_y) = NormalizeHand(right_hand_x, right_hand_y);

        // now filter the hand keypoints based on the hand indeces
        List<List<float>> filtered_left_hand = new List<List<float>>();
        List<List<float>> filtered_right_hand = new List<List<float>>();

        foreach (int index in hand_indices)
        {
            filtered_left_hand.Add(new List<float> { left_hand_x[index] - 0.5f, left_hand_y[index] - 0.5f });
            filtered_right_hand.Add(new List<float> { right_hand_x[index] - 0.5f, right_hand_y[index] - 0.5f });
        }

        // add the filtered keypoints together in one list
        List<List<float>> filtered_keypoints = new List<List<float>>();
        filtered_keypoints.AddRange(filtered_pose);
        filtered_keypoints.AddRange(filtered_left_hand);
        filtered_keypoints.AddRange(filtered_right_hand);

        keypointsBuffer.Add(filtered_keypoints);

        if (keypointsBuffer.Count > BUFFER_SIZE)
        {
            keypointsBuffer.RemoveAt(0);
        }
    }

    public List<List<List<float>>> GetKeypoints()
    {
        if (keypointsBuffer.Count < BUFFER_SIZE){
            return null;
        }
        return keypointsBuffer;
    }

}