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

public class ModelInfo
{
    public List<int> pose_landmarks;
    public List<int> hand_landmarks;
}

public class KeypointManager
{

    private ModelInfo modelInfo;
    private List<List<float>> keypointsBuffer;

    public KeypointManager(TextAsset modelInfoFile)
    {
        modelInfo = JsonUtility.FromJson<ModelInfo>(modelInfoFile.text);
        keypointsBuffer = new List<List<float>>();
    }

    private (List<float>, List<float>) NormalizeHand(List<float> hand_x, List<float> hand_y)
    {

        float min_x = hand_x.Min();
        float min_y = hand_y.Min();

        float max_x = hand_x.Max();
        float max_y = hand_y.Max();

        float width = max_x - min_x;
        float height = max_y - min_y;

        if (width == 0 || height == 0)
        {
            return (hand_x, hand_y);
        }

        float center_x = (min_x + max_x) / 2;
        float center_y = (min_y + max_y) / 2;

        List<float> normalized_x = new List<float>();
        List<float> normalized_y = new List<float>();

        for (int i = 0; i < hand_x.Count; i++)
        {
            normalized_x.Add((hand_x[i] - center_x) / width);
            normalized_y.Add((hand_y[i] - center_y) / height);
        }

        return (normalized_x, normalized_y);
    }

    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)
        {
            foreach (int landmark_index in modelInfo.pose_landmarks)
            {
                pose_x.Add(poseLandmarks.Landmark[landmark_index].X);
                pose_y.Add(poseLandmarks.Landmark[landmark_index].Y);
            }
        }
        else
        {
            foreach (int _ in modelInfo.pose_landmarks)
            {
                pose_x.Add(0);
                pose_y.Add(0);
            }
        }



        foreach (int landmark_index in modelInfo.hand_landmarks)
        {
            if (leftHandLandmarks == null)
            {
                left_hand_x.Add(0);
                left_hand_y.Add(0);
            }
            else
            {
                left_hand_x.Add(leftHandLandmarks.Landmark[landmark_index].X);
                left_hand_y.Add(leftHandLandmarks.Landmark[landmark_index].Y);
            }
            if (rightHandLandmarks == null)
            {
                right_hand_x.Add(0);
                right_hand_y.Add(0);
            }
            else
            {
                right_hand_x.Add(rightHandLandmarks.Landmark[landmark_index].X);
                right_hand_y.Add(rightHandLandmarks.Landmark[landmark_index].Y);
            }
        }

        // TODO: Add normalization
        //Debug.Log($"pose_landMarks = [{modelInfo.pose_landmarks.Aggregate("   ", (t, f) => $"{t}   {f}")}]");
        //Debug.Log($"hand_landmarks = [{modelInfo.hand_landmarks.Aggregate("   ", (t, f) => $"{t}   {f}")}]");
        (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);


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

        List<float> keypoints = new List<float>();
        for (int i = 0; i < pose_x.Count; i++)
        {
            keypoints.Add(pose_x[i]);
            keypoints.Add(pose_y[i]);
        }
        for (int i = 0; i < left_hand_x.Count; i++)
        {
            keypoints.Add(left_hand_x[i]);
            keypoints.Add(left_hand_y[i]);
        }
        for (int i = 0; i < right_hand_x.Count; i++)
        {
            keypoints.Add(right_hand_x[i]);
            keypoints.Add(right_hand_y[i]);
        }

        keypointsBuffer.Add(keypoints);
    }

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