#include "trilateration.h"

#include <cmath>
#include <iostream>

#include <Eigen/Eigen>

#include <unsupported/Eigen/NonLinearOptimization>
#include <unsupported/Eigen/NumericalDiff>

namespace Trilateration
{
    // see: https://github.com/Wayne82/Trilateration/blob/master/source/Trilateration.cpp

    Point2 peusodInverse2d(const std::vector<Point2>& positions, const std::vector<float>& distances)
    {
        // To locate position on a 2d plan, have to get at least 3 becaons,
        // otherwise return false.
        if (positions.size() < 3)
            assert(false);
        if (positions.size() != distances.size())
            assert(false);

        // Define the matrix that we are going to use
        size_t count = positions.size();
        size_t rows = count * (count - 1) / 2;
        Eigen::MatrixXd m(rows, 2);
        Eigen::VectorXd b(rows);

        // Fill in matrix according to the equations
        size_t row = 0;
        double x1, x2, y1, y2, r1, r2;

        for (size_t i = 0; i < count; ++i) {
            for (size_t j = i + 1; j < count; ++j) {
                x1 = positions[i].x, y1 = positions[i].y;
                x2 = positions[j].x, y2 = positions[j].y;
                r1 = distances[i];
                r2 = distances[j];
                m(row, 0) = x1 - x2;
                m(row, 1) = y1 - y2;
                b(row) = ((pow(x1, 2) - pow(x2, 2)) +
                    (pow(y1, 2) - pow(y2, 2)) -
                    (pow(r1, 2) - pow(r2, 2))) / 2;
                row++;
            }
        }

        // Then calculate to solve the equations, using the least square solution
        //Eigen::Vector2d location = m.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(b);
        Eigen::Vector2d pseudoInv = (m.transpose()*m).inverse() * m.transpose() *b;

        return Point2(pseudoInv.x(), pseudoInv.y());
    }

    Point3 pseudoInverse3d(const std::vector<Point3>& positions, const std::vector<float>& distances)
    {
        // To locate position in a 3D space, have to get at least 4 becaons
        if (positions.size() < 4)
            assert(false);
        if (positions.size() != distances.size())
            assert(false);

        // Define the matrix that we are going to use
        size_t count = positions.size();
        size_t rows = count * (count - 1) / 2;
        Eigen::MatrixXd m(rows, 3);
        Eigen::VectorXd b(rows);

        // Fill in matrix according to the equations
        size_t row = 0;
        double x1, x2, y1, y2, z1, z2, r1, r2;

        for (size_t i = 0; i < count; ++i) {
            for (size_t j = i + 1; j < count; ++j) {
                x1 = positions[i].x, y1 = positions[i].y, z1 = positions[i].z;
                x2 = positions[j].x, y2 = positions[j].y, z2 = positions[j].z;
                r1 = distances[i];
                r2 = distances[j];
                m(row, 0) = x1 - x2;
                m(row, 1) = y1 - y2;
                m(row, 2) = z1 - z2;
                b(row) = ((pow(x1, 2) - pow(x2, 2)) +
                    (pow(y1, 2) - pow(y2, 2)) +
                    (pow(z1, 2) - pow(z2, 2)) -
                    (pow(r1, 2) - pow(r2, 2))) / 2;
                row++;
            }
        }

        // Then calculate to solve the equations, using the least square solution
        Eigen::Vector3d location = m.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(b);

        return Point3(location.x(), location.y(), location.z());
    }


    // Generic functor
    // See http://eigen.tuxfamily.org/index.php?title=Functors
    // C++ version of a function pointer that stores meta-data about the function
    template<typename _Scalar, int NX = Eigen::Dynamic, int NY = Eigen::Dynamic>
    struct Functor
    {

        // Information that tells the caller the numeric type (eg. double) and size (input / output dim)
        typedef _Scalar Scalar;
        enum { // Required by numerical differentiation module
            InputsAtCompileTime = NX,
            ValuesAtCompileTime = NY
        };

        // Tell the caller the matrix sizes associated with the input, output, and jacobian
        typedef Eigen::Matrix<Scalar, InputsAtCompileTime, 1> InputType;
        typedef Eigen::Matrix<Scalar, ValuesAtCompileTime, 1> ValueType;
        typedef Eigen::Matrix<Scalar, ValuesAtCompileTime, InputsAtCompileTime> JacobianType;

        // Local copy of the number of inputs
        int m_inputs, m_values;

        // Two constructors:
        Functor() : m_inputs(InputsAtCompileTime), m_values(ValuesAtCompileTime) {}
        Functor(int inputs, int values) : m_inputs(inputs), m_values(values) {}

        // Get methods for users to determine function input and output dimensions
        int inputs() const { return m_inputs; }
        int values() const { return m_values; }

    };

    struct DistanceFunction : Functor<double>
    {
    private:
        const std::vector<Point2>& positions;
        const std::vector<float>& distances;

    public:
        DistanceFunction(const std::vector<Point2>& positions, const std::vector<float>& distances)
            : Functor<double>(positions.size(), positions.size()), positions(positions), distances(distances)
        {}

        int operator()(const Eigen::VectorXd &x, Eigen::VectorXd &fvec) const 
        {
            const Point2 p(x(0), x(1));

            for (size_t i = 0; i < positions.size(); i++)
            {
                fvec(i) = p.getDistance(positions[i]) - distances[i];
            }

            return 0;
        }
    };

    struct DistanceFunctionDiff : public Eigen::NumericalDiff<DistanceFunction>
    {
        DistanceFunctionDiff(const DistanceFunction& functor)
            : Eigen::NumericalDiff<DistanceFunction>(functor, 1.0e-6)
        {}
    };

    std::string lmStatusToString(const Eigen::LevenbergMarquardtSpace::Status stat)
    {
        switch (stat)
        {
        // Non-Erros
        case Eigen::LevenbergMarquardtSpace::NotStarted:
            return "Not started. No Error.";
        case Eigen::LevenbergMarquardtSpace::Running:
            return "Running. No Error.";
        // Error
        case Eigen::LevenbergMarquardtSpace::ImproperInputParameters:
            return "Error. Invalid input parameters";
        case Eigen::LevenbergMarquardtSpace::UserAsked:
            return "Error in user-implemented evaluation or gradient computation.";
        // Warnings
        case Eigen::LevenbergMarquardtSpace::RelativeReductionTooSmall:
            return "WARN: The cosine of the angle between fvec and any column of the jacobian is at most gtol in absolute value.";
        case Eigen::LevenbergMarquardtSpace::RelativeErrorTooSmall:
            return "WARN: Relative error too small.";
        case Eigen::LevenbergMarquardtSpace::RelativeErrorAndReductionTooSmall:
            return "WARN: Relative error and reduction too small.";
        case Eigen::LevenbergMarquardtSpace::CosinusTooSmall:
            return "WARN: The cosine of the angle between fvec and any column of the jacobian is at most gtol in absolute value.";
        case Eigen::LevenbergMarquardtSpace::TooManyFunctionEvaluation:
            return "WARN: Too many function evaluations done.";
        case Eigen::LevenbergMarquardtSpace::FtolTooSmall:
            return "WARN: ftol is too small. No further reduction in the sum of squares is possible";
        case Eigen::LevenbergMarquardtSpace::XtolTooSmall:
            return "WARN: xtol is too small. No further improvement in the approximate solution x is possible.";
        case Eigen::LevenbergMarquardtSpace::GtolTooSmall:
            return "WARN: gtol is too small.fvec is orthogonal to the columns of the jacobian to machine precision.";
        default:
            return "Unkown status. Status=" + std::to_string(static_cast<int>(stat));
        }
    }

    Point2 levenbergMarquardt(const std::vector<Point2>& positions, const std::vector<float>& distances)
    {
        const Point2 pseudoInvApprox = peusodInverse2d(positions, distances);

        Eigen::Vector2d startVal;
        startVal << pseudoInvApprox.x, pseudoInvApprox.y;
        //startVal << 0, 0;

        DistanceFunction functor(positions, distances);
        DistanceFunctionDiff numDiff(functor);
        Eigen::LevenbergMarquardt<DistanceFunctionDiff, double> lm(numDiff);

        // Parameters
        // factor : Sets the step bound for the diagonal shift
        // maxfev : Sets the maximum number of function evaluation
        // ftol   : Sets the tolerance for the norm of the vector function
        // xtol   : Sets the tolerance for the norm of the solution vector
        // gtol   : Sets the tolerance for the norm of the gradient of the error vector
        // epsilon: Sets the error precision
        lm.parameters.maxfev = 200;
        lm.parameters.xtol = 1.0e-10;
        

        Eigen::VectorXd z = startVal;
        auto status = lm.minimize(z);
        std::cout << "Levenberg Marquardt" << std::endl;
        std::cout << "status: " << lmStatusToString(status) << std::endl;
        std::cout << "iter count: " << lm.iter << std::endl;        
        std::cout << "levenMarq: " << z.transpose() << std::endl;
        std::cout << "pseudoInv: " << startVal.transpose() << std::endl;

        const Point2 levenPos(z(0), z(1));

        // Compute error        
        double errPseudo = 0;
        double errLeven = 0;
        for (size_t i = 0; i < positions.size(); i++)
        {
            double d1 = pseudoInvApprox.getDistance(positions[i]) - distances[i];
            errPseudo += d1 * d1;

            double d2 = levenPos.getDistance(positions[i]) - distances[i];
            errLeven += d2 * d2;
        }

        //assert(errLeven <= errPseudo);

        std::cout << "err pseud: " << errPseudo << std::endl;
        std::cout << "err leven: " << errLeven << std::endl << std::endl;

        return levenPos;
    }
}