#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <time.h>
#include <complex.h>
#include <fftw3-mpi.h>

#include "PAN_FFTW3.h"
#include "panphasia_functions.h"

#ifdef USE_OPENMP
#include <omp.h>
#endif


extern const int Nbasis;
extern const int irank_p[3][84];

extern size_t descriptor_order;
extern size_t descriptor_kk_limit;
extern size_t descriptor_base_size;

int PANPHASIA_compute_kspace_field_(size_t relative_level, ptrdiff_t N0_fourier_grid,
				     ptrdiff_t local_n0_fourier_return, ptrdiff_t local_0_start_fourier_return,
                                     FFTW_COMPLEX *return_field)
{

size_t copy_list[Nbasis];




int pmax = 6;

 int nsubdivide = 21; //(pmax%2==0)?pmax+1:pmax+2;

size_t ncopy = (pmax+1)*(pmax+2)*(pmax+3)/6;

if (ncopy%nsubdivide!=0) return(100010);
int nchunk = ncopy/nsubdivide;

int verbose = 1;
int flag_output_mode=2;
int error;
ptrdiff_t size_to_alloc_fourier;
ptrdiff_t size_to_alloc_pan;
 ptrdiff_t local_fourier_x_start, local_fourier_x_end;
FFTW_PLAN output_coeff_forward_plan; 

ptrdiff_t  N0_pan_grid = descriptor_base_size<<relative_level;

if (N0_fourier_grid%N0_pan_grid!=0) return (100015);

int    fdim = N0_fourier_grid/N0_pan_grid;
size_t nfft_dim = N0_fourier_grid;
size_t npan_dim = N0_pan_grid;

int SHARED_FOUR_PAN_SPACE = (nsubdivide==1)&&(fdim==1)&&(sizeof(PAN_REAL)==sizeof(FFTW_REAL));



////////////////////////////////////////////////////////////////////////////////////

if (pmax>descriptor_order) return(100020);

for (size_t i=0; i<Nbasis; i++) copy_list[i]=i;

//printf("Dimensions of FT   (%td,%td,%td)\n",N0_fourier_grid,N0_fourier_grid,N0_fourier_grid);
//printf("Dimensions of PG   (%td,%td,%td)\n",N0_pan_grid,N0_pan_grid,N0_pan_grid);
//printf("local_no %td local_0_start_fourier %td\n",local_n0_fourier_return, local_0_start_fourier_return);

//  Compute 1-D Spherical Bessel coefficients for each order //////////////////
//  These are needed for the convolutions below              //////////////////
size_t  n4dimen;



n4dimen=(nfft_dim%4==0) ? 4*(nfft_dim/4)+4 : 4*(nfft_dim/4)+5; 

double complex *sph_bessel_coeff = FFTW_MALLOC(sizeof(double complex)*n4dimen*(pmax+1)); 

if (sph_bessel_coeff==NULL) return(100030);

compute_sph_bessel_coeffs(nfft_dim, pmax, n4dimen, fdim, sph_bessel_coeff);

//printf("Reached here! ndimen4 %ld\n",n4dimen);
///////////////////////////////////////////////////////////////////////////////




//////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////
// Determine sizes of Fourier and Panphasia coefficient 3-D arrays //


ptrdiff_t local_n0_pan;
ptrdiff_t local_0_start_pan;
{
 
int rank =3;

ptrdiff_t local_n0_fourier;
ptrdiff_t local_0_start_fourier;
const ptrdiff_t ndimens_alloc_fourier[] = {N0_fourier_grid, N0_fourier_grid, N0_fourier_grid+2}; // Allocated for r2c
ptrdiff_t howmany = nchunk;
size_to_alloc_fourier = FFTW_MPI_LOCAL_SIZE_MANY(rank, ndimens_alloc_fourier, howmany,
 						    FFTW_MPI_DEFAULT_BLOCK,MPI_COMM_WORLD,
                                                    &local_n0_fourier,&local_0_start_fourier);
 
 if (local_0_start_fourier!=local_0_start_fourier_return){
   printf("Error local_0_start_fourier!=local_0_start_fourier_return\n");
   return(100032);
 };

if (local_n0_fourier!=local_n0_fourier_return){
   printf("Error local_n0_fourier!=local_n0_fourier_return\n");
   return(100033);
 };


 ptrdiff_t abs_fourier_x_start, abs_fourier_x_end;

 if (local_0_start_fourier_return%fdim==0){
   abs_fourier_x_start = local_0_start_fourier_return;
 }else{
   abs_fourier_x_start = local_0_start_fourier_return + fdim - (local_0_start_fourier_return%fdim);
 };


 if ((local_0_start_fourier_return + local_n0_fourier_return - 1)%fdim==0){
   abs_fourier_x_end = local_0_start_fourier_return + local_n0_fourier_return - 1;
 }else{
   abs_fourier_x_end = (local_0_start_fourier_return + local_n0_fourier_return - 1)-
     ((local_0_start_fourier_return + local_n0_fourier_return - 1)%fdim);
 };

 if ((abs_fourier_x_end-abs_fourier_x_start)%fdim!=0) return (100036); 

 local_0_start_pan = abs_fourier_x_start/fdim;
 local_n0_pan      = 1 + (abs_fourier_x_end-abs_fourier_x_start)/fdim;

 local_fourier_x_start = abs_fourier_x_start - local_0_start_fourier_return;
 local_fourier_x_end   = abs_fourier_x_end   - local_0_start_fourier_return;

const ptrdiff_t ndimens_alloc_pan[] = {N0_pan_grid, N0_pan_grid, N0_pan_grid+2}; 
 howmany = ncopy;

size_to_alloc_pan = howmany * local_n0_pan * N0_pan_grid * ( N0_pan_grid+2);


//printf("size_to_alloc_fourier = %td\n",size_to_alloc_fourier);
//printf("size_to_alloc_pan     = %td\n",size_to_alloc_pan);

 
 };
 /////////////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////////////





///////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////
 // Allocate arrays to store Panphasia coefficients and Fourier information
 // If nsubdivide==1 then use the same structure to store both.

 void *panphasia_coefficients= FFTW_MALLOC(sizeof(PAN_REAL)*size_to_alloc_pan);
 void *fourier_grids;

 if (panphasia_coefficients==NULL) return(100040);

 void *mode_weightings = FFTW_MALLOC(sizeof(FFTW_REAL)*size_to_alloc_fourier/nchunk);

 FFTW_REAL *ptr_mode_weightings;
 ptr_mode_weightings = mode_weightings;
 memset(ptr_mode_weightings, 0, sizeof(FFTW_REAL)*size_to_alloc_fourier/nchunk);


 FFTW_REAL     *ptr_real_fourier_grid;
 FFTW_REAL     *ptr_panphasia_coefficients = panphasia_coefficients;
 
 FFTW_COMPLEX *ptr_cmplx_fourier_grid;

 if (SHARED_FOUR_PAN_SPACE){
       ptr_real_fourier_grid  = panphasia_coefficients;
       ptr_cmplx_fourier_grid = panphasia_coefficients;

 }else{

    fourier_grids= FFTW_MALLOC(sizeof(FFTW_REAL)*size_to_alloc_fourier);
 
   if (fourier_grids==NULL) return(100041);

       ptr_real_fourier_grid  = fourier_grids;
       ptr_cmplx_fourier_grid = fourier_grids;
 };
///////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////




///////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////
//////////// Compute the Panphasia coefficients   ////////////////////////////////

{

size_t xorigin= local_0_start_pan, yorigin=0, zorigin=0;

size_t xextent =local_n0_pan, yextent = N0_pan_grid, zextent = N0_pan_grid;


 
if (error = PANPHASIA_compute_coefficients_(&xorigin,&yorigin,&zorigin,&xextent,&yextent,
		     &zextent, copy_list, &ncopy, panphasia_coefficients,
                    &flag_output_mode,&verbose)) return(100100+error);

 };
///////////////////////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////////////////////


 
//////////  Output diagnostics for small grids only

{

if (N0_pan_grid<-1){
 
int rank;
  MPI_Comm_rank(MPI_COMM_WORLD,&rank);
  char filename[100];
  sprintf(filename,"output_real_space_field.%d",rank);
  
  FILE *fp;
  printf("local_n0_pan %ld npan_dim %ld N0_pan_grid %ld\n",local_n0_pan,
	 npan_dim,N0_pan_grid);
 
  fp = fopen(filename,"w");

  for (int ix=0; ix<local_n0_pan; ix++)
    for (int iy=0; iy < npan_dim; iy++)
      for (int iz=0; iz < npan_dim; iz++){
	int index = ix*N0_pan_grid*(N0_pan_grid+2) + iy*(N0_pan_grid+2) + iz;
        fprintf(fp,"%6llu%6d%6d %14.8lf %d\n",ix+local_0_start_pan,iy,iz,
		ptr_panphasia_coefficients[index],index);
		      };

 fclose(fp);
 };
 };



//----------------------------------------------------------------------------------
////////////// Set up FTTW plan //////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////   

// printf("Making the plan ... \n");

//////////////////// Make plan for ncopy interleaved FTs ///////////////////////////
//////////////////////////////////////////////////////////////////////////   

{
  int rank = 3;
  const ptrdiff_t ndimens[3] =  {N0_fourier_grid, N0_fourier_grid, N0_fourier_grid};
  ptrdiff_t howmany = nchunk;
  ptrdiff_t block = FFTW_MPI_DEFAULT_BLOCK;
  ptrdiff_t tblock =  FFTW_MPI_DEFAULT_BLOCK;
  unsigned flags = FFTW_ESTIMATE;

   output_coeff_forward_plan = FFTW_MPI_PLAN_MANY_DTF_R2C(rank, ndimens,
                                 howmany, block, tblock,
                                 ptr_real_fourier_grid, ptr_cmplx_fourier_grid,
                                 MPI_COMM_WORLD, flags);
   if (output_coeff_forward_plan==NULL) {
   printf("Null plan\n");
   return(100051);
   };

 };

//printf("Plan completed ... \n");

//////////////////////////////////////////////////////////////////////////   
//////////////////////////////////////////////////////////////////////////   
//----------------------------------------------------------------------------------   


memset(return_field, 0, local_n0_fourier_return*N0_fourier_grid *(N0_fourier_grid +2) * sizeof(FFTW_REAL));
 

for (int iter = 0;  iter < nsubdivide; iter++){
 
int moffset = iter*nchunk;



if (!SHARED_FOUR_PAN_SPACE){

  memset(ptr_real_fourier_grid, 0, sizeof(FFTW_REAL)*size_to_alloc_fourier);

  // Copy Panphasia coefficients to Fourier grid with appropriate stride - fdim

size_t index_p,index_f;
int m;

 for (int ix_f = local_fourier_x_start, ix_p=0; ix_f <=local_fourier_x_end; ix_f+=fdim,ix_p++)
#ifdef USE_OPENMP
#pragma omp parallel for collapse(2) private (index_p,index_f,m)
#endif

   for(int iy_f=0, iy_p=0 ; iy_f<nfft_dim; iy_f+=fdim, iy_p++)
     for(int iz_f=0, iz_p=0; iz_f<nfft_dim; iz_f+=fdim, iz_p++){
	 index_p = ix_p*N0_pan_grid*(N0_pan_grid + 2) + iy_p*(N0_pan_grid + 2) + iz_p;
	 index_f = ix_f*N0_fourier_grid*(N0_fourier_grid + 2) + iy_f*(N0_fourier_grid + 2) + iz_f;

          for (m=0; m<nchunk; m++){
          ptr_real_fourier_grid[nchunk*index_f + m] =
	    ptr_panphasia_coefficients[ncopy*index_p + moffset + m];

	  };
       };
 };

 



FFTW_MPI_EXECUTE_DFT_R2C(output_coeff_forward_plan,
                         ptr_real_fourier_grid,ptr_cmplx_fourier_grid);


 

{    
  // Convolve and combine the FT of the Panphasia coefficient field

size_t index1,index2;
complex weight;
int m;


#ifdef USE_OPENMP
#pragma omp parallel for collapse(3)  \
  private (index1,index2,weight,m)
#endif
  for(int ix=0;ix<local_n0_fourier_return;ix++)
     for(int iy=0;iy<nfft_dim;iy++)
        for(int iz=0;iz<=nfft_dim/2;iz++){
	  index1 = ix*N0_fourier_grid*(N0_fourier_grid/2+1) + iy*(N0_fourier_grid/2+1) + iz;
        for (int m=0; m<nchunk; m++){
        index2 = nchunk*index1 + m;

	
        weight = sph_bessel_coeff[n4dimen*irank_p[0][m+moffset]+ix+local_0_start_fourier_return]*
	      sph_bessel_coeff[n4dimen*irank_p[1][m+moffset]+iy]*
	                                sph_bessel_coeff[n4dimen*irank_p[2][m+moffset]+iz];
        return_field[index1] += weight * ptr_cmplx_fourier_grid[index2];
        ptr_mode_weightings[index1]+=cabs(weight)*cabs(weight);
        };
	};
 };
 

 }; // End loop over iter


 
// printf("Reached here 10!\n");

   //Add phase shift and normalise field
{

   double complex phase_shift_and_scale;
   int  kx,ky,kz;
   const double pi = 4.0 * atan(1.0);
   size_t index1;

   double min_weight = 100.0;

#ifdef USE_OPENMP
#pragma omp parallel for collapse(3)  \
  private (index1,kx,ky,kz,phase_shift_and_scale)
#endif
  for(int ix=0;ix<local_n0_fourier_return;ix++)
    for(int iy=0;iy<nfft_dim;iy++)
      for(int iz=0;iz<=nfft_dim/2;iz++){
	index1 = ix*N0_fourier_grid*(N0_fourier_grid/2+1) + iy*(N0_fourier_grid/2+1) + iz;
        kx = (ix+local_0_start_fourier_return>nfft_dim/2) ? 
              ix + local_0_start_fourier_return - nfft_dim : ix + local_0_start_fourier_return;
        ky = (iy > nfft_dim/2) ? iy-nfft_dim : iy;
        kz = (iz > nfft_dim/2) ? iz-nfft_dim : iz;

        if ( (kx==nfft_dim/2)||(ky==nfft_dim/2)||(kz==nfft_dim/2)){
	  // Set Nyquist modes to zero - not used by IC_Gen anyway.
          phase_shift_and_scale = 0.0; //1.0/pow((double)nfft_dim,1.5);  // No phase shift
          ptr_mode_weightings[index1] = 0.0; // Set squared weight to zero
        }else{
	  phase_shift_and_scale = sqrt( (double)(fdim*fdim*fdim))*
              cexp( (double)fdim * (-I)*pi*(double)(kx + ky + kz)/sqrt(ptr_mode_weightings[index1])/
                          (double)nfft_dim)/pow((double)nfft_dim,1.5);
        };

	return_field[index1] *= phase_shift_and_scale;
        if (ptr_mode_weightings[index1]<min_weight) 
              min_weight=ptr_mode_weightings[index1];
     
      };

  //  printf("Minimum weight %lf\n",sqrt(min_weight));


 };

// printf("Reached here 11!\n");



 // Rescale selected Fourier modes to unit amplitude.
 // By default this part is not executed. 

 if (descriptor_kk_limit>0){
   size_t index1;
   complex weight;
   size_t ksquared;
   int kx,ky,kz;
#ifdef USE_OPENMP
#pragma omp parallel for collapse(3)  \
  private (index1,kx,ky,kz,ksquared,weight)
#endif
  for(int ix=0;ix<local_n0_fourier_return;ix++)
     for(int iy=0;iy<nfft_dim;iy++)
        for(int iz=0;iz<=nfft_dim/2;iz++){
          kx = (ix+local_0_start_fourier_return>nfft_dim/2) ? 
              ix + local_0_start_fourier_return - nfft_dim : ix + local_0_start_fourier_return;
          ky = (iy > nfft_dim/2) ? iy-nfft_dim : iy;
          kz = (iz > nfft_dim/2) ? iz-nfft_dim : iz;
          ksquared = kx*kx + ky*ky + kz*kz;
          if ( (kx!=nfft_dim/2)&&(ky!=nfft_dim/2)&&(kz!=nfft_dim/2)){ //Omit Nyquist modes

          if ((ksquared<=descriptor_kk_limit)&&(ksquared!=0)){
	    index1 = ix*N0_fourier_grid*(N0_fourier_grid/2+1) + iy*(N0_fourier_grid/2+1) + iz;
            weight = cabs(return_field[index1]);
            return_field[index1] /= weight;   
	  }; 	
	};
	};
 };
 
 //printf("Reached here 12!\n");


int rank;
  MPI_Comm_rank(MPI_COMM_WORLD,&rank);
  char filename[100];
  sprintf(filename,"output_k_space_alt.%d",rank);
  FILE *fp;



if (nfft_dim <-1){


  
  FILE *fp;

  fp = fopen(filename,"w");
  
  for (int ix=0; ix<local_n0_fourier_return; ix++)
    for (int iy=0; iy < nfft_dim; iy++)
      for (int iz=0; iz <= nfft_dim/2; iz++){

	int index = ix*N0_fourier_grid*(N0_fourier_grid/2+1) + iy*(N0_fourier_grid/2+1) + iz;
        fprintf(fp,"%6llu%6d%6d %14.8lf %14.8lf %14.8lf \n",ix+local_0_start_fourier_return,iy,iz,
		creal(return_field[index]),cimag(return_field[index]),sqrt(ptr_mode_weightings[index]));
		//		ptr_mode_weightings[index]);
		      };
  fclose(fp);



 }else{

  fp = fopen(filename,"w");
  
  for (int ix=0; ix<local_n0_fourier_return; ix++)
    for (int iy=0; iy < nfft_dim; iy++)
      for (int iz=0; iz <= nfft_dim/2; iz++){
        if (ix+iy+iz+local_0_start_fourier_return<100){
	int index = ix*N0_fourier_grid*(N0_fourier_grid/2+1) + iy*(N0_fourier_grid/2+1) + iz;
        fprintf(fp,"%6llu%6d%6d %14.8lf %14.8lf %14.8lf \n",ix+local_0_start_fourier_return,iy,iz,
		creal(return_field[index]),cimag(return_field[index]),sqrt(ptr_mode_weightings[index]));
		//   ptr_mode_weightings[index]);
		      };
      };
  fclose(fp);


 };


// Transpose output field

 {

   FFTW_PLAN transpose_output_plan;
   unsigned flags = FFTW_ESTIMATE;

   void *ptr_inter = return_field;

   FFTW_REAL  *ptr_return_as_real_field;
   ptr_return_as_real_field = ptr_inter;
   

   int rank = 2;

   ptrdiff_t size_to_transpose;
   ptrdiff_t howmany = N0_fourier_grid + 2;
   ptrdiff_t local_n0, local_0_start;
   ptrdiff_t local_n1, local_1_start;

   const ptrdiff_t ndimens[] = {N0_fourier_grid, N0_fourier_grid};


 size_to_transpose = FFTW_MPI_LOCAL_SIZE_MANY_TRANSPOSED(rank, ndimens, howmany, 
                 FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK, 
                     MPI_COMM_WORLD,&local_n0, &local_0_start,&local_n1, &local_1_start);

 // printf("size_to_transpose = %td\n",size_to_transpose);

  transpose_output_plan = FFTW_MPI_PLAN_MANY_TRANSPOSE(N0_fourier_grid, N0_fourier_grid,
                        howmany, FFTW_MPI_DEFAULT_BLOCK,FFTW_MPI_DEFAULT_BLOCK, 
		     ptr_return_as_real_field, ptr_return_as_real_field,
			MPI_COMM_WORLD, flags);

  //printf("Transpose plan completed.\n");

 FFTW_MPI_EXECUTE_R2R(transpose_output_plan,ptr_return_as_real_field,ptr_return_as_real_field);

 //printf("Transpose completed.\n");
 };








// Free all memory assigned by FFTW_MALLOC
 FFTW_FREE(mode_weightings);
FFTW_FREE(sph_bessel_coeff); 
FFTW_FREE(panphasia_coefficients);

if (!SHARED_FOUR_PAN_SPACE) FFTW_FREE(fourier_grids);

FFTW_DESTROY_PLAN(output_coeff_forward_plan);

//printf("Reached end of PANPHASIA_compute_kspace_field_\n");
return(0);


 };


//==========================================================================================
//==========================================================================================