#include <TRandom3.h>
#include <TGraph.h>
#include <TF1.h>

#include <iostream>
#include <cmath>
#include <vector>

namespace MySpline
{
  class MySpline
  {
  public:
    
    // The following function template and specialization implement a
    // recursive template for a Horner's-scheme-evaluated polynomial.
    // As long as the order of the polynomial is static, this generated
    // is the most efficient machine code possible.
    // NOTE: The N parameter here is the NUMBER OF TERMS in the
    // polynomial, it is NOT the order.  A quadratic has N = 3.
    
    // TODO: Add compile-time checks that N >= 1
    template <unsigned int N> 
    static double Poly(const double * const x, const double * const c)
    {
      return c[0] + x[0]*Poly<N-1>(x,c+1);
    }
    
    // This is a template for the nth derivative of an N-Polynomial
    // as evaluated by the other functions above.  Remember that N is the
    // number of coefficients in the polynomial, NOT the order.  
    // TODO: Add compile-time assertions that PolyDeriv<N,n> is 
    // never called with N < n.
    // http://stackoverflow.com/a/6765840
    template <unsigned int N, unsigned int n>
    static double PolyDeriv(const double * const x, const double * const c)
    { // nth derivative of Nth order polynomial.
      // f_N^{(n)} = \sum_{j = 0}^{N-n} (j+n)!/j! c_{j+n} x^j
      // The result is just an N-mth order polynomial with transformed 
      // coefficients.  The jth order coefficient is (j+n)!/j!*c_{j+n}.
      // TODO: Find a way to evaluate the coefficient multipliers (the
      // factorials) at compile-time.
      double cprime[N-n];
      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;
            }
        }
      return Poly<N-n>(x,cprime);
    }
  };
  
  // These specializations ensure that the recursion terminates and
  // that PolyDeriv<N,N> gives the right answer.
  template <>
  inline double MySpline::Poly<1>(const double * const, const double * const c)
  { return c[0]; }
  template <>
  inline double MySpline::Poly<0>(const double * const, const double * const)
  { return 0; }

  // This class represents a piecewise continuous or smooth polynomial
  // with a fixed number & position of junctions (knots).  The number
  // of coefficients of the segments N is also fixed and uniform.  The
  // knot x values are fixed, but not the y values.
  // When a PieceWise object is instantiated, you can ask its
  // n_c_total member how many parameters it needs.  The first N
  // parameters are the coefficients of the first segment.  Subsequent
  // parameters are the high-order coefficients of the other segments.
  // Continuity (smoothness order of 0) is enforced.
  template <unsigned int N_knots, unsigned int N>
  class PieceWise
  {
  public:
    PieceWise(const double * const knots) : knots_x(knots, knots+N_knots){};
    
    double operator()(const double * const x, const double * const c)
    { // Number of coefficients in c is n_c_total.
      
      // The first N are used for the first segment, right here:
      if(x[0] <= knots_x[0]) return MySpline::Poly<N>(x,c);
      
      // The following blocks of N-(s+1) elements of c are for each
      // following segment.  They are doing to be the HIGH-ORDER
      // elements of the segments, because the first s coefficients of the
      // later segments are determined by previous segments.
      
      // Because of the stupid ROOT * signatures we need to make a copy
      // of the coefficients to modify.
      double cprime[N] = {0};
      for(unsigned int i = 1; i < N_knots; i++)
        {
          if(x[0] <= knots_x[i])
            {
              // the +1 below in cprime+1 depends on the smoothness.
              std::copy(c+N+n_c_seg*(i-1), c+N+n_c_seg*i, cprime+1);
              cprime[0] = 0.0;
              // The fact that we ask for Poly here is because
              // smoothness = 0.  Nonzero smoothness would require PolyDeriv.
              cprime[0] = this->operator()(&(knots_x[i-1]),c) -
                MySpline::Poly<N>(&(knots_x[i-1]),cprime);
              return MySpline::Poly<N>(x,cprime);
            }
        }
      std::copy(c+n_c_total-(n_c_seg), c+n_c_total,cprime+1);
      cprime[0] = 0;
      cprime[0] = this->operator()(&(knots_x[N_knots-1]),c) -
        MySpline::Poly<N>(&(knots_x[N_knots-1]),cprime);
      return MySpline::Poly<N>(x,cprime);
    };
    
    static const unsigned int smoothness = 0;
    const unsigned int n_c_total = (N-(smoothness+1))*N_knots + N;
    const unsigned int n_c_seg = N-(smoothness+1); // Number of coeffs
                                                   // per additional segment.
    const std::vector<double> knots_x;
  };
 
  // This class specialization is in case the PieceWise function is
  // defined with zero knots.  Then it's just a simple polynomial.
  template <unsigned int N>
  class PieceWise<0,N>
  {
  public:
    PieceWise(const double * const = nullptr) {};
    
    double operator()(const double * const x, const double * const c)
    {
      return MySpline::Poly<N>(x,c);
    };
    static const unsigned int n_c_total = N;
    static const unsigned int n_c_seg = 0;
  };
  
  void test1()
  {
    const int N = 3;
    double x[] = {5};
    double c[N] = {1,2,3};
    std::cout << MySpline::Poly<N>(x,c) << std::endl;
    std::cout << MySpline::PolyDeriv<N,1>(x,c) << std::endl;
    std::cout << MySpline::PolyDeriv<N,2>(x,c) << std::endl;
    std::cout << MySpline::PolyDeriv<N,3>(x,c) << std::endl;
    //const double x[] = {0,1,2,3,4,5,6,7,8,9};
    //const double y[] = {3,6,2,7,8,4,2,6,8,4};
    
    const double k[3] = {2,5,7};
    double c2[] = {1,2,3, 4,5, 6,7, 8,9};
    PieceWise<3,3> pw(k);
    x[0] = 1.5;
    std::cout << pw(x,c2) << std::endl;
    x[0] = 2.5;
    std::cout << pw(x,c2) << std::endl;
    x[0] = 6.95;
    std::cout << pw(x,c2) << std::endl;
    x[0] = 7.0;
    std::cout << pw(x,c2) << std::endl;
    x[0] = 7.5;
    std::cout << pw(x,c2) << std::endl;
  }

  void test2()
  {
    const double k[2] = {1,2};
    PieceWise<2,3> *  pw = new PieceWise<2,3>(k);
    TF1 * fpw = new TF1("fpw",pw,0,3,pw->n_c_total);
    double c[7] = {0,1,0, //a + bx + cx^2 for first segment
                   -1,0,  //    bx + cx^2 for second segment. 
                   1,0};  //    bx + cx^2 for 3rd segment.
    std::cout << "TF1 needs " << pw->n_c_total << " parameters." << std::endl;
    fpw->SetParameters(c);
    fpw->SetNpx(1000);
    fpw->SetLineWidth(1);
    fpw->Draw();
  }
  
  void test3()
  {
    const double x[] = {0,1,2,3,4,5,6,7,8,9};
    const double y[] = {3,6,10,17,20,21,20,18,15,11};
    TGraph * g = new TGraph(10,x,y);
    
    const double k[1] = {4.5};
    PieceWise<1,3> * pw = new PieceWise<1,3>(k);
    double c[5] = {1,1,1,
                   1,1};
    TF1 * f_fit = new TF1("f_fit",pw,-1,11,pw->n_c_total);
    f_fit->SetParameters(c);
    f_fit->SetNpx(1000);
    f_fit->SetLineWidth(1);
    
    g->Fit(f_fit,"S");
    g->Draw("A*");
  }
  
  void test4(int seed = 1234)
  {
    const double k[5] = {1,2,3.5,4,7};
    PieceWise<5,6> * pw = new PieceWise<5,6>(k);
    TF1 * fpw = new TF1("fpw",pw,0,8,pw->n_c_total);
    double c[31] = {0,0,0,0,0,0,
                    0,0,0,0,0,
                    0,0,0,0,0,
                    0,0,0,0,0,
                    0,0,0,0,0,
                    0,0,0,0,0};
    TRandom3 tr(seed);
    unsigned int oldseg = -1;
    for(unsigned int i = 0; i < 31; i++)
      {
        unsigned int seg = i < 6 ? 0 : (i-6)/5 + 1;
        unsigned int order = i < 6 ? i : (i-6) % 5 + 1;
        c[i] = (tr.Rndm() - 0.5)/(i+1);
        if(oldseg != seg)
          {
            std::cout << "Segment " << seg << std::endl;
            oldseg = seg;
          }
        std::cout << "i: " << i << "(" << order << ")" << ": " << c[i] << std::endl;
      }
    
    std::cout << "TF1 needs " << pw->n_c_total << " parameters." << std::endl;
    fpw->SetParameters(c);
    fpw->SetNpx(1000);
    fpw->SetLineWidth(1);
    fpw->Draw();
  }
  
}// End of namespace MySpline