#ifndef MySpline2_C
#define MySpline2_C

#include <TMath.h>

#include <iostream>
#include <vector>
#include <cassert>

namespace MySpline2
{
    class Poly
  { // This class allows for efficient evaluation of mostly-fixed coefficient
    // polynomials.  It calculates the derived coefficients for derivatives upon
    // creation of updating of the coefficients and stores them in an internal vector.
  public:
    Poly(unsigned int N_, const double * const c_) : N(N_)
    { // N_: number of coefficients.
      // c_: array of N_ coefficients.  c_[0] is the constant term, 
      //     c_[N_-1] is the high-order term.
      SetCoeffs(N,c_);
    };
    Poly() : N(0) {};

    // Evaluate the polynomial at position x using the current coefficients.
    double Eval(double x) const { return Eval(N,x,coeffs[0].data()); };
    // Evaluate the nth derivative of the polynomial at the position x.
    double Deriv(unsigned int n, double x) const { return Eval(N-n,x,coeffs[n].data()); };


    void SetCoeffs(const double * const c)
    { // Set new coefficients for the polynomial.
      // This SetCoeffs does not change the number of coefficients.
      // Assumes enough coefficients are passed for the current N.
      coeffs[0].assign(c,c+N);
      for(unsigned int n = 1; n < N; n++)
      {
        for(unsigned int j = 0; j < N-n;j++)
        {
          coeffs[n][j] = c[j+n];
          for(unsigned int k = j+n; k > j; k--)
          {
            coeffs[n][j] *= k;
          }
        }
      }
    };

    void SetCoeffs(unsigned int N_, const double * const c)
    { // Set new coefficients for the polynomial, and change the number.
      // This overload is for when you want to change the number
      // of coefficients.  It's a bit faster than making a whole new Poly
      // object because it re-uses the coeff vector.
      N = N_;
      coeffs.clear();
      coeffs.emplace_back(c,c+N);
      for(unsigned int n = 1; n < N; n++)
      {
        std::vector<double> cprime(N-n,0);
        for(unsigned int j = 0; j < N-n;j++)
        {
          cprime[j] = c[j+n];
          for(unsigned int k = j+n; k > j; k--)
          {
            cprime[j] *= k;
          }
        }
        coeffs.push_back(std::move(cprime));
      }
    };

    // Retrieve a const pointer to the current coefficients.
    double const * GetCoeffs() const { return coeffs[0].data(); };

    // Print all the current coefficients & derived ones.
    void PrintCoeffs() const
    {
      for(unsigned int n = 0; n < N; n++)
      {
        std::cout << n << ": ";
        for(double val : coeffs[n])
        {
          std::cout << val << ", ";
        }
        std::cout << std::endl;
      }
    };

  private:
    inline static double Eval(unsigned int Nc, const double x, const double * const c)
    { // Recursive Horner's scheme polynomial evaluation.
      if( Nc > 1 ) { return c[0] + x*Eval(Nc-1,x,c+1); }
      else if( Nc == 1 ) { return c[0]; }
      else { return 0; }
    };

    unsigned int N; // I could just use coeffs[0].size()...
    std::vector<std::vector<double> > coeffs; 

  }; // End of class Poly
    class Spline
  { // This class represents a piecewise-defined polynomial with a certain smoothness
    // or continuity enforced.  The number of knots (the junctions between piecewise
    // segments) and X position of the knots must be specified ahead of time, but the 
    // coefficients of the polynomials may be specified later.  The spline then allows
    // for convenient calculation of the spline Y values and derivatives.
    // Depending on the smoothness required, each non-leading segment has a number of 
    // dependent coefficients.  These are automatically calculated in the constructor,
    // in SetCoeffs, and in the operator() and Eval overloads that take a cnew parameter.
  public:
    // The user is responsible for ensuring that "knots" has N_knots values and that they
    // are sorted.  The user is also responsible for ensuring that "coeffs_" has enough 
    // values for the requested spline.  The number of coefficients can be obtained from the
    // static method GetNCoeffsTotal.
    Spline() : N_knots(0), N(0), smoothness(0), n_c_total(0), n_c_seg(0) {};
    Spline(unsigned int N_knots_, unsigned int N_, unsigned int smoothness_,
           const double * const knots, const double * const coeffs_ = nullptr) : 
      N_knots(N_knots_), N(N_), smoothness(smoothness_),
      n_c_total((N-(smoothness+1))*N_knots + N), n_c_seg(N-(smoothness+1)), 
      knots_x(knots,knots+N_knots),
      //coeffs(coeffs_,coeffs_+n_c_total),
      scratch(n_c_total)
    {
      assert(N > smoothness && "smoothness argument cannot exceed poly coefficient count.");
      // fill coeffs with zeros if coefficients not provided.
      if(!coeffs_) coeffs.assign(n_c_total,0); 
      else coeffs.assign(coeffs_,coeffs_+n_c_total); 
      Poly junk(N,coeffs.data()); // Temporary Poly object with coefficients.
      segments.assign(N_knots+1,junk);  // Fill segments array with copies of junk.
      SetCoeffs(coeffs.data()); // Now calculate the correct coeffs for all segments.
    };

    // Convenience overloads.
    double operator()(const double * const x, const double * const cnew) 
    { return Eval(x[0],cnew);};
    double operator()(const double * const x) const { return Eval(x[0]); };
    double Eval(const double * const x, const double * const cnew)  
    { return Eval(x[0],cnew);};
    
    // This Eval looks up the correct segment and returns the spline Y-value there.
    double Eval(double x) const { return segments[GetSegment(x)].Eval(x); };

    // This Eval takes a cnew argument to change the coefficients.  It only updates
    // the derived coefficients if necessary.
    double Eval(double x, const double * const cnew)
    {
      TryNewCoeffs(cnew);
      return Eval(x);
    };

    // The deriv methods do the same thing as the Eval, but for derivatives.
    double Deriv(unsigned int n, const double * const x, const double * const cnew) 
    { return Deriv(n,x[0],cnew);};
    double Deriv(unsigned int n, double x) const
    {
      return segments[GetSegment(x)].Deriv(n,x);
    };
    double Deriv(unsigned int n, double x, const double * const cnew)
    {
      TryNewCoeffs(cnew);
      return Deriv(n,x);
    }

    void SetCoeffs(const double * const newcoeffs)
    { // Calculate new dependent coefficients for non-leading segments.

      // First we just copy the first N coefficients for the leading segment.
      segments[0].SetCoeffs(newcoeffs);
      // This loop is over each of the other dependent segments.
      for(unsigned int i_seg = 1; i_seg < N_knots+1; i_seg++)
      { // x_knot is the x value at the knot on the left of this segment.
        double x_knot = knots_x[i_seg-1]; 
        
        // copy over the independent coefficients for this segment.  This depends
        // on the order and smoothness.  The copying is done into the high-order
        // terms of the scratch array.
        std::copy(newcoeffs+N+n_c_seg*(i_seg-1),
                  newcoeffs+N+n_c_seg*i_seg,
                  scratch.data()+smoothness+1);
        // Temporarily fill the low-order terms of the scratch array with zeros.
        std::fill(scratch.data(),scratch.data()+smoothness+1,0);

        Poly & poly = segments[i_seg]; // Get a reference to Poly for this segment.
        // Set the segment poly to use the partially-filled coeffs.
        poly.SetCoeffs(scratch.data());
        for(unsigned int s = smoothness+1; s --> 0;) // "s goes to zero"
        {
          // This mysterious calculation & loop enforces the smoothness by setting
          // the low-order terms of the segment (previously set to zero) to the values
          // that make it match the previous segment.
          scratch[s] = segments[i_seg-1].Deriv(s,x_knot) - poly.Deriv(s,x_knot);
          poly.SetCoeffs(scratch.data());
        }
      }
      // I don't actually use the coeffs vector for anything, but it's useful for
      // recalling the current coefficients without needing to extract them from
      // the Poly segments that use the real coefficients.
      coeffs.assign(newcoeffs,newcoeffs+n_c_seg);
    };

    unsigned int GetSegment(double x) const
    { // Return index of segment for given x value.
      return std::lower_bound(knots_x.begin(), knots_x.begin()+N_knots,x) - knots_x.begin();
    }

    // Return pointer to current coeffs.
    double const * GetCoeffs() const { return coeffs.data(); };

    // Print values of all coefficients for all segments including derivatives.
    void PrintCoeffs() const
    {
      for(unsigned int n_seg = 0; n_seg < segments.size(); n_seg++)
      {
        std::cout << "Segment " << n_seg << ":" << std::endl;
        segments[n_seg].PrintCoeffs();
      }
    };

    // Boring getters for various private values.
    // I would have made the values public & const, but that makes the class
    // difficult to copy/assign and stuff, so this is easier.
    unsigned int GetN_knots() const {return N_knots;};
    unsigned int GetN() const {return N;};
    unsigned int GetSmoothness() const {return smoothness;};
    unsigned int GetNCoeffsTotal() const {return n_c_total;};
    unsigned int GetNCoeffsPerSegment() const {return n_c_seg;};
    std::vector<double> GetKnots() const {return knots_x;};

    // Static overload to give the number of coefficients required to construct
    // a Spline with the given number of knots, order, and smoothness.
    static unsigned int GetNCoeffsTotal(unsigned int N_knots, 
                                        unsigned int N, 
                                        unsigned int smoothness)
    { return (N-(smoothness+1))*N_knots + N ;};

private:
    unsigned int N_knots;
    unsigned int N;
    unsigned int smoothness;

    unsigned int n_c_total;
    unsigned int n_c_seg;
    std::vector<double> knots_x;
    std::vector<Poly> segments;
    std::vector<double> coeffs;
    std::vector<double> scratch;



    void TryNewCoeffs(const double * const cnew)
    {
      const double * const cold = coeffs.data();
      for(unsigned int i = 0; i < n_c_total; i++)
      {
        if(cold[i] != cnew[i])
        {
          SetCoeffs(cnew);
          return;
        }
      }
    }

  }; // End of class Spline
 
  class CorrectionVerifier
  {
  public:
    CorrectionVerifier(){slope = 0.0;};
    CorrectionVerifier(Spline & corr)
    {
      slope = corr.Eval(minimum_time)/minimum_time;
    };

    double Verify(double arr_time, double proposed) const
    {
      if(TMath::Abs(arr_time) > minimum_time)
      {
        return proposed;
      }
      else // This also catches NaN proposals.
      {
        return slope*arr_time;
      }
    };

    // The initialization of this static const data member cannot be in the class
    // definition for some reason, so look for it below this class.
    static const double minimum_time;
  private:
    double slope;
  };

  // For some reason the initialization of this static const data member
  // cannot be in the class definition, so I wrote it here.
  const double CorrectionVerifier::minimum_time = 0.01; // microseconds
  
} // End of namespace MySpline2

#if defined(__MAKECINT__)
#pragma link C++ defined_in MySpline2;
#pragma link C++ nestedclass;
//#pragma link C++ class MySpline2::MySpline2;

#pragma link C++ class MySpline2::Poly;
#pragma link C++ class MySpline2::Spline;

#pragma link C++ class MySpline2::CorrectionVerifier;
#endif

#endif // include guard