using System;
using UnityEngine;
using System.Linq;

// KD tree
//source: https://github.com/ericreg/Supercluster.KDTree
using Supercluster.KDTree;

// Linear algebra for C#
//source: https://numerics.mathdotnet.com/
using MathNet.Numerics.LinearAlgebra;
//NuGet in Unity editor source: https://github.com/GlitchEnzo/NuGetForUnity

public class ICP5 : MonoBehaviour
{
    /// <summary>Reference to AlignModel class</summary>
    [SerializeField] private AlignModels align;
    /// <summary>Reference to Model class</summary>
    [SerializeField] private RoscanModelUnity model1;
    /// <summary>Save the error</summary>
    private double previousErr = double.MaxValue;
    /// <summary>Save the model min Error</summary>
    private double modelMinErr = double.MaxValue;

    /// <summary>End of Alignment variable</summary>
    public bool endLoop = false;
    public int errorCounter = 0;

    /// <summary> normals of referent model </summary>
    public Matrix<double> normalref;
    /// <summary>new calculated Normals </summary>
    public Matrix<double> resultNormals;

    /// <summary>get rigid tranformation for point-to-plane ICP</summary>
    Matrix<double> point2planeICP(Matrix<double> source, Matrix<double> target, Matrix<double> normal)
    {
        // build A and b
        int pts_num = source.RowCount;
        Matrix<double> A = Matrix<double>.Build.Dense(pts_num, 6);
        Vector<double> b = Vector<double>.Build.Dense(pts_num, 0);
        b.ToRowMatrix();

        for (int idx = 0; idx < source.RowCount; idx++)
        {
            A[idx, 0] = normal[idx, 2] * target[idx, 1] - normal[idx, 1] * target[idx, 2];
            A[idx, 1] = normal[idx, 0] * target[idx, 2] - normal[idx, 2] * target[idx, 0];
            A[idx, 2] = normal[idx, 1] * target[idx, 0] - normal[idx, 0] * target[idx, 1];
            A[idx, 3] = normal[idx, 0];
            A[idx, 4] = normal[idx, 1];
            A[idx, 5] = normal[idx, 2];

            Matrix<double> diff = Matrix<double>.Build.Dense(pts_num, 3);

            diff = target - source;
            b[idx] = 0;

            for (int n = 0; n < 3; n++)
            {
                b[idx] = b[idx] - (diff[idx, n] * normal[idx, n]);
            }
        }

        // solve x
        Vector<double> x = ((A.Transpose() * A).Inverse()) * (A.Transpose()) * b;

        // rotation matrix
        Matrix<double> R = Matrix<double>.Build.Dense(3, 3);
        double sin_a = Math.Sin(x[0]);
        double cos_a = Math.Cos(x[0]);
        double sin_b = Math.Sin(x[1]);
        double cos_b = Math.Cos(x[1]);
        double sin_y = Math.Sin(x[2]);
        double cos_y = Math.Cos(x[2]);

        R[0, 0] = cos_y * cos_b;
        R[0, 1] = -sin_y * cos_a + cos_y * sin_b * sin_a;
        R[0, 2] = sin_y * sin_a + cos_y * sin_b * cos_a;
        R[1, 0] = sin_y * cos_b;
        R[1, 1] = cos_y * cos_a + sin_y * sin_b * sin_a;
        R[1, 2] = -cos_y * sin_a + sin_y * sin_b * cos_a;
        R[2, 0] = -sin_b;
        R[2, 1] = cos_b * sin_a;
        R[2, 2] = cos_b * cos_a;

        // translation vector
        Vector<double> T = Vector<double>.Build.Dense(3, 0);
        T[0] = x[3];
        T[1] = x[4];
        T[2] = x[5];

        Matrix<double> result = Matrix<double>.Build.Dense(3, 4);

        result.SetSubMatrix(0, 0, R);
        result.SetColumn(3, T);

        return result;
    }

    /// <summary>used to get plane normal</summary>
    public Matrix<double> getNormal(Matrix<double> source, Matrix<double> target)
    {
        int pts_num = source.RowCount;
        Matrix<double> normal = Matrix<double>.Build.Dense(pts_num, 3);

        // get meancenter of target mesh
        Vector<double> meanCenterRef = Vector<double>.Build.Dense(3, 1);

        Vector<double> m = target.ColumnSums();

        m.Divide(target.RowCount, meanCenterRef);
        
        for (int idx = 0; idx < pts_num; ++idx)
        {
            Vector3 tempVertex = new Vector3((float)target[idx, 0], (float)target[idx, 1], (float)target[idx, 2]);
            var tempInd = Array.IndexOf(model1.simplVert, tempVertex);

            normal[idx, 0] = model1.simplNorm[tempInd].x;
            normal[idx, 1] = model1.simplNorm[tempInd].y;
            normal[idx, 2] = model1.simplNorm[tempInd].z;

            // check the direction of normal vector
            if ((meanCenterRef[0] - target[idx, 0]) * normal[idx, 0] + (meanCenterRef[1] - target[idx, 1]) * normal[idx, 1] + (meanCenterRef[2] - target[idx, 2]) * normal[idx, 2] > 0)
            {
                normal[idx, 0] = -normal[idx, 0];
                normal[idx, 1] = -normal[idx, 1];
                normal[idx, 2] = -normal[idx, 2];
            }
        }
        return normal;
    }

    /// <summary>used to detect error</summary>
    private bool detectError(Matrix<double> source, Matrix<double> target, Matrix<double> RT)
    {
        Matrix<double> R = RT.SubMatrix(0, 3, 0, 3);
        Vector<double> T = RT.Column(3);
        Vector<double> diff2;

        Matrix<double> sourceTemp = source;

        for (int i = 0; i < source.RowCount; i++)
        {
            sourceTemp[i, 0] = source[i, 0] - T[0];
            sourceTemp[i, 1] = source[i, 1] - T[1];
            sourceTemp[i, 2] = source[i, 2] - T[2];
        }

        source = sourceTemp * R;

        diff2 = (source - target).ColumnSums();

        double error = diff2.L2Norm()/source.RowCount;

        if (error < Math.Exp(-11) || previousErr <= error || modelMinErr < error)
        {
            if (modelMinErr < previousErr || modelMinErr < error || error < Math.Exp(-11))
            {
                endLoop = true;
            }

            modelMinErr = previousErr;

            previousErr = double.MaxValue;

            return false;
        }
        else
        {
            previousErr = error;
            return true;
        }
    }

    /// <summary>
    /// This Function Calculates Final P_points under the conditions of either threshold or maximum iterations that given as input.
    /// Returns matched P_points.
    /// </summary>
    /// <param name="M_points"></param> // changable -v2
    /// <param name="P_points"></param> // ref - v1
    /// <returns></returns>
    public Matrix<double> ICP_run3(Matrix<double> M_points, Matrix<double> P_points, Matrix<double> normal)
    {
        int max_iterations = 25; //changable

        Matrix<double> tResult = align.changeRealPoints.Transpose();
        Matrix<double> tResult2 = align.changeRealPoints.Transpose();
        Matrix<double> result = align.changeRealPoints.Transpose();
        Matrix<double> Y2 = normal;
        Matrix<double> Y = normal;

        for (int itr = 1; itr <= max_iterations; itr++)
        {
            Y = KD_tree(P_points.Transpose(), M_points.Transpose());

            // get normals
            normalref = getNormal(M_points, Y.Transpose());

            Matrix<double> transformNew = point2planeICP(M_points, Y.Transpose(), normalref);

            // get error
            if (detectError(M_points, Y.Transpose(), transformNew))
            {
                Matrix<double> R = transformNew.SubMatrix(0, 3, 0, 3);
                Vector<double> T = transformNew.Column(3);

                Matrix<double> sourceTemp = M_points;
                resultNormals = normalref;

                for (int i = 0; i < M_points.RowCount; i++)
                {
                    sourceTemp[i, 0] = M_points[i, 0] - T[0];
                    sourceTemp[i, 1] = M_points[i, 1] - T[1];
                    sourceTemp[i, 2] = M_points[i, 2] - T[2];
                }

                for (int i = 0; i < align.changeRealPoints.RowCount; i++)
                {
                    tResult2[i, 0] = tResult[i, 0] - T[0];
                    tResult2[i, 1] = tResult[i, 1] - T[1];
                    tResult2[i, 2] = tResult[i, 2] - T[2];
                }

                sourceTemp = sourceTemp * R;
                result = tResult2 * R;

                M_points = sourceTemp;
                tResult = result;
            }
            else { break; }
        }
        return result.Transpose();
    }

    /// <summary>
    /// Calculates Nearest Neighbour for each points and returns distance and coordinate
    /// </summary>
    /// <param name="M_points"></param>
    /// <param name="P_points"></param>
    /// <returns></returns>
    public Matrix<double> KD_tree(Matrix<double> M_points, Matrix<double> P_points)
    {
        int Np = P_points.ColumnCount;

        /// this function is used for calculating the Euclidean norm
        Func<double[], double[], double> L2Norm = (x, y) =>
        {
            double dist = 0;
            for (int i = 0; i < x.Length; i++)
            {
                dist += (x[i] - y[i]) * (x[i] - y[i]);///squre root is not used to speed up
            }
            return dist;
        };
        ///create tree of M_points
        var treeData2 = M_points.ToColumnArrays();
        var treeNodes = treeData2.Select(p => p.ToString()).ToArray();
        var m = Matrix<double>.Build;
        Matrix<double> Y = m.Dense(3, Np, 0);
        Tuple<double[], string>[] test;
        var tree = new KDTree<double, string>(3, treeData2, treeNodes, L2Norm);
        var scan_data = P_points.ToColumnArrays();
        ///Calculate nearest neighbour for every element of cloud
        for (int i = 0; i < scan_data.Length; i++)
        {
            test = tree.NearestNeighbors(scan_data[i], 1);
            Y[0, i] = test[0].Item1[0];
            Y[1, i] = test[0].Item1[1];
            Y[2, i] = test[0].Item1[2];
        }
        return Y;///return result matrix
    }
}
