#include "interpolate.h"
#include "math.h"
#include "global.h"

/* ----------------------------------------------------------------------------
 * Constructor used to get info from caller, and prepare other necessary data
 * ---------------------------------------------------------------------------- */
Interpolate::Interpolate(int nx, int ny, int nz, int ndm, doublecomplex **DM)
{
   Nx = nx;
   Ny = ny;
   Nz = nz;
   Npt = Nx*Ny*Nz;
   ndim = ndm;
   memory = new Memory();
 
   which = UseGamma = 0;
 
   data = DM;
   Dfdx = Dfdy = Dfdz = D2fdxdy = D2fdxdz = D2fdydz = D3fdxdydz = NULL;
   flag_reset_gamma = flag_allocated_dfs = 0;
 
   return;
}

/* ----------------------------------------------------------------------------
 * Private method to initialize tricubic interpolations
 * ---------------------------------------------------------------------------- */
void Interpolate::tricubic_init()
{
   // prepare necessary data for tricubic
   if (flag_allocated_dfs == 0){
      memory->create(Dfdx, Npt, ndim, "Interpolate_Interpolate:Dfdx");
      memory->create(Dfdy, Npt, ndim, "Interpolate_Interpolate:Dfdy");
      memory->create(Dfdz, Npt, ndim, "Interpolate_Interpolate:Dfdz");
      memory->create(D2fdxdy, Npt, ndim, "Interpolate_Interpolate:D2fdxdy");
      memory->create(D2fdxdz, Npt, ndim, "Interpolate_Interpolate:D2fdxdz");
      memory->create(D2fdydz, Npt, ndim, "Interpolate_Interpolate:D2fdydz");
      memory->create(D3fdxdydz, Npt, ndim, "Interpolate_Interpolate:D2fdxdydz");
 
      flag_allocated_dfs = 1;
   }

   // get the derivatives
   int n=0;
   const double half = 0.5, one4 = 0.25, one8 = 0.125;
   for (int ii = 0; ii < Nx; ++ii)
   for (int jj = 0; jj < Ny; ++jj)
   for (int kk = 0; kk < Nz; ++kk){
      int ip = (ii+1)%Nx, jp = (jj+1)%Ny, kp = (kk+1)%Nz;
      int im = (ii-1+Nx)%Nx, jm = (jj-1+Ny)%Ny, km = (kk-1+Nz)%Nz;
   
      int p100 = (ip*Ny+jj)*Nz+kk;
      int p010 = (ii*Ny+jp)*Nz+kk;
      int p001 = (ii*Ny+jj)*Nz+kp;
      int p110 = (ip*Ny+jp)*Nz+kk;
      int p101 = (ip*Ny+jj)*Nz+kp;
      int p011 = (ii*Ny+jp)*Nz+kp;
      int pm00 = (im*Ny+jj)*Nz+kk;
      int p0m0 = (ii*Ny+jm)*Nz+kk;
      int p00m = (ii*Ny+jj)*Nz+km;
      int pmm0 = (im*Ny+jm)*Nz+kk;
      int pm0m = (im*Ny+jj)*Nz+km;
      int p0mm = (ii*Ny+jm)*Nz+km;
      int p1m0 = (ip*Ny+jm)*Nz+kk;
      int p10m = (ip*Ny+jj)*Nz+km;
      int p01m = (ii*Ny+jp)*Nz+km;
      int pm10 = (im*Ny+jp)*Nz+kk;
      int pm01 = (im*Ny+jj)*Nz+kp;
      int p0m1 = (ii*Ny+jm)*Nz+kp;
      int p111 = (ip*Ny+jp)*Nz+kp;
      int pm11 = (im*Ny+jp)*Nz+kp;
      int p1m1 = (ip*Ny+jm)*Nz+kp;
      int p11m = (ip*Ny+jp)*Nz+km;
      int pm1m = (im*Ny+jp)*Nz+km;
      int p1mm = (ip*Ny+jm)*Nz+km;
      int pmm1 = (im*Ny+jm)*Nz+kp;
      int pmmm = (im*Ny+jm)*Nz+km;
   
      for (int idim=0; idim<ndim; idim++){
         Dfdx[n][idim].r = (data[p100][idim].r - data[pm00][idim].r) * half;
         Dfdx[n][idim].i = (data[p100][idim].i - data[pm00][idim].i) * half;
         Dfdy[n][idim].r = (data[p010][idim].r - data[p0m0][idim].r) * half;
         Dfdy[n][idim].i = (data[p010][idim].i - data[p0m0][idim].i) * half;
         Dfdz[n][idim].r = (data[p001][idim].r - data[p00m][idim].r) * half;
         Dfdz[n][idim].i = (data[p001][idim].i - data[p00m][idim].i) * half;
         D2fdxdy[n][idim].r = (data[p110][idim].r - data[p1m0][idim].r - data[pm10][idim].r + data[pmm0][idim].r) * one4;
         D2fdxdy[n][idim].i = (data[p110][idim].i - data[p1m0][idim].i - data[pm10][idim].i + data[pmm0][idim].i) * one4;
         D2fdxdz[n][idim].r = (data[p101][idim].r - data[p10m][idim].r - data[pm01][idim].r + data[pm0m][idim].r) * one4;
         D2fdxdz[n][idim].i = (data[p101][idim].i - data[p10m][idim].i - data[pm01][idim].i + data[pm0m][idim].i) * one4;
         D2fdydz[n][idim].r = (data[p011][idim].r - data[p01m][idim].r - data[p0m1][idim].r + data[p0mm][idim].r) * one4;
         D2fdydz[n][idim].i = (data[p011][idim].i - data[p01m][idim].i - data[p0m1][idim].i + data[p0mm][idim].i) * one4;
         D3fdxdydz[n][idim].r = (data[p111][idim].r-data[pm11][idim].r - data[p1m1][idim].r - data[p11m][idim].r +
                                 data[p1mm][idim].r+data[pm1m][idim].r + data[pmm1][idim].r - data[pmmm][idim].r) * one8;
         D3fdxdydz[n][idim].i = (data[p111][idim].i-data[pm11][idim].i - data[p1m1][idim].i - data[p11m][idim].i +
                                 data[p1mm][idim].i+data[pm1m][idim].i + data[pmm1][idim].i - data[pmmm][idim].i) * one8;
      }
      n++;
   }
   return;
}

/* ----------------------------------------------------------------------------
 * Deconstructor used to free memory
 * ---------------------------------------------------------------------------- */
Interpolate::~Interpolate()
{
   data = NULL;
   memory->destroy(Dfdx);
   memory->destroy(Dfdy);
   memory->destroy(Dfdz);
   memory->destroy(D2fdxdy);
   memory->destroy(D2fdxdz);
   memory->destroy(D2fdydz);
   memory->destroy(D3fdxdydz);
   delete memory;

   return;
}

/* ----------------------------------------------------------------------------
 * Tricubic interpolation, by calling the tricubic library
 * ---------------------------------------------------------------------------- */
void Interpolate::tricubic(double *qin, doublecomplex *DMq)
{
   // qin should be in unit of 2*pi/L
   double q[3];
   for (int i = 0; i < 3; ++i) q[i] = qin[i];
   for (int i = 0; i < 3; ++i){
      while (q[i] < 0.)  q[i] += 1.;
      while (q[i] >= 1.) q[i] -= 1.;
   }
   
   int ix = int(q[0]*double(Nx));
   int iy = int(q[1]*double(Ny));
   int iz = int(q[2]*double(Nz));
   double x = q[0]*double(Nx)-double(ix);
   double y = q[1]*double(Ny)-double(iy);
   double z = q[2]*double(Nz)-double(iz);
   int ixp = (ix+1)%Nx, iyp = (iy+1)%Ny, izp = (iz+1)%Nz;
   vidx[0] = (ix*Ny+iy)*Nz+iz;
   vidx[1] = (ixp*Ny+iy)*Nz+iz;
   vidx[2] = (ix*Ny+iyp)*Nz+iz;
   vidx[3] = (ixp*Ny+iyp)*Nz+iz;
   vidx[4] = (ix*Ny+iy)*Nz+izp;
   vidx[5] = (ixp*Ny+iy)*Nz+izp;
   vidx[6] = (ix*Ny+iyp)*Nz+izp;
   vidx[7] = (ixp*Ny+iyp)*Nz+izp;
   for (int i=0; i<8; i++) if (vidx[i] == 0) UseGamma = 1;
 
   for (int idim = 0; idim < ndim; ++idim){
      for (int i = 0; i < 8; ++i){
         f[i] = data[vidx[i]][idim].r;
         dfdx[i] = Dfdx[vidx[i]][idim].r;
         dfdy[i] = Dfdy[vidx[i]][idim].r;
         dfdz[i] = Dfdz[vidx[i]][idim].r;
         d2fdxdy[i] = D2fdxdy[vidx[i]][idim].r;
         d2fdxdz[i] = D2fdxdz[vidx[i]][idim].r;
         d2fdydz[i] = D2fdydz[vidx[i]][idim].r;
         d3fdxdydz[i] = D3fdxdydz[vidx[i]][idim].r;
      }
      tricubic_get_coeff(&a[0],&f[0],&dfdx[0],&dfdy[0],&dfdz[0],&d2fdxdy[0],&d2fdxdz[0],&d2fdydz[0],&d3fdxdydz[0]); 
      DMq[idim].r = tricubic_eval(&a[0],x,y,z);
     
      for (int i = 0; i < 8; ++i){
         f[i] = data[vidx[i]][idim].i;
         dfdx[i] = Dfdx[vidx[i]][idim].i;
         dfdy[i] = Dfdy[vidx[i]][idim].i;
         dfdz[i] = Dfdz[vidx[i]][idim].i;
         d2fdxdy[i] = D2fdxdy[vidx[i]][idim].i;
         d2fdxdz[i] = D2fdxdz[vidx[i]][idim].i;
         d2fdydz[i] = D2fdydz[vidx[i]][idim].i;
         d3fdxdydz[i] = D3fdxdydz[vidx[i]][idim].i;
      }
      tricubic_get_coeff(&a[0],&f[0],&dfdx[0],&dfdy[0],&dfdz[0],&d2fdxdy[0],&d2fdxdz[0],&d2fdydz[0],&d3fdxdydz[0]); 
      DMq[idim].i = tricubic_eval(&a[0],x,y,z);
   }
   return;
}

/* ----------------------------------------------------------------------------
 * method to interpolate the DM at an arbitrary q point;
 * the input q should be a vector in unit of (2pi/a 2pi/b 2pi/c).
 * All q components will be rescaled into [0 1).
 * ---------------------------------------------------------------------------- */
void Interpolate::trilinear(double *qin, doublecomplex *DMq)
{
   // rescale q[i] into [0 1)
   double q[3];
   for (int i = 0; i < 3; ++i) q[i] = qin[i];
   for (int i = 0; i < 3; ++i){
      while (q[i] < 0.)  q[i] += 1.;
      while (q[i] >= 1.) q[i] -= 1.;
   }
 
   // find the index of the eight vertice
   int ix, iy, iz, ixp, iyp, izp;
   double x, y, z;
   q[0] *= double(Nx);
   q[1] *= double(Ny);
   q[2] *= double(Nz);
 
   ix = int(q[0])%Nx;
   iy = int(q[1])%Ny;
   iz = int(q[2])%Nz;
   ixp = (ix+1)%Nx;
   iyp = (iy+1)%Ny;
   izp = (iz+1)%Nz;
   x = q[0] - double(ix);
   y = q[1] - double(iy);
   z = q[2] - double(iz);
 
   //--------------------------------------
   vidx[0] = ((ix*Ny)+iy)*Nz + iz;
   vidx[1] = ((ixp*Ny)+iy)*Nz + iz;
   vidx[2] = ((ix*Ny)+iyp)*Nz + iz;
   vidx[3] = ((ix*Ny)+iy)*Nz + izp;
   vidx[4] = ((ixp*Ny)+iy)*Nz + izp;
   vidx[5] = ((ix*Ny)+iyp)*Nz + izp;
   vidx[6] = ((ixp*Ny)+iyp)*Nz + iz;
   vidx[7] = ((ixp*Ny)+iyp)*Nz + izp;
   for (int i = 0; i < 8; ++i) if (vidx[i] == 0) UseGamma = 1;
 
   double fac[8];
   fac[0] = (1.-x)*(1.-y)*(1.-z);
   fac[1] = x*(1.-y)*(1.-z);
   fac[2] = (1.-x)*y*(1.-z);
   fac[3] = (1.-x)*(1.-y)*z;
   fac[4] = x*(1.-y)*z;
   fac[5] = (1.-x)*y*z;
   fac[6] = x*y*(1.-z);
   fac[7] = x*y*z;
   
   // now to do the interpolation
   for (int idim = 0; idim < ndim; ++idim){
      DMq[idim].r = 0.;
      DMq[idim].i = 0.;
      for (int i = 0; i < 8; ++i){
         DMq[idim].r += data[vidx[i]][idim].r*fac[i];
         DMq[idim].i += data[vidx[i]][idim].i*fac[i];
      }
   }
 
   return;
}

/* ----------------------------------------------------------------------------
 * To invoke the interpolation
 * ---------------------------------------------------------------------------- */
void Interpolate::execute(double *qin, doublecomplex *DMq)
{
   UseGamma = 0;
   if (which == 1) // 1: tricubic
      tricubic(qin, DMq);
   else       // otherwise: trilinear
      trilinear(qin, DMq);

   return;
}

/* ----------------------------------------------------------------------------
 * Public method, to set/reset the interpolation method
 * ---------------------------------------------------------------------------- */
void Interpolate::set_method()
{
   char str[MAXLINE];
   int im = 1;
   printf("\n");for(int i=0; i<80; i++) printf("=");
   printf("\nWhich interpolation method would you like to use?\n");
   printf("  1. Tricubic;\n  2. Trilinear;\n");
   printf("Your choice [1]: ");
   fgets(str,MAXLINE,stdin);
   char *ptr = strtok(str," \t\n\r\f");
   if (ptr) im = atoi(ptr);
 
   which =2-im%2;
   printf("Your  selection: %d\n", which);
   for(int i=0; i<80; i++) printf("="); printf("\n\n");
 
   if (which == 1) tricubic_init();
 
   return;
}

/* ----------------------------------------------------------------------------
 * Public method, to reset gamma point data; in this case, the gamma point data
 * will be meaningless. should only be called once.
 * ---------------------------------------------------------------------------- */
void Interpolate::reset_gamma()
{
   if (flag_reset_gamma) return;
   flag_reset_gamma = 1;
 
   int p1 = 1%Nx, p2 = 2%Nx;
   int m1 = (Nx-1), m2 = (Nx-2+Nx)%Nx;
 
   int ip1 = p1*Ny*Nz, ip2 = p2*Ny*Nz;
   int im1 = m1*Ny*Nz, im2 = m2*Ny*Nz;
 
   double const one6 = -1./6., two3 = 2./3.;
 
   for (int idim=0; idim < ndim; idim++){
      data[0][idim].i = 0.;
      data[0][idim].r = (data[im2][idim].r + data[ip2][idim].r) * one6
                      + (data[im1][idim].r + data[ip1][idim].r) * two3;
   }
 
   return;
}
/* ---------------------------------------------------------------------------- */