#include <mkl_blas.h>
#include <mkl_lapack.h>
#include "Matrix.hh"
#include "SparseMatrix.hh"
#include <math.h>
#include <fftw.h>
#include <rfftw.h>
#include <iostream>

static double* DWORKSPACE(NULL);
static int* IWORKSPACE(NULL);
static size_t WORKSPACESIZE(0);

extern "C" void DGEMV(char *, 
		       int *, int *, 
		       double *, 
		       double *, int *,
		       double *, int *,
		       double *,
		       double *, int *);

extern "C" void dgemm_(char *, char *,
		       int *, int *, int *,
		       double *,
		       double *, int *,
		       double *, int *,
		       double *,
		       double *, int *);
		       

extern "C" void DGEEV(char *, char *, int *, double *, int *, 
		       double *, double *, double *, int *,
	  	       double *, int *, double *, int *, int *);

extern "C" void DGESV(int *, int *, double *, int *, int *,
		       double *, int *, int *);

extern "C" void DGESVD(char *, char *, int *, int *, double *, int *,
			double *, double *, int *, double *, int *,
			double *, int *, int *);

extern "C" void DGPADM(int *, int *, double *, 
			double *, int *,
			double *,
			int *,
			int *,
			int *, int *, int *);


extern "C"{
int sparseEntryComp(const void * ptr1, const void *ptr2);
	  }

 // Get working space of size = n
inline void GETSPACE(size_t n){

  if(WORKSPACESIZE < n){
    if(WORKSPACESIZE!=0){
      delete[] DWORKSPACE;
      delete[] IWORKSPACE;
    }
    DWORKSPACE = new double[n];
    IWORKSPACE = new int[n];
    WORKSPACESIZE = n;
  }
}

//free working space
inline void FREESPACE(){
  if(WORKSPACESIZE != 0){
    delete[] DWORKSPACE;
    delete[] IWORKSPACE;
  }
}

void GetSpace(size_t n){
  GETSPACE(n);
}

void FreeSpace(){
  FREESPACE();
  fft_destroy_plan();
}
  
			

//N becomes equal to M^T
void transpose(const Matrix& M, Matrix& N){
    size_t m = M.getm();
    size_t n = M.getn();
    assert(m==N.getn());
    assert(n==N.getm());

    for(size_t i=0; i < n; i++)
	for(size_t j=0; j < m; j++)
	    N(i,j) = M(j,i);
}

//prod = M*v
void mult(const Matrix& M, const Vector& v, Vector& prod){
  assert(M.getn()==v.getSize());
  assert(prod.getSize()==v.getSize());
  int m = M.getm(), n = M.getn();
  int incx = 1;
  int lda = M.lda;
  char trans[3] = "N ";
  double alpha = 1.0;
  double beta = 0.0;

  DGEMV(trans, 
	 &m, &n,
	 &alpha,
	 M.data, &lda,
	 v.data, &incx,
	 &beta,
	 prod.data, &incx);//BLAS call
}

//v = M*v
void mult(const Matrix& M, Vector& v){
  Vector vv(v.getSize());
  vv = v;
  mult(M,vv,v);
}

//prod = M1*M2
void mult(const Matrix& M1, const Matrix& M2, Matrix& prod){
  assert(M1.getn()==M2.getm());
  assert(M1.getm()==prod.getm());
  assert(M2.getn()==prod.getn());
  char trans[3] = "N ";
  int m = M1.getm(), n = M1.getn(), k = M2.getn();
  int lda1 = M1.lda;
  int lda2 = M2.lda;
  int lda3 = prod.lda;
  double alpha = 1.0, beta = 0.0;

  DGEMM(trans, trans, 
	 &m, &k, &n, 
	 &alpha,
	 M1.data, &lda1,
	 M2.data, &lda2,
	 &beta,
	 prod.data, &lda3);//BLAS call
}

//M2 = M1*M2
void mult(const Matrix& M1, Matrix & M2){
  assert(M1.getm()==M1.getn());
  Matrix MM(M2.getm(), M2.getn());
  MM = M2;
  mult(M1, MM, M2);
}
  
//finds soln to M*soln = v
void solve(const Matrix& M, const Vector& v, Vector& soln){
  size_t n = M.getm();
  assert((n==M.getn())&&(n==v.getSize()));
  assert((n==soln.getSize()));
  GETSPACE(n*n+n);
  Matrix A;
  A.size1 = n;
  A.size2 = n;
  A.lda = n;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M;
  soln = v;
  int nrhs = 1;
  int ldb = n;
  int info;

  DGESV((int *)&n, &nrhs,
	 A.data, (int *)&n,
	 IWORKSPACE,
	 soln.data, &ldb,
	 &info);//LAPACK call
}
  
//v becomes soln of M*x = v
void solve(const Matrix& M, Vector& v){
  solve(M,v,v);
}


//S becomes vector of singular values of M
void svd(const Matrix& M, Vector& S){
  char jobu[3] = "N ";
  char jobvt[3] = "N ";
  size_t m = M.getm();
  size_t n = M.getn();
#ifdef DEBUG
  size_t min_d = (m<=n)? m:n;
#endif
  assert(S.getSize()==min_d);
  GETSPACE(m*n+5*m+5*n);
  Matrix A;
  A.size1 = m;
  A.size2 = n;
  A.lda = m;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M;
  int lwork = 5*m+5*n;
  int dummy = n;
  int info;

  //matrix seen by this routine is A=M
  DGESVD(jobu, jobvt, 
	  (int *)&m, (int *)&n, 
	  A.data, (int *)&(A.lda),
	  S.data, 
	  (double *)(NULL), (int *)(&dummy),
	  (double *)NULL, (int *)(&dummy), 
	  DWORKSPACE+m*n, 
	  &lwork, &info);//LAPACK call
}

//S becomes vector of singular values of M
//U becomes matrix of left singular vectors of M
//V becomes matrix of right singular vectors of M
void svd(const Matrix& M, Vector& S, Matrix& U, Matrix& V){
  char jobu[3] = "N ";
  char jobvt[3] = "O ";
  size_t m = M.getm();
  size_t n = M.getn();
  size_t min_d = (m<=n)? m:n;
  assert(S.getSize()== min_d);
  assert(V.getm()== n);
  assert(V.getn()== min_d);
  assert(U.getm()==m);
  assert(U.getn()==min_d);
  GETSPACE(m*n+5*m+5*n);
  Matrix A; //becomes M
  A.size1 = m;
  A.size2 = n;
  A.lda = m;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M;
  int lwork = 5*m+5*n;
  int info;
  int dummy=m+n;
  //matrix seen by this routine is A=M
  DGESVD(jobu, jobvt,
	  (int *)&m, (int *)&n, 
	  A.data, (int *)&(A.lda),
	  S.data,
	  (double *)(NULL), (int *)(&dummy),
	  (double *)(NULL), (int *)(&dummy),
	  DWORKSPACE+m*n, &lwork, &info);
  
  transpose(A.getShadow(0,min_d-1,0,n-1), V);
 
  mult(M,V,U);
  for(size_t i=0; i < U.getm(); i++)
    for(size_t j=0; j < U.getn(); j++)
      U(i,j) = U(i,j)/S[j];
}

//V becomes matrix of eigenvectors of M (as in LAPACK)
//er+sqrt(-1)ei becomes the vector of eigenvalues
void eig(const Matrix& M, Matrix& V, Vector& er, Vector& ei){
  size_t n = M.getm();
  assert(n==M.getn());
  char jobvl[3] = "N ";
  char jobvr[3] = "V ";
  GETSPACE(n*n+5*n);
  Matrix A;
  A.size1 = n;
  A.size2 = n;
  A.lda = n;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M; //DGEEV overwrites matrix

  double * work = (DWORKSPACE)+n*n;
  int lwork=5*n;
  int info;
  int dummy=n;
   
  DGEEV(jobvl, jobvr,
	 (int *)&n, 
	 A.data, (int *)&(A.lda),
	 er.data, 
	 ei.data, 
	 (double *)NULL, (int *)(&dummy),
	 V.data, (int *)&(V.lda), 
	 work, &lwork, &info);//LAPACK call
}

//er+sqrt(-1)*ei becomes vector of eigenvalues
void eig(const Matrix& M, Vector& er, Vector& ei){
  size_t n = M.getm();
  assert(n==M.getn());
  assert(n==er.getSize());
  assert(n==ei.getSize());
  char jobvl[3] = "N ";
  char jobvr[3] = "N ";
  GETSPACE(n*n+5*n);
  Matrix A;
  A.size1 = n;
  A.size2 = n;
  A.lda = n;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M; //DGEEV overwrites matrix

  double * work = (DWORKSPACE)+n*n;
  int lwork=5*n;
  int info;
  int dummy=n;
  DGEEV(jobvl, jobvr, 
	 (int *)&n,
	 A.data, (int *)&(A.lda),
	 er.data,
	 ei.data,
	 (double *)NULL, (int *)(&dummy),
	 (double *)NULL, (int *)(&dummy), 
	 work, &lwork, &info);//LAPACK call
}


//M2 becomes inverse of M1.
void invert(const Matrix& M1, Matrix& M2){
  size_t n = M1.getm();
  assert(M1.getm()==M1.getn());
  assert(M2.getm()==M2.getn());
  assert(M1.getm()==M2.getm());
  GETSPACE(n*n+n);
  Matrix A;
  A.size1 = n;
  A.size2 = n;
  A.lda = n;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M1;
  for(size_t i=0; i < n; i++)
    for(size_t j=0; j < n; j++)
      M2(i,j) = (i==j)?1:0;
  int nrhs = n;
  int info;
  DGESV((int *)&n, &nrhs,
	 A.data, (int *)&(A.lda),
	 IWORKSPACE,
	 M2.data, (int *)&(M2.lda),
	 &info);//LAPACK call
}

double det(const Matrix& M){
  size_t n = M.getm();
  assert(M.getm()==M.getn());
  GETSPACE(n*n+n);
  Matrix A;
  A.size1 = n;
  A.size2 = n;
  A.lda = n;
  A.data = DWORKSPACE;
  A.owner = 0;
  A = M;
  int nrhs = 0;
  int lda = n;
  int info;
  DGESV((int *)&n, &nrhs,
	 A.data, (int *)&(A.lda),
	 IWORKSPACE,
	 (double *)NULL, (int *)&(lda),
	 &info);//LAPACK call

  double rVal = 1.0;
  for(size_t i=0; i < n; i++)
    rVal *= A(i,i);
  return rVal;
}


rfftw_plan getFwdPlan(int n){
  static const unsigned int N = 4;
  static rfftw_plan plan[N]; //circular queue of plans
  static int length[N]; //their size
  static size_t first = 0;
  static size_t count = 0;
  
  if(n < 0){ //destroy all the plans
    for(size_t i=0; i < count; i++)
      rfftw_destroy_plan(plan[(first+i)%N]);
    count = 0;
    first = 0;
    return NULL;
  }
  
  for(size_t i=0; i < count; i++)
    if(length[(first+i)%N]==n)
      return plan[(first+i)%N];
  
  first = (count>0)?(first+(N-1))%N:0;
  if(count < N)
    count++;
  else
    rfftw_destroy_plan(plan[first]);
  length[first] = n;
  plan[first] = rfftw_create_plan(n, FFTW_REAL_TO_COMPLEX,
				  FFTW_ESTIMATE);//fftw call
  return(plan[first]);
}

rfftw_plan getBwdPlan(int n){
  static const unsigned int N = 4;
  static rfftw_plan plan[N]; //circular queue of plans
  static int length[N]; //their length
  static size_t first = 0;
  static size_t count = 0;
  
  if(n < 0){ //destroy all the plans
    for(size_t i=0; i < count; i++)
      rfftw_destroy_plan(plan[(first+i)%N]);
    count = 0;
    first = 0;
    return NULL;
  }
  
  for(size_t i=0; i < count; i++)
    if(length[(first+i)%N]==n)
      return plan[(first+i)%N];
  
  first = (count>0)?(first+(N-1))%N:0;
  if(count < N)
    count++;
  else
    rfftw_destroy_plan(plan[first]);
  length[first] = n;
  plan[first] = rfftw_create_plan(n, FFTW_COMPLEX_TO_REAL,
				  FFTW_ESTIMATE);//fftw call
  return(plan[first]);
}


				    
void fft(const Vector& x, Vector& xf){
  assert(x.getSize()==xf.getSize());
  rfftw_one(getFwdPlan(x.getSize()), x.data, xf.data);//fftw call
  size_t n = xf.getSize();
  size_t i;
  for(i=0; i < n; i++)//normalize Fourier transform
    xf[i]/=n;
}



void ifft(const Vector& xf, Vector& x){
  assert(x.getSize()==xf.getSize());
  size_t n = x.getSize();
  GETSPACE(n);
  Vector v;
  v.size = n;
  v.data = DWORKSPACE;
  v.owner = 0;
  v = xf;
  rfftw_one(getBwdPlan(n), v.data, x.data);//fftw call
}


void diff(const Vector& x, Vector &xd){
  assert(x.getSize()==xd.getSize());
  size_t n = x.getSize();
  GETSPACE(n);
  Vector xf;
  xf.size = n;
  xf.data = DWORKSPACE;
  xf.owner = 0;
  rfftw_one(getFwdPlan(n), x.data, xf.data);
  size_t k;
  for(k=0; k < n; k++)
    xf(k)/=n;
  xf(0) = 0.0;
  for(k=1; k <= (n+1)/2-1; k++){//differentiate in Fourier domain
    xf(k) *= k;
    xf(n-k) *= k;
    double temp=xf(k);
    xf(k) = -xf(n-k);
    xf(n-k) = temp;
  }
  
  if(n%2==0)
    xf(n/2) = 0.0;
  
  rfftw_one(getBwdPlan(n), xf.data, xd.data);//fftw call
}

void diff(const Matrix& M, Matrix& N){
  size_t m = M.getm();
  size_t n = M.getn();
  assert(m==N.getm());
  assert(n==N.getn());
  Vector v(n), vd(n);
  size_t i,j;
  for(i=0; i < m; i++){
    for(j=0; j < n; j++)
      v(j) = M(i,j);
    diff(v,vd);
    for(j=0; j < n; j++)
      N(i,j) = vd(j);
  }
}

void shift(const Vector& x, double h, Vector& y){
  assert(x.getSize()==y.getSize());
  size_t n = x.getSize();
  GETSPACE(n);
  Vector xf;
  xf.size = n;
  xf.data = DWORKSPACE;
  xf.owner = 0;
  //cout<<"shiftv: point A"<<endl;
  rfftw_one(getFwdPlan(n), x.data, xf.data);
  //cout<<"shiftv: point B"<<endl;
  size_t k;
  for(k=0; k<n; k++)
    xf(k) /= n;

  for(k=1; k <= (n+1)/2 - 1; k++){//shift in Fourier domain
    double c = cos(k*h), s = sin(k*h);
    double a = xf(k), b = xf(n-k); 
    //(c+is)*(a+ib)
    xf(k) = c*a - s*b;
    xf(n-k) = c*b + s*a;
  }

  if(n%2==0)//this component cannot be shifted properly
    xf(n/2) *= cos(n*h/2.0);
  //cout<<"shiftv: point C"<<endl;
  rfftw_one(getBwdPlan(n), xf.data, y.data);
  //cout<<"shiftv: point D"<<endl;
}

void shift(const Matrix& M, double h, Matrix& N){
  //cout<<"shift: point A"<<endl;
  size_t m = M.getm();
  size_t n = M.getn();
  assert(m==N.getm());
  assert(n==N.getn());
  Vector v(n), w(n);
  size_t i,j;
  for(i=0; i < m; i++){
    for(j=0; j < n; j++)
      v(j) = M(i,j);
    shift(v,h,w);
    //cout<<"shift: point Bi"<<endl;
    for(j=0; j < n; j++)
      N(i,j) = w(j);
  }
  //cout<<"shift: point C"<<endl;
}


void filter(const Vector& x, double alpha, Vector& y){
  assert(x.getSize()==y.getSize());
  assert((0.0<=alpha)&&(alpha<=1.0));
  size_t n = x.getSize();
  GETSPACE(n);
  Vector xf;
  xf.size = n;
  xf.data = DWORKSPACE;
  xf.owner = 0;
  rfftw_one(getFwdPlan(n), x.data, xf.data);
  size_t k;
  for(k=0; k<n; k++)
    xf(k) /= n;
  if(n%2==0)
    xf(n/2) = 0;
  size_t w = (size_t)(alpha*n/2);
  if(n%2==0)
    for(k=1; k<w; k++){
      xf(n/2+k) = 0;
      xf(n/2-k) = 0;
    }
  else
    for(k=0; k<w; k++){
      xf(n/2-k) = 0;
      xf(n/2+1+k)=0;
    }
  rfftw_one(getBwdPlan(n), xf.data, y.data);
}

void fft_destroy_plan(){//destroy all plans
  getFwdPlan(-1);
  getBwdPlan(-1);
}


void expm(const Matrix& M1, Matrix& M2){
  assert(M1.size1 == M1.size2);
  assert(M2.size1 == M2.size2);
  assert(M1.size1 == M2.size1);
  int ideg = 6;
  int m = M1.getm();
  double t = 1.0;
  int ldh = M1.lda;
  GETSPACE(4*m*m+ideg+1);
  int lwsp = 4*m*m+ideg+1;
  int iexph;
  int ns;
  int iflag;

  DGPADM(&ideg, &m, &t, M1.data, &ldh,
	  DWORKSPACE, &lwsp,
	  IWORKSPACE, 
	  &iexph, &ns, &iflag);//calls EXPOKIT
  assert(iflag==0);
  
  Matrix tempM;
  tempM.size1 = m;
  tempM.size2 = m;
  tempM.lda = m;
  tempM.data = (DWORKSPACE+iexph-1);
  tempM.owner = 0;
  M2 = tempM;
}
	  

void mult(SparseMatrix& M, const Vector& v, Vector& prod){
  mult(M.A, v, prod);
}  


void  solve(SparseMatrix& M, const Vector& v, Vector& soln){
  solve(M.A, v, soln);
}

    

void nrspline(double* x, double *y, int n, double yp1, double ypn,
	      double *y2){
  int i,k;
  double p, qn, sig, un, *u;
  u = (double *)malloc((n-1)*sizeof(double));
  u = u-1;
  if(yp1 > 1e30)
    y2[1] = u[1] = 0.0;
  else{
    y2[1] = -0.5;
    u[1] = (3.0/(x[2]-x[1]))*((y[2]-y[1])/(x[2]-x[1])-yp1);
  }
  for(i=2; i<=n-1; i++){
    sig = (x[i]-x[i-1])/(x[i+1]-x[i-1]);
    p = sig*y2[i-1]+2.0;
    y2[i] = (sig-1.0)/p;
    u[i] = (y[i+1]-y[i])/(x[i+1]-x[i]) - (y[i]-y[i-1])/(x[i]-x[i-1]);
    u[i] = (6.0*u[i]/(x[i+1]-x[i-1])-sig*u[i-1])/p;
  }
  if(ypn>=1e30)
    qn=un=0.0;
  else{
    qn = 0.5;
    un = (3.0/(x[n]-x[n-1]))*(ypn-(y[n]-y[n-1])/(x[n]-x[n-1]));
  }
  y2[n] = (un-qn*u[n-1])/(qn*y2[n-1]+1.0);
  for(k=n-1;k>=1;k--)
    y2[k] = y2[k]*y2[k+1]+u[k];
  u = u+1;
  free(u);
}

void nrsplint(double *xa, double *ya, double *y2a, int n, double x,
	      double& y){
  int klo, khi, k;
  double h, b, a;
  klo = 1;
  khi = n;
  while(khi-klo > 1){
    k=(khi+klo) >> 1;
    if(xa[k] > x)
      khi = k;
    else
      klo = k;
  }
  h = xa[khi]-xa[klo];
  assert(!(h==0.0));
  a = (xa[khi]-x)/h;
  b = (x - xa[klo])/h;
  y = a*ya[klo]+b*ya[khi]+((a*a*a-a)*y2a[klo]+(b*b*b-b)*y2a[khi])*(h*h)/6.0;
}

void spline(const Vector& t, const Vector& v, 
	     double leftD, double rightD,
	     const Vector& tt, Vector& vv){
  assert(t.getSize()==v.getSize());
  assert(tt.getSize()==vv.getSize());
  int N = t.getSize();
  assert(N>=2);
  int i;
#ifdef DEBUG
  for(i=0; i < N-1; i++)
    assert(t(i) < t(i+1));
#endif 
  double yp1, ypn;
  yp1 = leftD;
  ypn = rightD;
  double* y2 = (double *)malloc(N*sizeof(double));
  nrspline(t.data-1, v.data-1, N, yp1, ypn, y2-1);
  for(i=0; i < vv.getSize(); i++)
    nrsplint(t.data-1, v.data-1, y2-1, N, tt(i), vv(i));
  free(y2);
}

void cspline(const Vector& t, const Vector& v,
	     const Vector& tt, Vector& vv){
  assert(t.getSize()==v.getSize());
  assert(tt.getSize()==vv.getSize());
  int N = t.getSize();
  double yp1 = (v(0)-v(1))/(t(0)-t(1));
  double ypn = (v(N-2)-v(N-1))/(t(N-2)-t(N-1));
  spline(t, v, yp1, ypn, tt, vv);
}
	 
void spline(const Vector& t, const Matrix& M, 
	    const Vector& leftD, const Vector& rightD,
	    const Vector& tt, Matrix& MM){
  assert(t.getSize()==M.getn());
  assert(tt.getSize()==MM.getn());
  assert(M.getm()==MM.getm());
  assert(M.getm()==leftD.getSize());
  assert(M.getm()==rightD.getSize());
  int m = M.getm();
  int n = M.getn();
  int nn = MM.getn();
  Vector v(n);
  Vector vv(nn);
  int i;
  for(i=0; i < m; i++){
    int j;
    for(j=0; j < n; j++)
      v(j) = M(i,j);
    spline(t, v, leftD(i), rightD(i), tt, vv);
    for(j=0; j < nn; j++)
      MM(i,j) = vv(j);
  }
}

void cspline(const Vector& t, const Matrix& M, const Vector& tt,
	     Matrix& MM){
  assert(t.getSize()==M.getn());
  assert(tt.getSize()==MM.getn());
  assert(M.getm()==MM.getm());
  int m = M.getm();
  int n = M.getn();
  int nn = MM.getn();
  Vector v(n);
  Vector vv(nn);
  int i;
  for(i=0; i < m; i++){
    int j;
    for(j=0; j < n; j++)
      v(j) = M(i,j);
    cspline(t, v, tt, vv);
    for(j=0; j < nn; j++)
      MM(i,j) = vv(j);
  }
}
  
void matrixsmooth(const Matrix& X, double alpha, Matrix& XX){
  size_t d = X.getm();
  size_t n = X.getn();
  assert(XX.getm()==d);
  assert(XX.getn()==n);
  size_t i, j;
  Vector vv(n), w(n);
  for(i=0; i < d; i++){
    for(j=0; j < n; j++)
      vv(j) = X(i,j);
    filter(vv,alpha,w);
    for(j=0; j < n; j++)
      XX(i,j) = w(j);
  }
}