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

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

namespace ns
{
namespace geo
{

SplineFunction::SplineFunction(const Eigen::VectorXd& x,
                               const Eigen::VectorXd& y, int degree)
{
    init(x, y, degree);
}

void SplineFunction::init(Eigen::VectorXd const& x, Eigen::VectorXd const& y,
                          int degree)
{
    x_min = x.minCoeff();
    auto dx = x.maxCoeff() - x_min;
    m_alpha = 1.0 / dx;

    NS_ASSERT(!std::isnan(m_alpha))
        << "Cannot create a spline function: range of x values is too small";

    m_spline = Eigen::SplineFitting<Eigen::Spline<double, 1>>::Interpolate(
        y.transpose(),
        degree, // caller guarantee smaller than 3
        scaled_values(x));

    m_der_alpha.resize(degree + 1);
    m_der_alpha(0) = 1.0;
    for (int i = 0; i < degree; ++i)
    {
        m_der_alpha(i + 1) = m_der_alpha(i) * m_alpha;
    }
}

double SplineFunction::operator()(double x) const
{
    // x values need to be scaled down in extraction as well.
    return m_spline(scaled_value(x))(0);
}

Eigen::VectorXd SplineFunction::operator()(const Eigen::VectorXd& x) const
{
    return x.unaryExpr([&](double t) { return m_spline(scaled_value(t))(0); });
}

Eigen::VectorXd SplineFunction::der(double x, int degree) const
{
    return m_spline.derivatives(scaled_value(x), degree).transpose() *
           m_der_alpha.head(degree + 1);
}

Eigen::MatrixXd SplineFunction::der(const Eigen::VectorXd& x, int degree) const
{
    auto n = x.size();
    Eigen::MatrixXd ret(degree + 1, n);
    for (int i = 0; i < n; ++i)
    {
        ret.col(i) = der(x(i), degree);
    }
    return ret;
}

double SplineFunction::scaled_value(double x) const
{
    return m_alpha * (x - x_min);
}

Eigen::RowVectorXd SplineFunction::scaled_values(Eigen::VectorXd const& x) const
{
    return x.unaryExpr([this](double lx) { return scaled_value(lx); })
        .transpose();
}

} // namespace geo
} // namespace ns