#include "PeriodicOrbit.hh"
#include "Vector.hh"
#include "Matrix.hh"
#include "SparseMatrix.hh"
#include <math.h>
#include "Clock.hh"

PeriodicOrbit::PeriodicOrbit(const Matrix& po, double t, 
			     vfTYPE vff, jTYPE jac):
  vfield(vff), jacobian(jac), pOrbit(po.getm(), po.getn()), T(t),
  dX(po.getm(), po.getn()), 
  A(po.getm()*po.getm(), po.getn()), 
  Res(po.getm(), po.getn()), 
  GradX(po.getm(), po.getn()),
  Y(po.getm()*po.getm(), po.getn()),
  f(po.getm(), po.getn()),
  g(po.getm(), po.getn())
{
  assert(T > 0);
  pOrbit = po;
  tinyFactor = 1e-6;
  spts = NULL;
  LPMatrix = NULL;
  YY = NULL;
  ff = NULL;
  gg = NULL;
  vf = NULL;
  updateCache();
  setSintvl();
}

PeriodicOrbit::~PeriodicOrbit(){
  if(spts != NULL)
    delete spts;
  if(LPMatrix != NULL)
    delete LPMatrix;
  if(YY != NULL)
    delete YY;
  if(ff != NULL)
    delete ff;
  if(gg != NULL)
    delete gg;
  if(vf != NULL)
    delete vf;
}

int PeriodicOrbit::getd(){
  return pOrbit.getm();
}

int PeriodicOrbit::getn(){
  return pOrbit.getn();
}

void PeriodicOrbit::addStep(const Matrix& step, double wstep){
  assert(step.getm()==pOrbit.getm());
  assert(step.getn()==pOrbit.getn());
  pOrbit.add(step);
  double w = 2*PI/T;
  w = w + wstep;
  T = 2*PI/w;
  updateCache();
}

double PeriodicOrbit::residualError(){
  return Res.fronorm();
}

double PeriodicOrbit::relativeResidualError(){
  double temp1 = Res.fronorm();
  double temp2 = dX.fronorm();
  double w = 2*PI/T;
  return temp1/(w*temp2);
}

double PeriodicOrbit::residualErrorAfterStep(const Matrix& step, double wstep){
  size_t d = pOrbit.getm();
  size_t n = pOrbit.getn();
  assert((step.getm()==d)&&(step.getn()==n));
  Matrix XX(d,n);
  Matrix RRes(d,n);
  double ww;
  XX = pOrbit; XX.add(step);
  double w = 2*PI/T;
  ww = w+wstep;
  ComputeResidual(XX,ww,RRes,vfield);
  return RRes.fronorm();
}

double PeriodicOrbit::predictedResidualError(const Matrix& step, double wstep){
  size_t d = pOrbit.getm();
  size_t n = pOrbit.getn();
  assert((step.getm()==d)&&(step.getn()==n));
  Matrix y(d,n);
  y = step;
  Matrix dy(d,n);
  diff(y, dy);
  double w = 2*PI/T;
  dy.scale(w);//dy = w y'(t)
  size_t i;
  Matrix shadowm;
  Vector v(d);
  for(i=0; i < n; i++){
    shadowm.shadow(A.getColumn(i),d,d);
    mult(shadowm, y.getColumn(i), v);
    y.getColumn(i) = v;
  }//y = A(t)*y(t)
  dy.sub(y); //dy = w y'(t) - A(t)*y(t)
  y = dX; 
  y.scale(wstep); //y = dw*x'(t)
  dy.add(y);//dy = w y'(t) - A(t)*y(t) + dw*x'(t)
  dy.sub(Res);
  return dy.fronorm();
} 


void PeriodicOrbit::smooth(double alpha, double beta){
  size_t d = pOrbit.getm();
  size_t n = pOrbit.getn();
  Matrix XX(d,n);
  Matrix RRes(d,n);
  double error1 = Res.fronorm();
  size_t i;
  for(i=0; i < 5; i++){
    matrixsmooth(pOrbit, alpha, XX);
    ComputeResidual(XX,2*PI/T,RRes,vfield);
    double error2 = RRes.fronorm();
    if(error2 < error1){
      pOrbit = XX;
      updateCache();
      return;
    }
    else
      alpha = alpha/beta;
  }
}

void PeriodicOrbit::getPeriodicOrbit(Matrix& M, double& Period){
  Period = T;
  M = pOrbit;
}

void PeriodicOrbit::updateCache(){
  diff(pOrbit, dX);
  ComputeResidual(pOrbit, 2*PI/T, dX, Res, vfield);
  ComputeLinearization(pOrbit, A, jacobian);
  ComputeGradDirn(2*PI/T, dX, Res, A, GradX, Gradw);
}

double PeriodicOrbit::dotStepAndGrad(const Matrix& step, double wstep){
  size_t d = step.getm();
  size_t n = step.getn();
  size_t i, j;
  double sum=0.0;
  for(j=0; j < n; j++)
    for(i=0; i < d; i++)
      sum+=step(i,j)*GradX(i,j);
  sum+=wstep*Gradw;
  return 2*sum;
}

void PeriodicOrbit::getPowellStep(Matrix& step, double& wstep){
  size_t d = pOrbit.getm();
  size_t n = pOrbit.getn();
  assert((step.getm()==d)&&(step.getn()==n));
  Matrix gX(d,n);
  Matrix dgX(d,n);
  double gw;
  double w = 2*PI/T;
  gX = GradX;
  gw = Gradw;
  //#ifdef DEBUG
  double relgradnorm = gX.fronorm();
  cout<<"Abs norm of gradient: "<<relgradnorm<<endl;
  relgradnorm /= pOrbit.fronorm();
  cout<<"Rel norm of gradient: "<<relgradnorm<<endl;
  //#endif
  step = gX;
  wstep = gw;
  size_t i;
  diff(gX, dgX);
  dgX.scale(w); //dgX is w*y'(t) where y(t) is gradient
  for(i=0; i < n; i++){
    Matrix shadowm;
    shadowm.shadow(A.getColumn(i),d,d);
    Vector shadowv;
    shadowv.shadow(gX.getColumn(i));
    mult(shadowm,shadowv);
  }//gX is A(t)*y(t)
  
  dgX.sub(gX);//dgX is w*y'(t) - A(t)y
  if((gw > 1e-50)||(gw < -1e-50)){
    dX.scale(gw);
    dgX.add(dX); //dgX  is w*y'(t) - A(t)y + dw*x'(t)
    dX.scale(1/gw);
  }
    
  double alpha;
  alpha = ComputeProjectionFactor(dgX, Res);
  step.scale(alpha);
  wstep *= alpha;
}


void PeriodicOrbit::setSintvl(int intvl){
  if(spts != NULL)
    delete spts;
  int n = pOrbit.getn();
  assert((intvl > 0) && (intvl <= n));
  numOfSpts = n/intvl;
  spts = new int[numOfSpts];
  int i;
  spts[0] = intvl;
  for(i=1; i < numOfSpts; i++)
    spts[i] = spts[i-1]+intvl;
  spts[numOfSpts-1] = n;
  allocateYYggffvf();
}

void PeriodicOrbit::allocateYYggffvf(){
  int d = pOrbit.getm();
  if(YY != NULL)
    delete YY;
  YY = new Matrix(d*d, numOfSpts);
  if(gg != NULL)
    delete gg;
  gg = new Matrix(d, numOfSpts);
  if(ff != NULL)
    delete ff;
  ff = new Matrix(d, numOfSpts);
  if(vf != NULL)
    delete vf;
  vf = new Matrix(d, numOfSpts);
}
  

void PeriodicOrbit::computeY(){
  double w = 2*PI/T;
  A.scale(1/w);
  FindFundSolnA(A,Y);
  A.scale(w);
}

void PeriodicOrbit::computef(){
  double w = 2*PI/T;
  A.scale(1/w);
  Res.scale(1/w);
  ApplyFundSolnC(A, Res, f);
  A.scale(w);
  Res.scale(w);
}

void PeriodicOrbit::computeg(){
  double w = 2*PI/T;
  A.scale(1/w);
  dX.scale(1/w);
  ApplyFundSolnC(A, dX, g);
  A.scale(w);
  dX.scale(w);
}

void PeriodicOrbit::computeYY(){
  int d = pOrbit.getm();
  int n = pOrbit.getn();
  Matrix I(d,d);
  Matrix y(d,d);
  int i, j;
  for(i=0; i < d; i++)
    for(j=0; j < d; j++)
      I(i,j) = (i==j)?1:0;
  y = I;
  j = 0;
  for(i=0; i <= n; i++){
    if(i==spts[j]){
      Matrix m1;
      m1.shadow((*YY).getColumn(j), d, d);
      m1 = y;
      j++;
      if(i==n)
	break;
      else
	y = I;
    }
    assert(i < n);
    Matrix m2;
    m2.shadow(Y.getColumn(i), d, d);
    mult(m2, y);
  }
}

void PeriodicOrbit::computeff(){
  int d = pOrbit.getm();
  int n = pOrbit.getn();
  Vector zz(d);
  Vector y(d);
  int i, j;
  for(i=0; i < d; i++)
    zz(i) = 0;
  y = zz;
  j = 0;
  for(i=0; i <= n; i++){
    if(i==spts[j]){
      (*ff).getColumn(j) = y;
      j++;
      if(i==n)
	break;
      else
	y = zz;
    }
    Matrix m1;
    m1.shadow(Y.getColumn(i), d, d);
    mult(m1, y);
    y.add(f.getColumn(i));
  }
}

void PeriodicOrbit::computegg(){
  int d = pOrbit.getm();
  int n = pOrbit.getn();
  Vector zz(d);
  Vector y(d);
  int i, j;
  for(i=0; i < d; i++)
    zz(i) = 0;
  y = zz;
  j = 0;
  for(i=0; i <= n; i++){
    if(i==spts[j]){
      (*gg).getColumn(j) = y;
      j++;
      if(i==n)
	break;
      else
	y = zz;
    }
    Matrix m1;
    m1.shadow(Y.getColumn(i), d, d);
    mult(m1, y);
    y.add(g.getColumn(i));
  }
}

void PeriodicOrbit::computevf(){
  int i;
  vfield(pOrbit.getColumn(0), (*vf).getColumn(0));
  for(i=0; i < numOfSpts-1; i++)
    vfield(pOrbit.getColumn(spts[i]),(*vf).getColumn(i+1));
  vf->scale(tinyFactor);
}

void PeriodicOrbit::computeLPMatrix(){
  int d = pOrbit.getm();
  int n = pOrbit.getn();
  if(LPMatrix != NULL)
    delete LPMatrix;
  int ns = numOfSpts;
  LPMatrix = new SparseMatrix(ns*d+1, (ns+1)*d*d+3*(ns*d+1));
  int i;
  if(ns > 1)
    for(i=0; i<ns*d; i++)//put ones on the diagonal
      LPMatrix->insert(i,i,1);
  
  YY->scale(-1);//put -YY(i) on subdiagonal
  for(i=0; i<ns-1; i++){
    Matrix m1;
    m1.shadow(YY->getColumn(i), d, d);
    LPMatrix->insert((i+1)*d, i*d, m1);
  }
  Matrix m1;
  m1.shadow(YY->getColumn(ns-1), d, d);
  if(ns > 1)
    LPMatrix->insert(0, (ns-1)*d, m1);
  else{ //insert -Y and I together
    Matrix mm(d,d);
    mm = m1;
    for(i=0; i < d; i++)
      mm(i,i) = mm(i,i)+1;
    LPMatrix->insert(0, 0, mm);
  }
  YY->scale(-1);

  
  for(i=0; i < ns-1; i++){//get gg(i) on last column
    m1.shadow(gg->getColumn(i), d, 1);
    LPMatrix->insert((i+1)*d, ns*d, m1);
  }
  m1.shadow(gg->getColumn(ns-1), d, 1);
  LPMatrix->insert(0, ns*d, m1);
  
  for(i=0; i < ns; i++){//put vf(i)s on last row
    m1.shadow(vf->getColumn(i), 1, d);
    LPMatrix->insert(ns*d, i*d, m1);
  }
  //LPMatrix->output("DBG/LP.dat");
  LPMatrix->formMatrix();
}

void PeriodicOrbit::computeShootingStep(Matrix& sstep, double& wstep){
  int ns = numOfSpts;
  int d = pOrbit.getm();
  Vector rhs(ns*d+1);
  rhs(ns*d) = 0;
  int i;
  for(i=0; i < ns-1; i++)
    rhs.getShadow((i+1)*d, (i+2)*d-1) = ff->getColumn(i);
  rhs.getShadow(0, d-1) = ff->getColumn(ns-1);
  Vector soln(ns*d+1);
  solve(*LPMatrix, rhs, soln);
  //rhs.output("DBG/rhs.dat");
  Matrix m1;
  m1.shadow(soln, d, ns);
  sstep = m1;
  wstep = soln(ns*d);
}
    
void PeriodicOrbit::propagatey(Vector& yi, int i, double wstep,
			       int k, Vector& yk){
  int j;
  int d = pOrbit.getm();
  Vector yj(d);
  yj = yi;
  assert(i<=k);
  for(j=i; j < k; j++){
    Matrix m1;
    m1.shadow(Y.getColumn(j), d, d);
    mult(m1, yj);
    yj.add(f.getColumn(j));
    Vector v1(d);
    v1 = g.getColumn(j);
    v1.scale(-wstep);
    yj.add(v1);
  }
  yk = yj;
}

void PeriodicOrbit::convertSStep2NStep(Matrix& sstep, double wstep,
				       Matrix& nstep){
  int d = pOrbit.getm();
  int n = pOrbit.getn();
  int i, j;
  j = 0;
  nstep.getColumn(0) = sstep.getColumn(0);
  for(i=1; i < n; i++){
    if(i==spts[j]){
      nstep.getColumn(i) = sstep.getColumn(j+1);
      j++;
    }
    else
      propagatey(nstep.getColumn(i-1), i-1, wstep, i, nstep.getColumn(i));
  }
}

static Clock clk;
static double timeA=0.0;
static double timeB=0.0;
static double timeC=0.0;

void PeriodicOrbit::getNewtonStep(Matrix& nstep, double& wstep, int flag){
  if(flag==0){
    computeY();
    computeg();
    clk.wtic();
    computeYY();
    computegg();
    computevf();
    timeA += clk.wtoc();
    clk.wtic();
    computeLPMatrix();
    timeB += clk.wtoc();
  }
  computef();
  clk.wtic();
  computeff();
  timeA += clk.wtoc();
  Matrix sstep(nstep.getm(), numOfSpts);
  clk.wtic();
  computeShootingStep(sstep, wstep);
  timeB += clk.wtoc();
  clk.wtic();
  convertSStep2NStep(sstep, wstep, nstep);
  timeA += clk.wtoc();
}
  
void PeriodicOrbit::getLinearResidual(Matrix& step, double wstep, 
				      Matrix& linRes){
  assert(step.getm()==pOrbit.getm());
  assert(step.getn()==pOrbit.getn());
  assert(linRes.getm()==step.getm());
  assert(linRes.getn()==step.getn());
  int d = step.getm();
  int n = step.getn();
  Matrix tstep(d,n);
  tstep = step;
  diff(step, linRes);
  linRes.scale(2.0*PI/T);//linRes = w y'(t)
  int i;
  for(i=0; i < n; i++){
    Matrix mm;
    mm.shadow(A.getColumn(i), d, d);
    Vector vv;
    vv.shadow(tstep.getColumn(i));
    mult(mm, vv);
  }
  linRes.sub(tstep);//linRes = wy'(t) - A(t)y
  tstep = dX;
  tstep.scale(wstep);
  linRes.add(tstep);//linRes = wy'(t) - A(t)Y + dw x'(t)
  linRes.sub(Res);//linRes = wy'(t) - A(t)Y + dw x'(t) - r(t)
}

void PeriodicOrbit::setResidual(Matrix& res){
  Res = res;
}

void printSTUFF(){
  cout<<"shooting time: "<<timeA<<endl;
  cout<<"superlu time: "<<timeB<<endl;
  timeA = 0;
  timeB = 0;
}

void SolveQuadEqn(double a, double b, double c, 
		  double& r1, double& r2){
  double d = b*b - 4*a*c;
  assert(d>0);
  d = sqrt(d);
  if(d*b > 0){
    r1 = (-b-d)/(2*a);
    r2 = (2*c)/(-b-d);
  }
  else{
    r1 = (-b+d)/(2*a);
    r2 = (2*c)/(-b+d);
  }
}

void dogleg(const Matrix& pstep, double pwstep,
	    const Matrix& nstep, double nwstep,
	    double delta,
	    Matrix& step, double& wstep){
#ifdef DEBUG
  cout<<"dogleg step ..."<<endl; //this will be in dogleg region
#endif
  size_t d = step.getm();
  size_t n = step.getn();
  Matrix tempM(d,n);
  double c1;
  tempM=nstep; tempM.sub(pstep);
  c1 = tempM.fronorm();
  c1 = c1*c1 + (nwstep - pwstep)*(nwstep - pwstep);
  double c2;
  tempM = nstep;
  tempM.sub(pstep);
  tempM.mul(pstep);
  tempM.scale(2);
  size_t i,j;
  c2 = 0;
  for(i=0; i < d; i++)
    for(j=0; j < n; j++)
      c2 += tempM(i,j);
  c2+=2*(nwstep-pwstep)*pwstep;
  double c3;
  c3 = pstep.fronorm();
  c3 = c3*c3 + pwstep*pwstep - delta*delta;
  double root1, root2;
  SolveQuadEqn(c1,c2,c3,root1,root2);
  double root;
  if((0<=root1)&&(root1<=1))
    root = root1;
  else if((0<=root2)&&(root2<=1))
    root = root2;
  else
    cout<<"something strange with leg of dog"<<endl;
  wstep = (pwstep)+root*(nwstep - pwstep);
  step = nstep;
  step.sub(pstep);
  step.scale(root);
  step.add(pstep);
}

void GetStepOfPowell(const Matrix& pstep, double pwstep, 
		     const Matrix& nstep, double nwstep, 
		     int InnerOuterFlag, 
		     double delta,
		     Matrix& step, double& wstep){
  assert(delta >= 0);
  if(InnerOuterFlag==0){
    double pstepsize = pstep.fronorm();
    pstepsize = sqrt(pstepsize*pstepsize+pwstep*pwstep);
    double factor = delta/pstepsize;
    assert(factor<=1+1e-7);
    step = pstep; wstep = pwstep;
    step.scale(factor); wstep *= factor;
  }
  else if(InnerOuterFlag==1){
    double pstepsize = pstep.fronorm();
    pstepsize = sqrt(pstepsize*pstepsize+pwstep*pwstep); 
    double nstepsize = nstep.fronorm();
    nstepsize = sqrt(nstepsize*nstepsize+nwstep*nwstep);
    assert(nstepsize > pstepsize);
    if(delta >= nstepsize){
      step = nstep;
      wstep = nwstep;
      return;
    }
    else
      dogleg(pstep, pwstep, nstep, nwstep, delta,
	     step, wstep);
  }
}


int stepOk(PeriodicOrbit& po, Matrix& step, double wstep){
  double error = po.residualError();
  double errorstep = po.residualErrorAfterStep(step, wstep);
  if(errorstep < error)
    return 1;
  else 
    return 0;
}

int stepGood(PeriodicOrbit& po, Matrix& step, double wstep){
  double CONST = 0.1;
  double error = po.residualError();
  double errorstep = po.residualErrorAfterStep(step, wstep);
  double dotprod = po.dotStepAndGrad(step, wstep);
  if(dotprod > 0)
    return 0;
  if(errorstep*errorstep < error*error + CONST*dotprod)
    return 1;
  else
    return 0;
}

int predictionGood(PeriodicOrbit& po, Matrix& step, double wstep){
  double CONST = 0.05;
  double errorstep = po.residualErrorAfterStep(step, wstep);
  double predictederror = po.predictedResidualError(step, wstep);
  if(fabs(predictederror - errorstep)/errorstep < CONST)
    return 1;
  else
    return 0;
}

double returnStepSize(Matrix& step, double wstep){
  double size = step.fronorm();
  size = sqrt(size*size + wstep*wstep);
  return size;
}


void IncreaseDeltaAndComputeStep(PeriodicOrbit& po, double& delta,
				 Matrix& pstep, double pwstep,
				 Matrix& nstep, double nwstep,
				 Matrix& step, double& wstep){
  double pstepsize = returnStepSize(pstep, pwstep);
  double nstepsize = returnStepSize(nstep, nwstep);
  while(predictionGood(po, step, wstep) &&
	stepGood(po, step, wstep) && (delta < nstepsize)){
#ifdef DEBUG
    cout<<"increasing delta..."<<endl;
#endif
    delta = 2*delta;
    int InnerOuterFlag = (delta <= pstepsize)? 0:1;
    GetStepOfPowell(pstep, pwstep, nstep, nwstep, InnerOuterFlag, delta,
		  step, wstep);
  }
  if(!stepOk(po,step, wstep)){
    delta = delta/2;
    int InnerOuterFlag = (delta <= pstepsize)? 0:1;
    GetStepOfPowell(pstep, pwstep, nstep, nwstep, InnerOuterFlag, delta,
		  step, wstep);
  }
}
				 
void DecreaseDeltaAndComputeStep(PeriodicOrbit& po, double& delta,
				 Matrix& pstep, double pwstep,
				 Matrix& nstep, double nwstep,
				 Matrix& step, double& wstep){

  int LIMIT = 5;
  int count = 0;
  double pstepsize = returnStepSize(pstep, pwstep);
  while((!stepOk(po, step, wstep)) && (count < LIMIT)){
#ifdef DEBUG
    cout<<"halving delta..."<<endl;
#endif
    delta = delta/2; 
    int InnerOuterFlag = (delta <= pstepsize)? 0:1;
    count++;
    GetStepOfPowell(pstep, pwstep, nstep, nwstep, InnerOuterFlag, delta,
		  step, wstep);
    Matrix sstep(step.getm(), step.getn());
    sstep = step;
    matrixsmooth(sstep, 0.5, step);
  }
  if(count==LIMIT){
#ifdef DEBUG
    cout<<"Trust region step failed"<<endl;
#endif
    int i, j;
    for(i=0; i < step.getm(); i++)
      for(j=0; j < step.getn(); j++)
	step(i, j) = 0.0;
    wstep = 0;
  }
}

//po is the periodic orbit
//delta is the trust region radius at input
//delta is the new trust region radius at output
//if delta is -1 the Powell step is used to find the trust region radius
//step contains the trust region step
//wstep contains the timestep
void ComputeTrustStep(PeriodicOrbit& po, double& delta,
		      Matrix& step, double& wstep){
 
  int InnerOuterFlag=-1;//flag=0 => inner and flag=1 => outer
  int d = step.getm();
  int n = step.getn();
  Matrix pstep(d,n); double pwstep;
  Matrix nstep(d,n); double nwstep;
  assert((delta>0)||(delta == -1));
  po.getPowellStep(pstep, pwstep);
  double pstepsize = returnStepSize(pstep, pwstep);
  //#ifdef DEBUG
  cout<<"pstepsize: "<<pstepsize<<endl;
  //#endif
  if(delta == -1){
    delta = pstepsize;
    InnerOuterFlag = 0;
  }
  else if(delta < pstepsize)
    InnerOuterFlag = 0;
  else{
    InnerOuterFlag = 1;
    po.getNewtonStep(nstep, nwstep);
    double nstepsize = returnStepSize(nstep, nwstep);
    //#ifdef DEBUG
    cout<<"nstepsize: "<<nstepsize<<endl;
    //#endif
    if(delta > 1.05*nstepsize)
      delta = 1.05*nstepsize;
  }
  GetStepOfPowell(pstep, pwstep, nstep, nwstep, InnerOuterFlag, delta,
		  step, wstep);
  if(stepGood(po, step, wstep))
    if(predictionGood(po, step, wstep)){
      if(InnerOuterFlag == 0)
	po.getNewtonStep(nstep, nwstep);
      double nstepsize = returnStepSize(nstep, nwstep);
      //#ifdef DEBUG
      cout<<"nstepsize: "<<returnStepSize(nstep, nwstep)<<endl;
     //#endif 
      if(delta > 1.05*nstepsize)
	delta = 1.05*nstepsize;
      IncreaseDeltaAndComputeStep(po, delta, pstep, pwstep, 
				  nstep, nwstep, step, wstep);
    }
    else
      return;
  else{
    delta /= 1.2;
    DecreaseDeltaAndComputeStep(po, delta, pstep, pwstep,
				nstep, nwstep, step, wstep);
  }
}