#include "./polyfit.h"
#define USE_EIGEN3
#ifdef USE_EIGEN3
#include <Eigen/Dense>
#endif

/* ----------------------------------------------------------------------
   Evaluate polynomials
*/

vector<R> Polyfit1::coeffs() const
{
  return {c0, c1};
}

Polyfit1 Polyfit2::deriv() const
{
  return Polyfit1(c1, c2*2.0);
}

vector<R> Polyfit2::coeffs() const
{
  return {c0, c1, c2};
}


Polyfit2 Polyfit3::deriv() const
{
  return Polyfit2(c1, c2*2.0, c3*3.0);
}

vector<R> Polyfit3::coeffs() const
{
  return {c0, c1, c2, c3};
}


Polyfit3 Polyfit4::deriv() const
{
  return Polyfit3(c1, c2*2.0, c3*3.0, c4*4.0);
}

vector<R> Polyfit4::coeffs() const
{
  return {c0, c1, c2, c3, c4};
}


Polyfit4 Polyfit5::deriv() const
{
  return Polyfit4(c1, c2*2.0, c3*3.0, c4*4.0, c5*5.0);
}

vector<R> Polyfit5::coeffs() const
{
  return {c0, c1, c2, c3, c4, c5};
}


ostream & operator <<(ostream &s, const Polyfit1 &a)
{ 
  return s <<
    "Polyfit1("s << a.c0 <<
    ", " << a.c1 <<
    ")";
}

ostream & operator <<(ostream &s, const Polyfit2 &a)
{
  return s <<
    "Polyfit2("s << a.c0 <<
    ", " << a.c1 <<
    ", " << a.c2 <<
    ")";
}

ostream & operator <<(ostream &s, const Polyfit3 &a)
{
  return s <<
    "Polyfit3("s << a.c0 <<
    ", " << a.c1 <<
    ", " << a.c2 <<
    ", " << a.c3 <<
    ")";
}

ostream & operator <<(ostream &s, const Polyfit4 &a)
{
  return s <<
    "Polyfit4("s << a.c0 <<
    ", " << a.c1 <<
    ", " << a.c2 <<
    ", " << a.c3 <<
    ", " << a.c4 <<
    ")";
}

ostream & operator <<(ostream &s, const Polyfit5 &a)
{
  return s <<
    "Polyfit5("s << a.c0 <<
    ", " << a.c1 <<
    ", " << a.c2 <<
    ", " << a.c3 <<
    ", " << a.c4 <<
    ", " << a.c5 <<
    ")";
}


/*
  Fit a polynomial to some X and Y data. That is, return a Polyfit{1,3,5} p so that getValue(p, X) approximates Y.
 */

#ifdef USE_EIGEN3
Polyfit1 mkPolyfit1(vector<R> const &xs, vector<R> const &ys)
{
  if (xs.size() != ys.size()) throw runtime_error("incompatible arrays"s);
  if (xs.size() < 2) throw runtime_error("not enough data"s);

  Eigen::MatrixXd xsm(xs.size(), 2);
  Eigen::MatrixXd ysm(ys.size(), 1);

  for (size_t ri = 0; ri < xs.size(); ri++) {
    double x = xs[ri];
    double y = ys[ri];
    xsm(ri, 0) = 1;
    xsm(ri, 1) = x;
    ysm(ri, 0) = y;
  }

  // Throws runtime_error if no solution
  Eigen::MatrixXd coeffs = xsm.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ysm);
  return Polyfit1(coeffs(0, 0), coeffs(1, 0));
}

Polyfit2 mkPolyfit2(vector<R> const &xs, vector<R> const &ys)
{
  if (xs.size() != ys.size()) throw runtime_error("incompatible arrays"s);
  if (xs.size() < 4) throw runtime_error("not enough data"s);

  Eigen::MatrixXd xsm(xs.size(), 3);
  Eigen::MatrixXd ysm(ys.size(), 1);

  for (size_t ri = 0; ri < xs.size(); ri++) {
    double x = xs[ri];
    double y = ys[ri];
    xsm(ri, 0) = 1;
    xsm(ri, 1) = x;
    xsm(ri, 2) = x * x;
    ysm(ri, 0) = y;
  }

  // Throws runtime_error if no solution
  Eigen::MatrixXd coeffs = xsm.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ysm);
  return Polyfit2(coeffs(0, 0), coeffs(1, 0), coeffs(2, 0));
}

Polyfit3 mkPolyfit3(vector<R> const &xs, vector<R> const &ys)
{
  if (xs.size() != ys.size()) throw runtime_error("incompatible arrays"s);
  if (xs.size() < 4) throw runtime_error("not enough data"s);

  Eigen::MatrixXd xsm(xs.size(), 4);
  Eigen::MatrixXd ysm(ys.size(), 1);

  for (size_t ri = 0; ri < xs.size(); ri++) {
    double x = xs[ri];
    double y = ys[ri];
    xsm(ri, 0) = 1;
    xsm(ri, 1) = x;
    xsm(ri, 2) = x * x;
    xsm(ri, 3) = x * x * x;
    ysm(ri, 0) = y;
  }

  // Throws runtime_error if no solution
  Eigen::MatrixXd coeffs = xsm.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ysm);
  return Polyfit3(coeffs(0, 0), coeffs(1, 0), coeffs(2, 0), coeffs(3, 0));
}

Polyfit4 mkPolyfit4(vector<R> const &xs, vector<R> const &ys)
{
  if (xs.size() != ys.size()) throw runtime_error("incompatible arrays"s);
  if (xs.size() < 6) throw runtime_error("not enough data"s);

  Eigen::MatrixXd xsm(xs.size(), 5);
  Eigen::MatrixXd ysm(ys.size(), 1);

  for (size_t ri = 0; ri < xs.size(); ri++) {
    double x = xs[ri];
    double y = ys[ri];
    xsm(ri, 0) = 1;
    xsm(ri, 1) = x;
    xsm(ri, 2) = x * x;
    xsm(ri, 3) = x * x * x;
    xsm(ri, 4) = x * x * x * x;
    ysm(ri, 0) = y;
  }

  // Throws runtime_error if no solution
  Eigen::MatrixXd coeffs = xsm.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ysm);
  return Polyfit4(coeffs(0, 0), coeffs(1, 0), coeffs(2, 0), coeffs(3, 0), coeffs(4, 0));
}

Polyfit5 mkPolyfit5(vector<R> const &xs, vector<R> const &ys)
{
  if (xs.size() != ys.size()) throw runtime_error("incompatible arrays"s);
  if (xs.size() < 6) throw runtime_error("not enough data"s);

  Eigen::MatrixXd xsm(xs.size(), 6);
  Eigen::MatrixXd ysm(ys.size(), 1);

  for (size_t ri = 0; ri < xs.size(); ri++) {
    double x = xs[ri];
    double y = ys[ri];
    xsm(ri, 0) = 1;
    xsm(ri, 1) = x;
    xsm(ri, 2) = x * x;
    xsm(ri, 3) = x * x * x;
    xsm(ri, 4) = x * x * x * x;
    xsm(ri, 5) = x * x * x * x * x;
    ysm(ri, 0) = y;
  }

  // Throws runtime_error if no solution
  Eigen::MatrixXd coeffs = xsm.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(ysm);
  return Polyfit5(coeffs(0, 0), coeffs(1, 0), coeffs(2, 0), coeffs(3, 0), coeffs(4, 0), coeffs(5, 0));
}
#endif