#include <stdio.h>
#include <vector>
#include "optimize-gsl.h"
#include "closest-point-vtk.h"
#include "lidardata.h"
#include "mathutil.h"


static LidarPointCloud g_source;
static LidarPointCloud g_target;
static double g_sourceCentroid[3];
static bool g_translationOnly;
static bool g_rotationOnly;

// The purpose of the following factor is as follows. When solving the optimization
// we want the values of the independent variables to be approximately the same size.
// Now the 4 elements of the rotation quaternion are usually less then 1. However, if
// the camera is far away from the asteroid, then the values of the spacecraft position
// may be orders of magnitude greater than the quaternion values. This difference can
// mess up the optimization. Therefore, before starting the optmization, we divide the
// position by g_MULT_FACTOR so as to make the values on the order of 1. g_MULT_FACTOR is
// set to the norma of the position. Then before each time we make use of the position,
// we multiply it by g_MULT_FACTOR to get back the correct value. This way the optimization
// sees all the independent variables as having approximately the same size.
static double g_MULT_FACTOR = 1.0;

/* Return distance squared between points x and y */
static double vdist2(const double x[3], const double y[3])
{
    return ( ( x[0] - y[0] ) * ( x[0] - y[0] )
             + ( x[1] - y[1] ) * ( x[1] - y[1] )
             + ( x[2] - y[2] ) * ( x[2] - y[2] ) );
}

/* Add vectors v1 and v2 and store result in vout */
static void vadd(const double v1[3], const double v2[3], double vout[3])
{
    vout[0] = v1[0] + v2[0];
    vout[1] = v1[1] + v2[1];
    vout[2] = v1[2] + v2[2];
}

static void centroid(const LidarPointCloud& points, double* centroid)
{
    centroid[0] = 0.0;
    centroid[1] = 0.0;
    centroid[2] = 0.0;
    int n = points.size();
    int i;
    for (i=0; i<n; ++i)
    {
        vadd(centroid, points[i].targetpos, centroid);
    }

    centroid[0] /= n;
    centroid[1] /= n;
    centroid[2] /= n;
}

static void findAllClosestPoints(const double* transformation)
{
    const double* translation = transformation;
    const double* rotation = &transformation[3];

    g_target.clear();
    int n = g_source.size();
    int found;
    for (int i=0; i<n; ++i)
    {
        Point s = g_source[i];

        transformPoint(s.targetpos, g_sourceCentroid, translation, rotation, s.targetpos);

        Point closestPoint;
        findClosestPointVtk(s.targetpos, closestPoint.targetpos, &found);
        g_target.push_back(closestPoint);
    }
}

static double func(const double* transformation, void*)
{
    const double translation[3] = {
        g_MULT_FACTOR * transformation[0],
        g_MULT_FACTOR * transformation[1],
        g_MULT_FACTOR * transformation[2] };
    double rotation[4];

    normalizeQuaternion(&transformation[3], rotation);

    int n = g_source.size();
    double ssd = 0.0;
    for (int i=0; i<n; ++i)
    {
        Point s = g_source[i];
        Point t = g_target[i];

        transformPoint(s.targetpos, g_sourceCentroid, translation, rotation, s.targetpos);

        double dist2 = vdist2(s.targetpos, t.targetpos);
        ssd += dist2;
    }

    return ssd;
}

static void grad(const double* transformation, double* df, void*)
{
    const double* translation = transformation;
    const double* rotation = &transformation[3];

    double x = translation[0];
    double y = translation[1];
    double z = translation[2];
    double q0 = rotation[0];
    double q1 = rotation[1];
    double q2 = rotation[2];
    double q3 = rotation[3];

    double cor0 = g_sourceCentroid[0];
    double cor1 = g_sourceCentroid[1];
    double cor2 = g_sourceCentroid[2];

    double m = g_MULT_FACTOR;

    double q12 = pow( (q1*q1)+(q2*q2)+(q3*q3)+(q0*q0),-(1.0/2.0));
    double q32 = pow( (q1*q1)+(q2*q2)+(q3*q3)+(q0*q0),-(3.0/2.0));

    for (int i=0; i<7; ++i)
        df[0] = 0.0;

    int n = g_source.size();
    for (int i=0; i<n; ++i)
    {
        const Point& s = g_source[i];
        const Point& t = g_target[i];

        double s0 = s.targetpos[0];
        double s1 = s.targetpos[1];
        double s2 = s.targetpos[2];
        double t0 = t.targetpos[0];
        double t1 = t.targetpos[1];
        double t2 = t.targetpos[2];

        // The following crazy formulas for gradient components have been autogenerated by computer algebra system called GiNaC (http://www.ginac.de/).
        // See the file gradient.ginsh for the ginsh script used to generate these formulas.

        if (g_rotationOnly == false)
        {
            // df/dx
            df[0] += -2.0*m*( s0+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q12+-2.0*cor0-x*m+t0+-2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1);

            // df/dy
            df[1] += -2.0*( t1+s1+-2.0*q0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+-2.0*q12*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)-m*y+-2.0*cor1)*m;

            // df/dz
            df[2] += -2.0*( s2+-2.0*q3*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+t2+-2.0*cor2-z*m)*m;
        }

        if (g_translationOnly == false)
        {
            // df/dq0
            df[3] += 2.0*( 2.0*q3*q12*( q0*q32*( s1-cor1)*q2+q3*q0*q32*( s2-cor2)+q0*q32*( s0-cor0)*q1)+2.0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+2.0*q3*q0*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+-2.0*( q12*( s2-cor2)-q0*q32*q2*( s0-cor0)+q0*q32*( s1-cor1)*q1-(q0*q0)*q32*( s2-cor2))*q0*q12+-2.0*(q0*q0)*q32*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0)))*( s2+-2.0*q3*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+t2+-2.0*cor2-z*m)+2.0*( t1+s1+-2.0*q0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+-2.0*q12*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)-m*y+-2.0*cor1)*( 2.0*q0*q32*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*(q0*q0)*q32*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+2.0*q0*q12*( (q0*q0)*q32*( s1-cor1)-q3*q0*q32*( s0-cor0)-q12*( s1-cor1)+q0*q32*( s2-cor2)*q1)+2.0*q12*q2*( q0*q32*( s1-cor1)*q2+q3*q0*q32*( s2-cor2)+q0*q32*( s0-cor0)*q1)+-2.0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1)))+2.0*( 2.0*q0*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1+2.0*( q3*q0*q32*( s1-cor1)-q12*( s0-cor0)-q0*q32*q2*( s2-cor2)+(q0*q0)*q32*( s0-cor0))*q0*q12+2.0*q12*( q0*q32*( s1-cor1)*q2+q3*q0*q32*( s2-cor2)+q0*q32*( s0-cor0)*q1)*q1+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q12+2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*(q0*q0)*q32)*( s0+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q12+-2.0*cor0-x*m+t0+-2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1);

            // df/dq1
            df[4] += 2.0*( -2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*(q1*q1)+2.0*( q3*q32*( s1-cor1)*q1-q32*q2*( s2-cor2)*q1+q0*q32*( s0-cor0)*q1)*q0*q12+2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q32*q1+2.0*q12*( q3*q32*( s2-cor2)*q1-q12*( s0-cor0)+q32*( s1-cor1)*q2*q1+q32*( s0-cor0)*(q1*q1))*q1)*( s0+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q12+-2.0*cor0-x*m+t0+-2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1)+2.0*( 2.0*( q32*q2*( s0-cor0)*q1-q32*( s1-cor1)*(q1*q1)+q12*( s1-cor1)+q0*q32*( s2-cor2)*q1)*q0*q12+-2.0*q0*q32*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))*q1+2.0*q3*q12*( q3*q32*( s2-cor2)*q1-q12*( s0-cor0)+q32*( s1-cor1)*q2*q1+q32*( s0-cor0)*(q1*q1))+2.0*q3*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1)*( s2+-2.0*q3*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+t2+-2.0*cor2-z*m)+2.0*( 2.0*q0*q32*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))*q1+2.0*q12*q2*( q3*q32*( s2-cor2)*q1-q12*( s0-cor0)+q32*( s1-cor1)*q2*q1+q32*( s0-cor0)*(q1*q1))+2.0*q32*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1+-2.0*( q3*q32*( s0-cor0)*q1+q12*( s2-cor2)-q32*( s2-cor2)*(q1*q1)-q0*q32*( s1-cor1)*q1)*q0*q12)*( t1+s1+-2.0*q0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+-2.0*q12*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)-m*y+-2.0*cor1);

            // df/dq2
            df[5] += 2.0*( 2.0*q0*q12*( q12*( s2-cor2)+q3*q32*( s1-cor1)*q2+q0*q32*q2*( s0-cor0)-q32*(q2*q2)*( s2-cor2))+2.0*q32*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1+2.0*q12*( q32*q2*( s0-cor0)*q1+q3*q32*q2*( s2-cor2)-q12*( s1-cor1)+q32*( s1-cor1)*(q2*q2))*q1+2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q32*q2)*( s0+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q12+-2.0*cor0-x*m+t0+-2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1)+2.0*( s2+-2.0*q3*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+t2+-2.0*cor2-z*m)*( -2.0*q0*q12*( q12*( s0-cor0)-q0*q32*q2*( s2-cor2)-q32*(q2*q2)*( s0-cor0)+q32*( s1-cor1)*q2*q1)+2.0*q3*q12*( q32*q2*( s0-cor0)*q1+q3*q32*q2*( s2-cor2)-q12*( s1-cor1)+q32*( s1-cor1)*(q2*q2))+-2.0*q0*q32*q2*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+2.0*q3*q32*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2))+2.0*( -2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q32*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))*q2+2.0*q0*q12*( q32*q2*( s2-cor2)*q1-q3*q32*q2*( s0-cor0)+q0*q32*( s1-cor1)*q2)+2.0*q32*(q2*q2)*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q12*q2*( q32*q2*( s0-cor0)*q1+q3*q32*q2*( s2-cor2)-q12*( s1-cor1)+q32*( s1-cor1)*(q2*q2)))*( t1+s1+-2.0*q0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+-2.0*q12*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)-m*y+-2.0*cor1);

            // df/dq3
            df[6] += 2.0*( 2.0*q3*q0*q32*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+2.0*q12*( q3*q32*( s0-cor0)*q1-q12*( s2-cor2)+q3*q32*( s1-cor1)*q2+(q3*q3)*q32*( s2-cor2))*q2+2.0*( q3*q32*( s2-cor2)*q1+q3*q0*q32*( s1-cor1)+q12*( s0-cor0)-(q3*q3)*q32*( s0-cor0))*q0*q12+2.0*q3*q32*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2))*( t1+s1+-2.0*q0*q12*( q12*( s2-cor2)*q1-q3*q12*( s0-cor0)+q0*q12*( s1-cor1))+-2.0*q12*q2*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)-m*y+-2.0*cor1)+2.0*( 2.0*q0*q12*( (q3*q3)*q32*( s1-cor1)-q3*q32*q2*( s2-cor2)+q3*q0*q32*( s0-cor0)-q12*( s1-cor1))+2.0*q12*( q3*q32*( s0-cor0)*q1-q12*( s2-cor2)+q3*q32*( s1-cor1)*q2+(q3*q3)*q32*( s2-cor2))*q1+2.0*q3*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1+2.0*q3*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q32)*( s0+-2.0*( q0*q12*( s0-cor0)-q12*q2*( s2-cor2)+q3*q12*( s1-cor1))*q0*q12+-2.0*cor0-x*m+t0+-2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)*q1)+2.0*( s2+-2.0*q3*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+2.0*q0*q12*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+t2+-2.0*cor2-z*m)*( -2.0*q12*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2)+-2.0*q3*q0*q32*( q12*( s1-cor1)*q1-q0*q12*( s2-cor2)-q12*q2*( s0-cor0))+2.0*q3*q12*( q3*q32*( s0-cor0)*q1-q12*( s2-cor2)+q3*q32*( s1-cor1)*q2+(q3*q3)*q32*( s2-cor2))+-2.0*q0*q12*( q3*q32*( s1-cor1)*q1-q3*q32*q2*( s0-cor0)-q3*q0*q32*( s2-cor2))+2.0*(q3*q3)*q32*( q12*( s0-cor0)*q1+q3*q12*( s2-cor2)+q12*( s1-cor1)*q2));
        }
    }
}

static void samplePointSpatially(const LidarPointCloud& source, int maxNumPointsToUse, double endFractionToIgnore)
{
    g_source.clear();

    if (maxNumPointsToUse <= 0)
    {
        g_source = source;
        return;
    }

    // First compute the total "distance" of the point cloud. I.e. the total distance when
    // going in order through all the points.
    int totalNumPoints = source.size();
    double totalDistance = 0.0;
    for (int i=1; i<totalNumPoints; ++i)
    {
        double d = sqrt(Distance2BetweenPoints(source[i-1].targetpos, source[i].targetpos));
        totalDistance += d;
    }

    double endDistanceToIgnore = endFractionToIgnore * totalDistance;
    double step = (totalDistance-2.0*endDistanceToIgnore) / maxNumPointsToUse;
    double distance = 0.0;
    int count = 1;
    for (int i=1; i<totalNumPoints; ++i)
    {
        double d = sqrt(Distance2BetweenPoints(source[i-1].targetpos, source[i].targetpos));
        distance += d;
        if (distance-endDistanceToIgnore >= count*step &&
            distance >= endDistanceToIgnore &&
            distance <= (totalDistance-endDistanceToIgnore))
        {
            g_source.push_back(source.at(i));
            ++count;
        }
    }
}

/**
   Perform ICP algorithm on source and target points, both of size
   n. The optimal translation that maps the source points into target
   points is calculated and placed in translation.
 */
void icpGsl(LidarPointCloud& source, double translation[3], double rotation[4], double centerOfRotation[3], int maxNumPointsToUse, double endFractionToIgnore, bool translationOnly, bool rotationOnly)
{
    /* For the ICP, we do the following:

       First compute an initial transformation between the 2 point sets
       by computing the centroid of each set and translating the source
       by that amount to the target

       Then iterate the following:
       1. For each point in source, find the closest point in target.
       2. Minimize the ssd between these corresponding points
    */

    g_source.clear();
    g_target.clear();

    g_translationOnly = translationOnly;
    g_rotationOnly = rotationOnly;

    samplePointSpatially(source, maxNumPointsToUse, endFractionToIgnore);

    centroid(g_source, g_sourceCentroid);

    // First 3 are translation, next 4 are rotation quaternion
    double transformation[7];
    transformation[0] = 0.0;
    transformation[1] = 0.0;
    transformation[2] = 0.0;
    transformation[3] = 1.0;
    transformation[4] = 0.0;
    transformation[5] = 0.0;
    transformation[6] = 0.0;

    if (rotationOnly == false)
    {
        // The following finds all the target points using the identity transformation
        findAllClosestPoints(transformation);

        Point targetCentroid;
        centroid(g_target, targetCentroid.targetpos);

        // Now update the current translation to be the displacement vector between centroids of the
        // source and target points
        transformation[0] = targetCentroid.targetpos[0] - g_sourceCentroid[0];
        transformation[1] = targetCentroid.targetpos[1] - g_sourceCentroid[1];
        transformation[2] = targetCentroid.targetpos[2] - g_sourceCentroid[2];
    }

#ifndef NDEBUG
    printf("Beginning ICP\n");
    printf("Initial value of objective function with no translation: %f\n",
           func(transformation, NULL));
#endif

    /* Now iterate */

    int maxIter = 5000;
    int i;
    for (i=0;i<maxIter;++i)
    {
#ifndef NDEBUG
        printf("\nStarting iteration %d of ICP\n", i+1);
#endif

        findAllClosestPoints(transformation);

        double prevssd = func(transformation, NULL);

        // See comment above for why we need to normalize the translation vector
        g_MULT_FACTOR = Normalize(transformation);
        optimizeGsl(func, grad, transformation, 7, NULL);
        MultiplyScalar(transformation, g_MULT_FACTOR);
        g_MULT_FACTOR = 1.0;

        // Normalize the quaternion since the quaternion probably got denormalized
        // during the optimization.
        normalizeQuaternion(&transformation[3], &transformation[3]);

        double ssd = func(transformation, NULL);

        if (ssd >= prevssd)
            break;
    }

    translation[0]      = transformation[0];
    translation[1]      = transformation[1];
    translation[2]      = transformation[2];
    rotation[0]         = transformation[3];
    rotation[1]         = transformation[4];
    rotation[2]         = transformation[5];
    rotation[3]         = transformation[6];
    centerOfRotation[0] = g_sourceCentroid[0];
    centerOfRotation[1] = g_sourceCentroid[1];
    centerOfRotation[2] = g_sourceCentroid[2];

    // transform the source target points and spacecraft positions using the optimal translation and rotation
    int totalNumPoints = source.size();
    for (int i=0; i<totalNumPoints; ++i)
    {
        transformPoint(source[i].targetpos, centerOfRotation, translation, rotation, source[i].targetpos);
        transformPoint(source[i].scpos, centerOfRotation, translation, rotation, source[i].scpos);
    }

#ifndef NDEBUG
    printf("Final transformation: %f %f %f %f %f %f %f\n", translation[0], translation[1], translation[2],
            rotation[0], rotation[1], rotation[2], rotation[3]);
#endif
}
