#include "util/util.h"
#include "geo/private/math/spline_parametric_curve.h"

#include "Eigen/LU"
#include "Eigen/Dense"
#include "unsupported/Eigen/NonLinearOptimization"

namespace ns
{
namespace geo
{

/**
 * if points' num less than 4, then use a line interpolation to generate 4 ctrl
 * points
 * @param points origin ctrl points
 * @return 4 ctrl points
 */
Eigen::MatrixXd get_ctrl_points(const Eigen::MatrixXd& points)
{
    using namespace Eigen;

    Index n = points.cols();
    NS_ASSERT(n > 1);
    NS_ASSERT(n < 4);

    MatrixXd ctrls(points.rows(), 4);
    if (n == 2)
    {
        VectorXd v = (points.col(1) - points.col(0)) / 3;
        for (int i = 0; i < 4; ++i)
        {
            ctrls.col(i) = points.col(0) + i * v;
        }
    }
    else if (n == 3)
    {
        RowVectorXd sd =
            (points.rightCols(2) - points.leftCols(2)).colwise().squaredNorm();
        Index i = (sd(0) > sd(1)) ? 1 : 2;
        ctrls.leftCols(i) = points.leftCols(i);
        ctrls.col(i) = (points.col(i - 1) + points.col(i)) / 2;
        ctrls.rightCols(3 - i) = points.rightCols(3 - i);
    }
    return ctrls;
}

SplineParametricCurve::SplineParametricCurve(const Eigen::MatrixXd& points)
{
    NS_ASSERT(points.cols() >= 2) << "SplineParametricCurve: need at least 2 "
                                     "points to create a spline (got "
                                  << points.cols() << ")." << std::endl;

    const Eigen::Index min_points = 4;
    if (points.cols() < min_points)
    {
        init_with_matrix(get_ctrl_points(points));
    }
    else
    {
        init_with_matrix(points);
    }
}

const Eigen::Spline<double, 2, 3>& SplineParametricCurve::get_spline() const
{
    return m_spline;
}

void SplineParametricCurve::init_with_matrix(const Eigen::MatrixXd& ctrls)
{
    using namespace Eigen;

    int num_controls = static_cast<int>(ctrls.cols()); // n + 1 control points
    int num_knots = num_controls + 3 + 1;              // m+1, m = n + p + 1

    VectorXd knots(num_knots);
    knots.head(3).setZero();
    knots.tail(3).setConstant(1);

    // internal knots are uniformly spaced
    knots.segment(3, num_knots - 6) = VectorXd::LinSpaced(num_knots - 6, 0, 1);

    m_spline = Spline<double, 2, 3>(knots, ctrls);
}

Eigen::Vector2d SplineParametricCurve::operator()(double tau) const
{
    return m_spline(tau);
}

Eigen::MatrixXd SplineParametricCurve::operator()(
    const Eigen::VectorXd& tau) const
{
    Eigen::MatrixXd ret(2, tau.size());
    for (int i = 0; i < tau.size(); ++i)
    {
        ret.col(i) = m_spline(tau(i));
    }
    return ret;
}

Eigen::Vector2d SplineParametricCurve::tangent(double tau) const
{
    Eigen::Vector2d ret = m_spline.derivatives(tau, 1).col(1);
    return ret.normalized();
}

Eigen::Vector2d SplineParametricCurve::norm(double tau) const
{
    Eigen::Vector2d t = tangent(tau);
    return {-t(1), t(0)};
}

Eigen::MatrixXd SplineParametricCurve::norm(const Eigen::VectorXd& tau) const
{
    auto n = tau.size();
    Eigen::MatrixXd ret(2, n);
    for (int i = 0; i < n; ++i)
    {
        ret.col(i) = norm(tau(i));
    }

    return ret;
}

Eigen::MatrixXd SplineParametricCurve::curvature(
    const Eigen::VectorXd& tau) const
{
    auto n = tau.size();
    Eigen::VectorXd ret(n);
    Eigen::Matrix2d der;
    for (int i = 0; i < n; ++i)
    {
        der = m_spline.derivatives(tau(i), 2).rightCols(2);
        double t = der.col(0).norm();
        ret(i) = (t < 1e-10) ? 0 : der.determinant() / (t * t * t);
    }

    return ret;
}

/*
 * Get the closest τ value for given point
 *
 * Let F(s) to be the square distance from point (x(s), y(s)) to p, this is to
 * get a τ value so that F(τ) minimize F(s).
 *
 * The implementation uses Levenberg-Marquardt Algorithm (LMA). For details of
 * the algorithm, search it on the wikipedia.
 *
 * F(s) = |p(s) - p|^2
 *
 * The jacobi matrix in matrix form is
 * J(s) = 2(p(s) - p) ⋅ p'(s)
 *
 * , which is a 1 x 1 matrix
 *
 * LMA is robust for smooth one dimensional problem like this case. However, be
 * aware that it could fall into local minimal.
 */
double SplineParametricCurve::closest_tau(const Eigen::Vector2d& p,
                                          double t0) const
{
    using namespace Eigen;
    using CubicSpline = Spline<double, 2, 3>;
    struct Functor
    {
        Functor(const CubicSpline& spline, const Eigen::Vector2d& p)
            : m_spline(spline), m_point(p)
        {
        }

        int inputs() const
        {
            return 1;
        }

        int values() const
        {
            return 1;
        }

        int operator()(const VectorXd& x, VectorXd& fvec) const
        {
            fvec(0) = (m_spline(x(0)).matrix() - m_point).squaredNorm();
            return 0;
        }

        int df(const VectorXd& x, MatrixXd& fjac) const
        {
            double t = x(0);
            Matrix2d der = m_spline.derivatives(t, 1);
            fjac(0, 0) = 2 * (der.col(0) - m_point).dot(der.col(1));
            return 0;
        }

    private:
        const CubicSpline& m_spline;
        const Eigen::Vector2d& m_point;
    };

    Functor functor(m_spline, p);
    LevenbergMarquardt<Functor> lm(functor);
    lm.parameters.ftol = 1.e-4;
    lm.parameters.xtol = 1.e-4;

    VectorXd x(1);
    x(0) = t0;

    // for smooth cubic function, LMA should be able to find a solution
    (void)lm.minimize(x);

    return x(0);
}

std::tuple<bool, double, double> SplineParametricCurve::closest_tau_and_offset(
    const Eigen::Vector2d& pos, const double acceptable_dis) const
{
    bool found = false;
    double opt_tau, opt_offset = 0.;

    // try with multiple initial guesses
    for (int i = 0; i <= 5; ++i)
    {
        double t0 = i * 0.2;
        double tau = closest_tau(pos, t0);
        if (tau >= 0. && tau < 1.0)
        {
            // 0 means we disallow extrapolation
            // note that if |tau| is too large, the result become unreliable

            double offset = (pos - m_spline(tau).matrix()).dot(norm(tau));
            if (std::abs(offset) < acceptable_dis)
            {
                return std::make_tuple(true, tau, offset);
            }
            else if (!found)
            {
                opt_offset = offset;
                opt_tau = tau;
                found = true;
            }
            else if (std::abs(offset) < std::abs(opt_offset))
            {
                opt_offset = offset;
                opt_tau = tau;
            }
        }
    }
    return std::make_tuple(found, opt_tau, opt_offset);
}

} // namespace geo
} // namespace ns