#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 calculateLocation2d(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 calculateLocation3d(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)
        {}
    };

    Point2 levenbergMarquardt(const std::vector<Point2>& positions, const std::vector<float>& distances)
    {
        Point2 pseudoInvApprox = calculateLocation2d(positions, distances);
        Eigen::Vector2d initVal;
        initVal << pseudoInvApprox.x, pseudoInvApprox.y;

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

        DistanceFunction functor(positions, distances);
        DistanceFunctionDiff numDiff(functor);
        Eigen::LevenbergMarquardt<DistanceFunctionDiff, double> lm(numDiff);
        lm.parameters.maxfev = 2000;
        lm.parameters.xtol = 1.0e-10;
        std::cout << lm.parameters.maxfev << std::endl;

        Eigen::VectorXd z = startVal;
        int ret = lm.minimize(z);
        std::cout << "iter count: " << lm.iter << std::endl;
        std::cout << "return status: " << ret << std::endl;
        std::cout << "zSolver: " << z.transpose() << std::endl;
        std::cout << "pseudoInv: " << initVal.transpose() << std::endl;


        Point2 bla(z(0), z(1));

        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 = bla.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 Point2(z(0), z(1));

    }
}