/* FFT interpolation of 3D real grid data */

/* Also trilinear and tricubic interpolation */

/* Copyright (c) 2014-2022 MJ Rutter 
 * 
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3
 * of the Licence, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see http://www.gnu.org/licenses/
 */ 

#include<stdio.h>
#include<stdlib.h>
#include<math.h>

#include "c2xsf.h"

void fft3d(double *c, int *ngptar, int dir);
void pad_recip(double *o, int fft[3], double **ptr, int nfft[3]);
double interpolate0d_fft(struct grid *gptr, double x_in[3], int comp,
			 double *rec);

/* Do FFT interpolation of real space real data */
void interpolate3d(struct grid *old_grid, struct grid *new_grid){
  int ffft[3],fft[3],nfft[3];
  int i,ic;
  long  old_size,new_size;
  double *o,*n,scale;

  /* A new grid dimension of zero means leave as was */
  for(i=0;i<3;i++)
    if (new_grid->size[i]==0) new_grid->size[i]=old_grid->size[i];

  for(i=0;i<3;i++) fft[i]=old_grid->size[i];
  for(i=0;i<3;i++) nfft[i]=new_grid->size[i];

  if (debug>1)
    fprintf(stderr,"Interpolating real data from %dx%dx%d to %dx%dx%d\n",
            fft[0],fft[1],fft[2],nfft[0],nfft[1],nfft[2]);

  old_size=fft[0]*fft[1]*fft[2];
  new_size=nfft[0]*nfft[1]*nfft[2];
  new_grid->data=malloc(old_grid->comps*new_size*sizeof(double));
  if (!new_grid->data) error_exit("Malloc error for final grid in interpolate");
  new_grid->comps=old_grid->comps;
  
  if((fft[0]==nfft[0])&&(fft[1]==nfft[1])&&(fft[2]==nfft[2])){
    if (debug>1) fprintf(stderr,"Null interpolation reduced to copy.\n");
    for(i=0;i<old_grid->comps*old_size;i++)
      new_grid->data[i]=old_grid->data[i];
    return;
  }

  /* Pad real data to complex */

  o=malloc(2*old_size*sizeof(double));
  if (!o) error_exit("Malloc error for first grid in interpolate");

  for(ic=0;ic<old_grid->comps;ic++){
    for(i=0;i<old_size;i++){
      o[2*i]=old_grid->data[i+ic*old_size];
      o[2*i+1]=0.0;
    }

  /* FFT to reciprocal space */

  /* A FORTRAN data order ... */

    ffft[0]=fft[2];
    ffft[1]=fft[1];
    ffft[2]=fft[0];

    if (debug>1) fprintf(stderr,"first FFT in interpolate\n");

    fft3d(o,ffft,-1);

  /* Pad onto interpolated reciprocal space grid */

  /* Assume all bits zero is a double zero */

    if (debug>1) fprintf(stderr,"padding in interpolate\n");

    pad_recip(o,fft,&n,nfft);

  /* FFT back to real space */

    if (debug>1) fprintf(stderr,"second FFT in interpolate\n");

    ffft[0]=nfft[2];
    ffft[1]=nfft[1];
    ffft[2]=nfft[0];

    fft3d(n,ffft,1);

    if (debug>1) fprintf(stderr,"end of second FFT in interpolate\n");

  /* Convert back to real and rescale */

    scale=1.0/old_size;

    for(i=0;i<new_size;i++)
      new_grid->data[i+ic*new_size]=scale*n[2*i];
    
    free(n);

  }
  free(o);
  
}

/* cubic interpolation of four datapoints */
/* p(-1), p(0), p(1), p(2).  0<=x<=1 */ 
static inline double cubic_int(double p[4], double x){
  return 0.5*(((-p[0]+3*p[1]-3*p[2]+p[3])*x+
	       (2*p[0]-5*p[1]+4*p[2]-p[3]))*x+
	      (-p[0]+p[2]))*x+p[1];
}

/* Tricubic interpolation to a point given in fractional co-ords */
double tricubic0d(struct grid *gptr, double x_in[3], int comp){
  int i,j,k,n[3],nx,ny[4],nz[4],ngx,ngy,ngz,offset;
  double x[3],g0[4][4][4],g1[4][4],g2[4],*data;

  for(i=0;i<3;i++){
    x[i]=fmod(x_in[i],1.0);
    if (x[i]<0) x[i]+=1;
    x[i]=x[i]*gptr->size[i];
    n[i]=(int)x[i];
    x[i]=x[i]-n[i];
    /* leave n[i] pointing at lowest co-ord to consider */
    n[i]=(n[i]+gptr->size[i]-1)%gptr->size[i];
  }

  ngx=gptr->size[0];
  ngy=gptr->size[1];
  ngz=gptr->size[2];
  data=gptr->data+(comp-1)*ngx*ngy*ngz;

  /* pre-compute, as % is expensive on some machines */
  ny[0]=n[1];
  if (n[1]+3<ngy){
    for(i=1;i<4;i++)
      ny[i]=n[1]+i;
  }
  else{
    for(i=1;i<4;i++)
      ny[i]=(n[1]+i)%ngy;
  }
  
  if (n[2]+3>=ngz){  /* This code is readable and general */
      //  if (1){
  /* Fill g0 with surrounding 64 points */
    nz[0]=n[2];
    for(i=1;i<4;i++)
      nz[i]=(n[2]+i)%ngz;
    for(i=0;i<4;i++){
      nx=(n[0]+i)%ngx;
      for(j=0;j<4;j++){
	for(k=0;k<4;k++){
	  g0[i][j][k]=data[nz[k]+ngz*(ny[j]+ngy*nx)];
	}
      }
    }

    /* remove z */
  
    for(i=0;i<4;i++){
      for(j=0;j<4;j++){
	//      for(k=0;k<4;k++)
	//	p[k]=g0[i][j][k];
	//      g1[i][j]=cubic_int(p,x[2]);
	g1[i][j]=cubic_int(g0[i][j],x[2]);
      }
    }
  }
  else{ /* This collapsing of the above code is less readable, but faster */
    for(i=0;i<4;i++){
      nx=(n[0]+i)%ngx;
      offset=nx*ngy*ngz+n[2];
      for(j=0;j<4;j++)
      	g1[i][j]=cubic_int(data+offset+ngz*ny[j],x[2]);
    }
  }
    

  /* remove y */

  for(i=0;i<4;i++){
    //    for(k=0;k<4;k++)
    //      p[k]=g1[i][k];
    //    g2[i]=cubic_int(p,x[1]);
    g2[i]=cubic_int(g1[i],x[1]);
  }

  return(cubic_int(g2,x[0]));
}

/* Trilinear interpolation to a point given in fractional co-ords */
double interpolate0d(struct grid *gptr, double x_in[3], int comp){
  int i,n[3],nx,ny,nz,plus_1,ngx,ngy,ngz,offset;
  double x[3],g[2],g1,g2,*data;

  if ((comp<=0)||(comp>gptr->comps)){
    fprintf(stderr,"Invalid value of comp in interpolate0d. "
	    "Have %d, valid 1 to %d\n",comp,gptr->comps);
    exit(1);
  }

  if (flags&FFT) return interpolate0d_fft(gptr,x_in,comp,NULL);
  
  if (flags&TRICUBIC) return tricubic0d(gptr,x_in,comp);

  for(i=0;i<3;i++){
    x[i]=fmod(x_in[i],1.0);
    if (x[i]<0) x[i]+=1;
    x[i]=x[i]*gptr->size[i];
    n[i]=(int)x[i];
    x[i]=x[i]-n[i];
  }

  ngx=gptr->size[0];
  ngy=gptr->size[1];
  ngz=gptr->size[2];

  data=gptr->data+(comp-1)*ngx*ngy*ngz;
  
  nz=n[2];
  plus_1=1;
  if (nz==ngz-1) plus_1=1-ngz;
  for(i=0;i<2;i++){
      nx=(n[0]+i)%ngx;
      ny=n[1];
      offset=nx*ngy*ngz+ny*ngz+nz;
      g1=data[offset]+x[2]*(data[offset+plus_1]-data[offset]);
      ny=(ny+1)%ngy;
      offset=nx*ngy*ngz+ny*ngz+nz;
      g2=data[offset]+x[2]*(data[offset+plus_1]-data[offset]);
      g[i]=g1+x[1]*(g2-g1);
  }   

  return(g[0]+x[0]*(g[1]-g[0]));
}

void vinterpolate0d(struct grid *gptr, double x_in[3], double *z){
  int i;

  for(i=0;i<gptr->comps;i++)
    z[i]=interpolate0d(gptr,x_in,i+1);
  
}

void interpolate1d(struct grid *gptr, double st[3], double end[3],
                   int npts, double *out){
  int i,j,nn[3];
  double x[3],*rec,*ptr;

  rec=NULL;
  if (flags&FFT){
    rec=malloc(2*npts*sizeof(double));
    if (!rec) error_exit("malloc error in interpolate1d");

    ptr=rec;
    for(i=0;i<npts;i++){
      *(ptr++)=gptr->data[i];
      *(ptr++)=0;
    }

    nn[0]=gptr->size[2];
    nn[1]=gptr->size[1];
    nn[2]=gptr->size[0];
    fft3d(rec,nn,-1);
  }
  
  for(i=0;i<npts;i++){
    for(j=0;j<3;j++)
      x[j]=st[j]+(i/(double)(npts-1))*(end[j]-st[j]);
    if (flags&FFT)
      out[i]=interpolate0d_fft(gptr,x,1,rec);
    else
      out[i]=interpolate0d(gptr,x,1);
  }

  if (rec) free(rec);
}

/* Pad into array assumed to be zeroed */
/* Deal with case of target being both larger and smaller than source */
void pad_recip(double *o, int fft[3], double **nptr, int nfft[3]){
  int i,j,k,ii,jj,kk,nii,njj,nkk;
  int imin,imax,jmin,jmax,kmin,kmax;
  int ind,nind;
  double *n;

  for(i=0;i<3;i++)
    if (nfft[i]==0) nfft[i]=fft[i];

  for(i=0;i<3;i++)
    if (nfft[i]<1) {
      fprintf(stderr,"Invalid grid size %dx%dx%d\n",nfft[0],nfft[1],nfft[2]);
      exit(1);
    }

  if (debug) fprintf(stderr,"Moving from %dx%dx%d to %dx%dx%d recip grid\n",
                     fft[0],fft[1],fft[2],nfft[0],nfft[1],nfft[2]);

  *nptr=calloc(nfft[0]*nfft[1]*nfft[2],2*sizeof(double));
  if (!*nptr) error_exit("Malloc error for interpolated grid\n");
  n=*nptr;

  imax=min(fft[0]/2,nfft[0]/2);
  imin=max((1-fft[0])/2,(1-nfft[0])/2);

  jmax=min( fft[1]/2,nfft[1]/2);
  jmin=max((1-fft[1])/2,(1-nfft[1])/2);

  kmax=min(fft[2]/2,nfft[2]/2);
  kmin=max((1-fft[2])/2,(1-nfft[2])/2);
  
  for(i=imin;i<=imax;i++){
    if (i>=0) {
      ii=i;
      nii=i;
    }
    else{
      ii=fft[0]+i;
      nii=nfft[0]+i;
    }
    for(j=jmin;j<=jmax;j++){
      if (j>=0) {
        jj=j;
        njj=j;
      }
      else{
        jj=fft[1]+j;
        njj=nfft[1]+j;
      }
      for(k=kmin;k<=kmax;k++){
        if (k>=0) {
          kk=k;
          nkk=k;
        }
        else{
          kk=fft[2]+k;
          nkk=nfft[2]+k;
        }
        ind=2*(kk+fft[2]*(jj+ii*fft[1]));
        nind=2*(nkk+nfft[2]*(njj+nii*nfft[1]));
        n[nind]=o[ind];
        n[nind+1]=o[ind+1];
      }
    }
  }
}

double interpolate0d_fft(struct grid *gptr, double x_in[3], int comp,
			 double *rec_in){
  int i,j,k,ii,jj,kk,npts,nn[3],off;
  double *rec,*ptr,x[3],z,ph;

  if (!gptr->data) error_exit("no data for interpolation");

  npts=gptr->size[0]*gptr->size[1]*gptr->size[2];

  rec=rec_in;
  if (!rec){
    rec=malloc(2*npts*sizeof(double));
    if (!rec) error_exit("malloc error in interpolate0d_fft");

    ptr=rec;
    for(i=0;i<npts;i++){
      *(ptr++)=gptr->data[i+(comp-1)*npts];
      *(ptr++)=0;
    }

    nn[0]=gptr->size[2];
    nn[1]=gptr->size[1];
    nn[2]=gptr->size[0];
    fft3d(rec,nn,-1);
  }
  
  for(i=0;i<3;i++) x[i]=2*M_PI*x_in[i];
  z=0;
  for(i=0;i<gptr->size[0];i++){
    ii=i;
    if (i>gptr->size[0]/2) ii-=gptr->size[0];
    for(j=0;j<gptr->size[1];j++){
      jj=j;
      if (j>gptr->size[1]/2) jj-=gptr->size[1];
      for(k=0;k<gptr->size[2];k++){
	kk=k;
	if (k>gptr->size[2]/2) kk-=gptr->size[2];
	ph=ii*x[0]+jj*x[1]+kk*x[2];
	off=2*(k+gptr->size[2]*(j+gptr->size[1]*i));
	z+=cos(ph)*rec[off]-sin(ph)*rec[off+1];
      }
    }
  }

  if (!rec_in) free(rec);
  return(z/npts);
  
}