#include "Solve.hh"
#include <math.h>
#include <assert.h>

Solve::Solve(void (*vf)(const Vector& , Vector& ),
	     void (*jac)(const Vector& , Matrix& ),
	     size_t dimension,
	     int maxsteps, SolverType s)
  :vfield(vf), jacobian(jac),
   d(dimension), MaxNumberOfSteps(maxsteps),
   NumberOfStepsTaken(0), step(0.0), 
   initflag(0), solver(s), //soln[0] not set; solver is s
   solndataspace(d*(maxsteps+1))
{
  assert(dimension > 0);
  assert(maxsteps > 0);
  soln = new Vector[MaxNumberOfSteps+1];
  for(size_t i=0; i < MaxNumberOfSteps+1; i++)
    (soln[i]).shadow(solndataspace, i*d, (i+1)*d-1);
}


void Solve::setStart(const Vector& v, double h){
  assert(v.getSize() == d);
  initflag = 1;
  soln[0] = v;
  step = h;
  NumberOfStepsTaken = 0; //if reinitialization, old computation is lost
}

void Solve::getSoln(int m, Vector& v) const{
  assert((0<=m)&&(m<=(int)NumberOfStepsTaken));
  v = soln[m];
}

void Solve::getSoln(int m, int n, Vector& v) const{
  assert(m<=n);
  assert((0<=m)&&(m<=(int)NumberOfStepsTaken));
  assert((0<=n)&&(n<=(int)NumberOfStepsTaken));
  for(int i=0; i < (n-m)+1; i++)
    v.getShadow(i*d,(i+1)*d-1) = soln[m+i];
}

void Solve::getSolnAtT(double t, Vector& v){
  int n = (int)floor(t/step);
  assert((0<=n)&&(n<=(int)NumberOfStepsTaken));
  switch(solver){
  case RK4: 
    rk4step(soln[n], (t-step*n), v); break;
  case RK7:
    rk7step(soln[n], (t-step*n), v); break;
  }
}


void Solve::advance(int n){
  assert(initflag==1);
  size_t base = NumberOfStepsTaken;
  assert(base+n <= MaxNumberOfSteps);
  int i;
  switch(solver){
  case RK4:
    for(i=0; i < n; i++)
      rk4step(soln[base+i], step, soln[base+i+1]);
    break;
  case RK7:
    for(i=0; i < n; i++)
      rk7step(soln[base+i], step, soln[base+i+1]);
    break;
  }
  NumberOfStepsTaken += n;
}


Solve::~Solve(){
  delete[] soln;
}

void Solve::rk4step(const Vector& y0, double h, Vector& yout){
  static Vector y1(d);
  static Vector y2(d);
  static Vector y3(d);
  static Vector f0(d);
  static Vector f1(d);
  static Vector f2(d);
  static Vector f3(d);
  
  static Vector temp1(d);

  vfield(y0,f0);
  y1 = f0;
  y1.scale(h/2.0);
  y1.add(y0);

  vfield(y1,f1);
  y2 = f1;
  y2.scale(h/2.0);
  y2.add(y0);

  vfield(y2,f2);
  y3 = f2;
  y3.scale(h);
  y3.add(y0);
  
  vfield(y3,f3);

  yout = y0;
  temp1 = f0; temp1.scale(h/6.0); yout.add(temp1);
  temp1 = f1; temp1.scale(h/3.0); yout.add(temp1);
  temp1 = f2; temp1.scale(h/3.0); yout.add(temp1);
  temp1 = f3; temp1.scale(h/6.0); yout.add(temp1);
}
  

void Solve::rk7step(const Vector& y0, double h, Vector& yout){
  static Vector y1(d);
  static Vector y2(d);
  static Vector y3(d);
  static Vector y4(d);
  static Vector y5(d);
  static Vector y6(d);
  static Vector y7(d);
  static Vector y8(d);
  static Vector y9(d);
  static Vector y10(d);
  static Vector y11(d);
  static Vector y12(d);
  static Vector f0(d);
  static Vector f1(d);
  static Vector f2(d);
  static Vector f3(d);
  static Vector f4(d);
  static Vector f5(d);
  static Vector f6(d);
  static Vector f7(d);
  static Vector f8(d);
  static Vector f9(d);
  static Vector f10(d);
  static Vector f11(d);
  static Vector f12(d);

  static Vector temp1(d);

  vfield(y0,f0);

  y1 = y0;
  temp1 = f0; temp1.scale(2*h/27); y1.add(temp1);

  vfield(y1,f1);

  y2 = y0;
  temp1 = f0; temp1.scale(h/36); y2.add(temp1);
  temp1 = f1; temp1.scale(h/12); y2.add(temp1);

  vfield(y2,f2);

  y3 = y0;
  temp1 = f0; temp1.scale(h/24); y3.add(temp1);
  temp1 = f2; temp1.scale(h/8);  y3.add(temp1);

  vfield(y3,f3);

  y4 = y0;
  temp1 = f0; temp1.scale(5*h/12); y4.add(temp1);
  temp1 = f2; temp1.scale(-25*h/16); y4.add(temp1);
  temp1 = f3; temp1.scale(25*h/16); y4.add(temp1);

  vfield(y4,f4);

  y5 = y0;
  temp1 = f0; temp1.scale(h/20); y5.add(temp1);
  temp1 = f3; temp1.scale(h/4); y5.add(temp1);
  temp1 = f4; temp1.scale(h/5); y5.add(temp1);
  
  vfield(y5,f5);

  y6 = y0;
  temp1 = f0; temp1.scale(-25*h/108); y6.add(temp1);
  temp1 = f3; temp1.scale(125*h/108); y6.add(temp1);
  temp1 = f4; temp1.scale(-65*h/27); y6.add(temp1);
  temp1 = f5; temp1.scale(125*h/54); y6.add(temp1);

  vfield(y6,f6);

  y7 = y0;
  temp1 = f0; temp1.scale(31*h/300); y7.add(temp1);
  temp1 = f4; temp1.scale(61*h/225); y7.add(temp1);
  temp1 = f5; temp1.scale(-2*h/9); y7.add(temp1);
  temp1 = f6; temp1.scale(13*h/900); y7.add(temp1);

  vfield(y7,f7);

  y8 = y0;
  temp1 = f0; temp1.scale(2*h); y8.add(temp1);
  temp1 = f3; temp1.scale(-53*h/6); y8.add(temp1);
  temp1 = f4; temp1.scale(704*h/45); y8.add(temp1);
  temp1 = f5; temp1.scale(-107*h/9); y8.add(temp1);
  temp1 = f6; temp1.scale(67*h/90); y8.add(temp1);
  temp1 = f7; temp1.scale(3*h); y8.add(temp1);

  vfield(y8,f8);

  y9=y0;
  temp1 = f0; temp1.scale(-91*h/108); y9.add(temp1);
  temp1 = f3; temp1.scale(23*h/108); y9.add(temp1);
  temp1 = f4; temp1.scale(-976*h/135); y9.add(temp1);
  temp1 = f5; temp1.scale(311*h/54); y9.add(temp1);
  temp1 = f6; temp1.scale(-19*h/60); y9.add(temp1);
  temp1 = f7; temp1.scale(17*h/6); y9.add(temp1);
  temp1 = f8; temp1.scale(-h/12); y9.add(temp1);

  vfield(y9,f9);
  
  y10=y0;
  temp1 = f0; temp1.scale(2383*h/4100); y10.add(temp1);
  temp1 = f3; temp1.scale(-341*h/164); y10.add(temp1);
  temp1 = f4; temp1.scale(4496*h/1025); y10.add(temp1);
  temp1 = f5; temp1.scale(-301*h/82); y10.add(temp1);
  temp1 = f6; temp1.scale(2133*h/4100); y10.add(temp1);
  temp1 = f7; temp1.scale(45*h/82); y10.add(temp1);
  temp1 = f8; temp1.scale(45*h/164); y10.add(temp1);
  temp1 = f9; temp1.scale(18*h/41); y10.add(temp1);

  vfield(y10,f10);
  
  //seventh order formula, mercifully, does not use stages 11 and 12

  yout = y0;
  temp1 = f0; temp1.scale(41*h/840); yout.add(temp1);
  temp1 = f5; temp1.scale(34*h/105); yout.add(temp1);
  temp1 = f6; temp1.scale(9*h/35); yout.add(temp1);
  temp1 = f7; temp1.scale(9*h/35); yout.add(temp1);
  temp1 = f8; temp1.scale(9*h/280); yout.add(temp1);
  temp1 = f9; temp1.scale(9*h/280); yout.add(temp1);
  temp1 = f10; temp1.scale(41*h/840); yout.add(temp1);
  
}




/**********Extensions to compute Sensitivity Matrix start here**********/


//can try direct call to BLAS here
void Solve::ex_vfield(const Vector& x, Vector& y){
  static Matrix X, Y;
  static Matrix jac(d,d);
  static Vector vx, vy;

  assert(x.getSize()==d*d+d);
  assert(y.getSize()==d*d+d);
  vx.shadow(x,0,d-1);
  vy.shadow(y,0,d-1);
  X.shadow(x, d, d+d*d-1, d, d);
  Y.shadow(y, d, d+d*d-1, d, d);

  vfield(vx,vy);
  jacobian(vx, jac);
  mult(jac, X, Y);
} 


void Solve::ex_rk4step(const Vector& y0, double h, Vector& yout){
  
  static size_t dd = d*d + d;
  static Vector y1(dd);
  static Vector y2(dd);
  static Vector y3(dd);
  static Vector f0(dd);
  static Vector f1(dd);
  static Vector f2(dd);
  static Vector f3(dd);
  
  static Vector temp1(dd);

  ex_vfield(y0,f0);
  y1 = f0;
  y1.scale(h/2.0);
  y1.add(y0);

  ex_vfield(y1,f1);
  y2 = f1;
  y2.scale(h/2.0);
  y2.add(y0);

  ex_vfield(y2,f2);
  y3 = f2;
  y3.scale(h);
  y3.add(y0);
  
  ex_vfield(y3,f3);

  yout = y0;
  temp1 = f0; temp1.scale(h/6.0); yout.add(temp1);
  temp1 = f1; temp1.scale(h/3.0); yout.add(temp1);
  temp1 = f2; temp1.scale(h/3.0); yout.add(temp1);
  temp1 = f3; temp1.scale(h/6.0); yout.add(temp1);
}

inline void Solve::ex_rk7step(const Vector& y0, double h, Vector& yout){
  static size_t dd = d*d+d;
  static Vector y1(dd);
  static Vector y2(dd);
  static Vector y3(dd);
  static Vector y4(dd);
  static Vector y5(dd);
  static Vector y6(dd);
  static Vector y7(dd);
  static Vector y8(dd);
  static Vector y9(dd);
  static Vector y10(dd);
  static Vector y11(dd);
  static Vector y12(dd);
  static Vector f0(dd);
  static Vector f1(dd);
  static Vector f2(dd);
  static Vector f3(dd);
  static Vector f4(dd);
  static Vector f5(dd);
  static Vector f6(dd);
  static Vector f7(dd);
  static Vector f8(dd);
  static Vector f9(dd);
  static Vector f10(dd);
  static Vector f11(dd);
  static Vector f12(dd);

  static Vector temp1(dd);

  ex_vfield(y0,f0);

  y1 = y0;
  temp1 = f0; temp1.scale(2*h/27); y1.add(temp1);

  ex_vfield(y1,f1);

  y2 = y0;
  temp1 = f0; temp1.scale(h/36); y2.add(temp1);
  temp1 = f1; temp1.scale(h/12); y2.add(temp1);

  ex_vfield(y2,f2);

  y3 = y0;
  temp1 = f0; temp1.scale(h/24); y3.add(temp1);
  temp1 = f2; temp1.scale(h/8);  y3.add(temp1);

  ex_vfield(y3,f3);

  y4 = y0;
  temp1 = f0; temp1.scale(5*h/12); y4.add(temp1);
  temp1 = f2; temp1.scale(-25*h/16); y4.add(temp1);
  temp1 = f3; temp1.scale(25*h/16); y4.add(temp1);

  ex_vfield(y4,f4);

  y5 = y0;
  temp1 = f0; temp1.scale(h/20); y5.add(temp1);
  temp1 = f3; temp1.scale(h/4); y5.add(temp1);
  temp1 = f4; temp1.scale(h/5); y5.add(temp1);
  
  ex_vfield(y5,f5);

  y6 = y0;
  temp1 = f0; temp1.scale(-25*h/108); y6.add(temp1);
  temp1 = f3; temp1.scale(125*h/108); y6.add(temp1);
  temp1 = f4; temp1.scale(-65*h/27); y6.add(temp1);
  temp1 = f5; temp1.scale(125*h/54); y6.add(temp1);

  ex_vfield(y6,f6);

  y7 = y0;
  temp1 = f0; temp1.scale(31*h/300); y7.add(temp1);
  temp1 = f4; temp1.scale(61*h/225); y7.add(temp1);
  temp1 = f5; temp1.scale(-2*h/9); y7.add(temp1);
  temp1 = f6; temp1.scale(13*h/900); y7.add(temp1);

  ex_vfield(y7,f7);

  y8 = y0;
  temp1 = f0; temp1.scale(2*h); y8.add(temp1);
  temp1 = f3; temp1.scale(-53*h/6); y8.add(temp1);
  temp1 = f4; temp1.scale(704*h/45); y8.add(temp1);
  temp1 = f5; temp1.scale(-107*h/9); y8.add(temp1);
  temp1 = f6; temp1.scale(67*h/90); y8.add(temp1);
  temp1 = f7; temp1.scale(3*h); y8.add(temp1);

  ex_vfield(y8,f8);

  y9=y0;
  temp1 = f0; temp1.scale(-91*h/108); y9.add(temp1);
  temp1 = f3; temp1.scale(23*h/108); y9.add(temp1);
  temp1 = f4; temp1.scale(-976*h/135); y9.add(temp1);
  temp1 = f5; temp1.scale(311*h/54); y9.add(temp1);
  temp1 = f6; temp1.scale(-19*h/60); y9.add(temp1);
  temp1 = f7; temp1.scale(17*h/6); y9.add(temp1);
  temp1 = f8; temp1.scale(-h/12); y9.add(temp1);

  ex_vfield(y9,f9);
  
  y10=y0;
  temp1 = f0; temp1.scale(2383*h/4100); y10.add(temp1);
  temp1 = f3; temp1.scale(-341*h/164); y10.add(temp1);
  temp1 = f4; temp1.scale(4496*h/1025); y10.add(temp1);
  temp1 = f5; temp1.scale(-301*h/82); y10.add(temp1);
  temp1 = f6; temp1.scale(2133*h/4100); y10.add(temp1);
  temp1 = f7; temp1.scale(45*h/82); y10.add(temp1);
  temp1 = f8; temp1.scale(45*h/164); y10.add(temp1);
  temp1 = f9; temp1.scale(18*h/41); y10.add(temp1);

  ex_vfield(y10,f10);
  
  //seventh order formula, mercifully, does not use stages 11 and 12

  yout = y0;
  temp1 = f0; temp1.scale(41*h/840); yout.add(temp1);
  temp1 = f5; temp1.scale(34*h/105); yout.add(temp1);
  temp1 = f6; temp1.scale(9*h/35); yout.add(temp1);
  temp1 = f7; temp1.scale(9*h/35); yout.add(temp1);
  temp1 = f8; temp1.scale(9*h/280); yout.add(temp1);
  temp1 = f9; temp1.scale(9*h/280); yout.add(temp1);
  temp1 = f10; temp1.scale(41*h/840); yout.add(temp1);
  
}


void Solve::getSMatrix(double t, Matrix& S){
  int n = (int)floor(t/step);
  assert(n<=(int)NumberOfStepsTaken);
  static Vector yin(d*d+d);
  static Vector yout(d*d+d);
  static Matrix Min, Mout;
  Min.shadow(yin, d, d*d+d-1, d, d);
  Mout.shadow(yout, d, d*d+d-1, d, d);
  yin.getShadow(0,d-1) = soln[0];
  size_t i, j;
  for(i=0; i < d; i++)
    for(j=0; j < d; j++)
      Min(i,j) = (i==j)?1.0:0.0;
  switch(solver){
  case RK4:
    for(i=0; i < n; i++){
      ex_rk4step(yin, step, yout);
      yin = yout;
    }
    ex_rk4step(yin, (t-n*step), yout);
    break;
  case RK7:
    for(i=0; i < n; i++){
      ex_rk7step(yin, step, yout);
      yin = yout;
    }
    ex_rk7step(yin, (t-n*step), yout);
    break;
  }
  S = Mout;//Mout shadows yout
}