using System.Collections.Generic;
using UnityEngine;

namespace LibMMD.Motion
{
    public class MmdMotion
    {
        public string Name { get; set; }

        public int Length { get; set; }

        public Dictionary<string, List<KeyValuePair<int, BoneKeyframe>>> BoneMotions { get; set; }
        public Dictionary<string, List<KeyValuePair<int, MorphKeyframe>>> MorphMotions { get; set; }
        public Dictionary<string, List<KeyValuePair<int, bool>>> IkMotions { get; set; }

        public MmdMotion()
        {
            BoneMotions = new Dictionary<string, List<KeyValuePair<int, BoneKeyframe>>>();
            MorphMotions = new Dictionary<string, List<KeyValuePair<int, MorphKeyframe>>>();
            IkMotions = new Dictionary<string, List<KeyValuePair<int, bool>>>();
        }

        public BonePose GetBonePose(string boneName, int frame)
        {
            List<KeyValuePair<int, BoneKeyframe>> keyFrames;
            BoneMotions.TryGetValue(boneName, out keyFrames);
            if (keyFrames == null || keyFrames.Count == 0)
            {
                return new BonePose
                {
                    Translation = Vector3.zero,
                    Rotation = Quaternion.identity
                };
            }

            if (keyFrames[0].Key >= frame)
            {
                var key = keyFrames[0].Value;
                return new BonePose
                {
                    Translation = key.Translation,
                    Rotation = key.Rotation
                };
            }

            if (keyFrames[keyFrames.Count - 1].Key <= frame)
            {
                var key = keyFrames[keyFrames.Count - 1].Value;
                return new BonePose
                {
                    Translation = key.Translation,
                    Rotation = key.Rotation
                };
            }

            var toSearch = new KeyValuePair<int, BoneKeyframe>(frame, null);
            var rightBoundIndex = keyFrames.BinarySearch(toSearch, BoneKeyframeSearchComparator.Instance);
            if (rightBoundIndex < 0)
            {
                rightBoundIndex = ~rightBoundIndex;
            }
            int leftBoundIndex;
            if (rightBoundIndex == 0)
            {
                leftBoundIndex = 0;
            } else if (rightBoundIndex >= keyFrames.Count)
            {
                rightBoundIndex = leftBoundIndex = keyFrames.Count - 1;
            }
            else
            {
                leftBoundIndex = rightBoundIndex - 1;
            }
            var rightBound = keyFrames[rightBoundIndex];
            var rightFrame = rightBound.Key;
            var rightKey = rightBound.Value;
            var leftBound = keyFrames[leftBoundIndex];
            var leftFrame = leftBound.Key;
            var leftKey = leftBound.Value;
            if (leftFrame == rightFrame)
            {
                return new BonePose
                {
                    Translation = leftKey.Translation,
                    Rotation = leftKey.Rotation
                };
            }
            var baryPos = (frame - leftFrame) / (float) (rightFrame - leftFrame);
            var translation = new Vector3();
            // In VMD, a key stores the Bezier curve for the segment that arrives at that
            // key from the previous one. Therefore the right/destination key owns all four
            // interpolators for this interval. Using the left key is barely visible in dense
            // motions but badly distorts sparse center/groove and leg tracks.
            var lambda = rightKey.XInterpolator.Calculate(baryPos);
            translation.x = leftKey.Translation.x * (1 - lambda) + rightKey.Translation.x * lambda;
            lambda = rightKey.YInterpolator.Calculate(baryPos);
            translation.y = leftKey.Translation.y * (1 - lambda) + rightKey.Translation.y * lambda;
            lambda = rightKey.ZInterpolator.Calculate(baryPos);
            translation.z = leftKey.Translation.z * (1 - lambda) + rightKey.Translation.z * lambda;
            lambda = rightKey.RInterpolator.Calculate(baryPos);
            var rotation = Quaternion.Slerp(leftKey.Rotation, rightKey.Rotation, lambda);
            return new BonePose
            {
                Translation = translation,
                Rotation = rotation
            };
        }

        public BonePose GetBonePose(string boneName, double time)
        {
            List<KeyValuePair<int, BoneKeyframe>> keyFrames;
            BoneMotions.TryGetValue(boneName, out keyFrames);
            if (keyFrames == null || keyFrames.Count == 0)
            {
                return new BonePose
                {
                    Translation = Vector3.zero,
                    Rotation = Quaternion.identity
                };
            }

            var dFrame = time * 30.0;

            if (keyFrames[0].Key >= dFrame)
            {
                var key = keyFrames[0].Value;
                return new BonePose
                {
                    Translation = key.Translation,
                    Rotation = key.Rotation
                };
            }

            if (keyFrames[keyFrames.Count - 1].Key <= dFrame)
            {
                var key = keyFrames[keyFrames.Count - 1].Value;
                return new BonePose
                {
                    Translation = key.Translation,
                    Rotation = key.Rotation
                };
            }

            // Keep the fractional frame when selecting the surrounding keys. Casting dFrame
            // to int makes an exact key at floor(dFrame) look like the right-hand key; the
            // interpolation factor then exceeds 1 until the next whole frame and snaps back.
            var rightBoundIndex = LowerBound(keyFrames, dFrame);
            int leftBoundIndex;
            if (rightBoundIndex == 0)
            {
                leftBoundIndex = 0;
            } else if (rightBoundIndex >= keyFrames.Count)
            {
                rightBoundIndex = leftBoundIndex = keyFrames.Count - 1;
            }
            else
            {
                leftBoundIndex = rightBoundIndex - 1;
            }
            var rightBound = keyFrames[rightBoundIndex];
            var rightFrame = rightBound.Key;
            var rightKey = rightBound.Value;
            var leftBound = keyFrames[leftBoundIndex];
            var leftFrame = leftBound.Key;
            var leftKey = leftBound.Value;
            if (leftFrame == rightFrame)
            {
                return new BonePose
                {
                    Translation = leftKey.Translation,
                    Rotation = leftKey.Rotation
                };
            }
            var baryPos = (float) (dFrame - leftFrame) / (rightFrame - leftFrame);
            var translation = new Vector3();
            var lambda = rightKey.XInterpolator.Calculate(baryPos);
            translation.x = leftKey.Translation.x * (1 - lambda) + rightKey.Translation.x * lambda;
            lambda = rightKey.YInterpolator.Calculate(baryPos);
            translation.y = leftKey.Translation.y * (1 - lambda) + rightKey.Translation.y * lambda;
            lambda = rightKey.ZInterpolator.Calculate(baryPos);
            translation.z = leftKey.Translation.z * (1 - lambda) + rightKey.Translation.z * lambda;
            lambda = rightKey.RInterpolator.Calculate(baryPos);
            var rotation = Quaternion.Slerp(leftKey.Rotation, rightKey.Rotation, lambda);
            return new BonePose
            {
                Translation = translation,
                Rotation = rotation
            };
        }

        public MorphPose GetMorphPose(string morphName, int frame)
        {
            List<KeyValuePair<int, MorphKeyframe>> keyFrames;
            MorphMotions.TryGetValue(morphName, out keyFrames);
            if (keyFrames == null || keyFrames.Count == 0)
            {
                return new MorphPose(0.0f);
            }
            if (keyFrames[0].Key >= frame)
            {
                var key = keyFrames[0].Value;
                return new MorphPose(key.Weight);
            }

            if (keyFrames[keyFrames.Count - 1].Key <= frame)
            {
                var key = keyFrames[keyFrames.Count - 1].Value;
                return new MorphPose(key.Weight);
            }

            var toSearch = new KeyValuePair<int, MorphKeyframe>(frame, null);
            var rightBoundIndex = keyFrames.BinarySearch(toSearch, MorphKeyframeSearchComparator.Instance);
            if (rightBoundIndex < 0)
            {
                rightBoundIndex = ~rightBoundIndex;
            }
            int leftBoundIndex;
            if (rightBoundIndex == 0)
            {
                leftBoundIndex = 0;
            } else if (rightBoundIndex >= keyFrames.Count)
            {
                rightBoundIndex = leftBoundIndex = keyFrames.Count - 1;
            }
            else
            {
                leftBoundIndex = rightBoundIndex - 1;
            }
            var rightBound = keyFrames[rightBoundIndex];
            var rightFrame = rightBound.Key;
            var rightKey = rightBound.Value;
            var leftBound = keyFrames[leftBoundIndex];
            var leftFrame = leftBound.Key;
            var leftKey = leftBound.Value;
            if (leftFrame == rightFrame)
            {
                return new MorphPose(leftKey.Weight);
            }
            var baryPos = (frame - leftFrame) / (float) (rightFrame - leftFrame);
            var lWeight = leftKey.Weight;
            var rWeight = rightKey.Weight;
            var lambda = leftKey.WInterpolator.Calculate(baryPos);

            return new MorphPose(lWeight * (1 - lambda) + rWeight * lambda);
        }

        public MorphPose GetMorphPose(string morphName, double time)
        {
            List<KeyValuePair<int, MorphKeyframe>> keyFrames;
            MorphMotions.TryGetValue(morphName, out keyFrames);
            if (keyFrames == null || keyFrames.Count == 0)
            {
                return new MorphPose(0.0f);
            }

            var iFrame = time * 30.0;

            if (keyFrames[0].Key >= iFrame)
            {
                var key = keyFrames[0].Value;
                return new MorphPose(key.Weight);
            }

            if (keyFrames[keyFrames.Count - 1].Key <= iFrame)
            {
                var key = keyFrames[keyFrames.Count - 1].Value;
                return new MorphPose(key.Weight);
            }

            var rightBoundIndex = LowerBound(keyFrames, iFrame);
            int leftBoundIndex;
            if (rightBoundIndex == 0)
            {
                leftBoundIndex = 0;
            } else if (rightBoundIndex >= keyFrames.Count)
            {
                rightBoundIndex = leftBoundIndex = keyFrames.Count - 1;
            }
            else
            {
                leftBoundIndex = rightBoundIndex - 1;
            }
            var rightBound = keyFrames[rightBoundIndex];
            var rightFrame = rightBound.Key;
            var rightKey = rightBound.Value;
            var leftBound = keyFrames[leftBoundIndex];
            var leftFrame = leftBound.Key;
            var leftKey = leftBound.Value;
            if (leftFrame == rightFrame)
            {
                return new MorphPose(leftKey.Weight);
            }
            var baryPos = (float) (iFrame - leftFrame) / (rightFrame - leftFrame);
            var lWeight = leftKey.Weight;
            var rWeight = rightKey.Weight;
            var lambda = leftKey.WInterpolator.Calculate(baryPos);

            return new MorphPose(lWeight * (1 - lambda) + rWeight * lambda);
        }

        public bool IsBoneRegistered(string boneName)
        {
            return BoneMotions.ContainsKey(boneName);
        }

        public bool IsMorphRegistered(string morphName)
        {
            return MorphMotions.ContainsKey(morphName);
        }

        public bool IsIkRegistered(string boneName)
        {
            return IkMotions.ContainsKey(boneName);
        }

        public bool GetIkEnabled(string boneName, int frame)
        {
            return GetIkEnabled(boneName, (double) frame);
        }

        public bool GetIkEnabled(string boneName, double frame)
        {
            if (!IkMotions.TryGetValue(boneName, out var keyFrames) || keyFrames.Count == 0)
                return true;

            var index = LowerBound(keyFrames, frame);
            if (index < keyFrames.Count && keyFrames[index].Key <= frame)
                return keyFrames[index].Value;
            if (index == 0) return true;
            return keyFrames[index - 1].Value;
        }

        private static int LowerBound<T>(List<KeyValuePair<int, T>> keyFrames, double frame)
        {
            var first = 0;
            var count = keyFrames.Count;
            while (count > 0)
            {
                var step = count / 2;
                var index = first + step;
                if (keyFrames[index].Key < frame)
                {
                    first = index + 1;
                    count -= step + 1;
                }
                else
                {
                    count = step;
                }
            }
            return first;
        }

        private class BoneKeyframeSearchComparator : IComparer<KeyValuePair<int, BoneKeyframe>>
        {
            public static readonly BoneKeyframeSearchComparator Instance = new BoneKeyframeSearchComparator();

            private BoneKeyframeSearchComparator()
            {
            }

            public int Compare(KeyValuePair<int, BoneKeyframe> x, KeyValuePair<int, BoneKeyframe> y)
            {
                return x.Key.CompareTo(y.Key);
            }
        }

        private class MorphKeyframeSearchComparator : IComparer<KeyValuePair<int, MorphKeyframe>>
        {
            public static readonly MorphKeyframeSearchComparator Instance = new MorphKeyframeSearchComparator();

            private MorphKeyframeSearchComparator()
            {
            }

            public int Compare(KeyValuePair<int, MorphKeyframe> x, KeyValuePair<int, MorphKeyframe> y)
            {
                return x.Key.CompareTo(y.Key);
            }
        }
    }
}