using Psychomix.EngineMath;
using Psychomix.Mesh;
using SharpDX;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Psychomix.Animation
{
    using SharpDX.Toolkit;
    using SharpDX.Toolkit.Graphics;
    using System.Diagnostics;
    using TribeGame.Animation;

    public class SkinningModel
    {
        public class SkinnedBone
        {
            public string Name;
            public int IndexGlobal;
            public SRTMatrix BindTransform;
            public SRTMatrix DefaultTransform;
        }

        public class Skeleton
        {
            public SkinnedBone[] Bones;
            public int[][] Hierarchy;
            public int RootBone;

            public int[] HierarchyOrderAccess;

            private Dictionary<string, int> _indexHashtable;

            public void BuildIndexHash()
            {
                _indexHashtable = new Dictionary<string, int>();
                for (int i = 0; i < Bones.Length; i++)
                {
                    _indexHashtable.Add(Bones[i].Name, i);
                }
            }

            public int GetBoneIndex(string name)
            {
                if(!_indexHashtable.ContainsKey(name))
                    return -1;

                return _indexHashtable[name];
            }

            public Matrix[] LocalToWorld(Matrix[] array)
            {
                Matrix[] matrix = new Matrix[array.Length];

                for (int i = 0; i < matrix.Length; i++)
                {
                    int[] hierarchy = Hierarchy[i];

                    matrix[i] = array[i];

                    for (int h = 0; h < hierarchy.Length; h++)
                    {
                        matrix[i] *= array[hierarchy[h]];
                    }
                }

                return matrix;
            }

            public Matrix[] WorldToLocal(Matrix[] array)
            {
                Matrix[] matrix = new Matrix[array.Length];

                for (int i = 0; i < matrix.Length; i++)
                {
                    int boneID = HierarchyOrderAccess[i];
                    int[] hierarchy = Hierarchy[boneID];

                    Matrix localMatrix = array[boneID];
                    for (int h = 0; h < hierarchy.Length; h++)
                    {
                        Matrix invertH = Matrix.Invert(matrix[hierarchy[hierarchy.Length - h - 1]]);
                        localMatrix *= invertH;
                    }

                    matrix[boneID] = localMatrix;
                }

                return matrix;
            }
        }

        public class GeometryMesh
        {
            public string Name;
            public SkinnedBone[] SkinnedBones;
            public Buffer<SkinningVertex> VertexBuffer;
            public Buffer<uint> IndexBuffer;
            public int IndexCount;
        }

        public class Mesh
        {
            //TODO: NEED HIERARCHY
            public SRTMatrix ModelWorldTransform;

            public string Name;
            public GeometryMesh[] GeometryMeshes;
        }

        private class BonesPresentation
        {
            public List<string> Names;
            public int BoneCount { get { return Names.Count; } }
            public BonesPresentation()
            {
                Names = new List<string>();
            }

            public void AddBone(string name)
            {
                if (!Names.Contains(name))
                    Names.Add(name);
            }

            public int FindBone(string name)
            {
                for (int i = 0; i < Names.Count; i++)
                    if (Names[i].Equals(name)) return i;
                return -1;
            }
        }

        public GeometryMesh[] _geometryMeshes;
        public Mesh[] _modelMeshes;
        public Skeleton _modelSkeleton;
        public AnimationClip[] _animationClips;

        //TODO: SKINNING MODEL DRAW, OTHER METHODS
        public int[] BuildChannelMap(AnimationClip clip)
        {
            int[] channelMap = new int[clip.Bones.Length];

            for (int i = 0; i < clip.Bones.Length; i++)
            {
                channelMap[i] = -1;

                foreach (SkinningModel.SkinnedBone bone in _modelSkeleton.Bones)
                {
                    if (bone.Name == clip.Bones[i].Name)
                    {
                        channelMap[i] = bone.IndexGlobal;
                        break;
                    }
                }
            }

            return channelMap;
        }

        //private Mesh
        public static Assimp.PostProcessSteps DefaultPPS = Assimp.PostProcessSteps.FlipUVs |
                                                Assimp.PostProcessSteps.FlipWindingOrder |
                                                Assimp.PostProcessSteps.LimitBoneWeights |
                                                Assimp.PostProcessSteps.SplitByBoneCount |
                                                Assimp.PostProcessSteps.CalculateTangentSpace |
                                                Assimp.PostProcessSteps.Triangulate;

        public static SkinningModel Load(GraphicsDevice graphics, string path)
        {
            SkinningModel model = new SkinningModel();

            {
                Stopwatch stopwatch = new Stopwatch();

                stopwatch.Start();

                Assimp.AssimpImporter importer = new Assimp.AssimpImporter();
                Assimp.PostProcessSteps steps = DefaultPPS;


                Assimp.Scene scene = importer.ImportFile(path, steps);

                stopwatch.Stop();
                Debug.WriteLine("Time (SkinnedModel) [ASSIMP-LOAD]: " + stopwatch.Elapsed.TotalMilliseconds + "ms");
                stopwatch.Restart();

                if (!scene.HasAnimations)
                    throw new InvalidOperationException("Model don't have animations!");

                BonesPresentation bp = new BonesPresentation();
                for (int m = 0; m < scene.MeshCount; m++)
                    foreach (Assimp.Bone assimpBone in scene.Meshes[m].Bones)
                        //if (assimpBone.HasVertexWeights)
                        {
                            bp.AddBone(assimpBone.Name);
                        }

                model._modelSkeleton = new Skeleton();
                model._modelSkeleton.Bones = new SkinnedBone[bp.BoneCount];

                /* GEOMETRY MESH LOAD */
                model._geometryMeshes = new GeometryMesh[scene.Meshes.Length];

                for (int m = 0; m < model._geometryMeshes.Length; m++)
                {
                    Assimp.Mesh assimpMesh = scene.Meshes[m];

                    if(assimpMesh.Vertices.Length != assimpMesh.Normals.Length ||
                       assimpMesh.Vertices.Length != assimpMesh.Tangents.Length ||
                       assimpMesh.Vertices.Length != assimpMesh.BiTangents.Length)
                        throw new InvalidOperationException("Problem with count Vertices and TBN data!");

                    Assimp.Vector3D[] uvChannel = assimpMesh.GetTextureCoords(0);

                    SkinningVertex[] vertices = new SkinningVertex[assimpMesh.Vertices.Length];
                    uint[] indices = new uint[assimpMesh.Faces.Length * 3];

                    for (int i = 0; i < assimpMesh.Faces.Length; i++)
                    {
                        indices[3 * i + 0] = assimpMesh.Faces[i].Indices[0];
                        indices[3 * i + 1] = assimpMesh.Faces[i].Indices[1];
                        indices[3 * i + 2] = assimpMesh.Faces[i].Indices[2];
                    }

                    /* PROCESS BLEND IW */
                    bool hasWeights = false;
                    int[] skinningCount = new int[assimpMesh.VertexCount];
                    Int4[] skinningIndices = new Int4[assimpMesh.VertexCount];
                    Vector4[] skinningWeights = new Vector4[assimpMesh.VertexCount];
                    List<SkinnedBone> skinnedBones = new List<SkinnedBone>();

                    for (int b = 0; b < assimpMesh.Bones.Length; b++)
                    {
                        Assimp.Bone assimpBone = assimpMesh.Bones[b];

                            int skinnedBoneIndex = skinnedBones.Count;

                            SkinnedBone bone = new SkinnedBone();
                            bone.Name = assimpBone.Name;
                            bone.IndexGlobal = bp.FindBone(bone.Name);
                            System.Diagnostics.Debug.Assert(bone.IndexGlobal != -1);

                            bone.BindTransform = MathCrossoverHelper._srtm(assimpBone.OffsetMatrix);
                            model._modelSkeleton.Bones[bone.IndexGlobal] = bone;

                            if (assimpBone.HasVertexWeights)
                            {
                                skinnedBones.Add(bone);

                                for (int j = 0; j < assimpBone.VertexWeightCount; j++)
                                {
                                    var weights = assimpBone.VertexWeights[j];
                                    var vertexSkinningCount = skinningCount[weights.VertexID];

                                    skinningIndices[weights.VertexID][vertexSkinningCount] = skinnedBoneIndex;
                                    skinningWeights[weights.VertexID][vertexSkinningCount] = weights.Weight;
                                    skinningCount[weights.VertexID] = ++vertexSkinningCount;
                                }

                                hasWeights = true;
                            }
                    }

                    /* VERTICES LOAD */
                    for (int i = 0; i < vertices.Length; i++)
                    {
                        Vector3 position = MathCrossoverHelper._vc3(assimpMesh.Vertices[i]);
                        Vector3 tangent = MathCrossoverHelper._vc3(assimpMesh.Tangents[i]);
                        Vector3 bitangent = MathCrossoverHelper._vc3(assimpMesh.BiTangents[i]);
                        Vector3 normal = MathCrossoverHelper._vc3(assimpMesh.Normals[i]);
                        Vector2 uv = MathCrossoverHelper._vc2(uvChannel[i]);

                        vertices[i] = new SkinningVertex(position, uv, tangent, bitangent, normal);

                        System.Diagnostics.Debug.Assert(hasWeights);
                        if (hasWeights)
                        {
                            Int4 blendIndices = skinningIndices[i];
                            Vector4 blendWeight = skinningWeights[i];

                            vertices[i].BlendIndices = blendIndices;
                            vertices[i].BlendWeights = blendWeight;
                        }
                    }


                    GeometryMesh mesh = new GeometryMesh();
                    mesh.Name = assimpMesh.Name;
                    mesh.SkinnedBones = skinnedBones.ToArray();

                    // TODO: PACK VERTICES
                    mesh.VertexBuffer = Buffer.New<SkinningVertex>(graphics, vertices, BufferFlags.VertexBuffer, SharpDX.Direct3D11.ResourceUsage.Immutable);
                    mesh.IndexBuffer = Buffer.New<uint>(graphics, indices, BufferFlags.IndexBuffer, SharpDX.Direct3D11.ResourceUsage.Immutable);

                    mesh.IndexCount = indices.Length;

                    model._geometryMeshes[m] = mesh;
                }

                model._modelSkeleton.BuildIndexHash();

                /* LOAD MESHES */
                model._modelMeshes = _collectMesh(scene, ref model._geometryMeshes);

                model._modelSkeleton.Hierarchy = new int[bp.BoneCount][];

                /* LOAD SKELETON HIERARCHY */
                for (int i = 0; i < model._modelSkeleton.Bones.Length; i++)
                {
                    SkinnedBone bone = model._modelSkeleton.Bones[i];

                    Assimp.Node bnode = _recursivityFindNode(scene.RootNode, bone.Name);

                    System.Diagnostics.Debug.Assert(bnode != null);

                    bone.DefaultTransform = MathCrossoverHelper._srtm(bnode.Transform);

                    List<int> hierarchy = new List<int>();

                    while (bnode.Parent != null)
                    {
                        SkinnedBone parent = null;
                        /* FIND IN SKINNED CHUNK */
                        for (int j = 0; j < model._modelSkeleton.Bones.Length; j++)
                        {
                            if (model._modelSkeleton.Bones[j].Name == bnode.Parent.Name)
                            {
                                parent = model._modelSkeleton.Bones[j];
                                parent.DefaultTransform = MathCrossoverHelper._srtm(bnode.Parent.Transform);

                                break;
                            }
                        }

                        if (parent == null)
                            break;


                        hierarchy.Add(parent.IndexGlobal);

                        bnode = bnode.Parent;
                    }

                    model._modelSkeleton.Hierarchy[i] = hierarchy.ToArray();
                }

                model._modelSkeleton.HierarchyOrderAccess = new int[model._modelSkeleton.Hierarchy.Length];
                for (int i = 0; i < model._modelSkeleton.HierarchyOrderAccess.Length; i++)
                {
                    model._modelSkeleton.HierarchyOrderAccess[i] = i;
                }

                model._modelSkeleton.HierarchyOrderAccess = model._modelSkeleton.HierarchyOrderAccess.OrderBy(
                    c => model._modelSkeleton.Hierarchy[c].Length).ToArray();

                model._animationClips = AnimationClip.LoadAnimations(scene);

                Debug.WriteLine("Time (SkinnedModel) [PROCESSED]: " + stopwatch.Elapsed.TotalMilliseconds + "ms");

            }


            return model;
        }

        private static Mesh[] _collectMesh(Assimp.Scene scene, ref GeometryMesh[] geometryMeshes)
        {
            List<Mesh> meshes = new List<Mesh>();
            _recursivityMeshCollect(scene.RootNode, ref meshes, ref geometryMeshes, Matrix.Identity);
            return meshes.ToArray();
        }

        private static void _recursivityMeshCollect(Assimp.Node node, ref List<Mesh> meshes, ref GeometryMesh[] geometryMeshes, Matrix previousMatrix)
        {
            if (node.HasMeshes)
            {
                Mesh mesh = new Mesh();
                SRTMatrix localTransform = MathCrossoverHelper._srtm(node.Transform);
                mesh.ModelWorldTransform = previousMatrix * localTransform;

                mesh.Name = node.Name;

                mesh.GeometryMeshes = new GeometryMesh[node.MeshCount];
                for (int i = 0; i < node.MeshCount; i++)
                    mesh.GeometryMeshes[i] = geometryMeshes[node.MeshIndices[i]];

                meshes.Add(mesh);
            }

            if (node.HasChildren)
            {
                foreach (Assimp.Node children in node.Children)
                    _recursivityMeshCollect(children, ref meshes, ref geometryMeshes, MathCrossoverHelper._srtm(node.Transform));
            }
        }

        private static Assimp.Node _recursivityFindNode(Assimp.Node node, string nodeName)
        {
            if (node.Name == nodeName)
                return node;

            if (node.HasChildren)
            {
                foreach (Assimp.Node children in node.Children)
                {
                    Assimp.Node result = _recursivityFindNode(children, nodeName);
                    if (result != null)
                        return result;
                }
            }

            return null;
        }
    }
}