#include <stdlib.h>
#include <math.h>
#include "geostat.h"


/*FAST FOURIER TRANSFORM MOVING AVERAGE METHOD */
/*Turns a Gaussian white noise vector into a   */
/*spatially correlated vector                  */
/*input:                                       */
/*variogram: structure defining the variogram  */
/*           model                             */
/*grid: structure defining the grid            */
/*n: vector with the number of cells along the */
/*   X, Y and Z axes for the underlying grid   */
/*   i = [0 1 2]                               */
/*   --> 0 0 0        : n will be computed and */
/*   updated as output                         */
/*   --> nx ny nz: these dimensions are used   */
/*realin: structure defining a realization -   */
/*        must be a Gaussian white noise       */
/*output:                                      */
/*realout: structure defining a realization -  */



void FFTMA(struct vario_mod variogram,struct grid_mod grid,int n[3],struct realization_mod *realin,struct realization_mod *realout)

{
	int NTOT,i,j,k,NMAX,NDIM,ntot,nmax,NXYZ,nxyz,maille0,maille1;
	double *table,*covar,*workr,*worki,*realization,temp;
	
	
	/*test over the input realization*/
	/*if ((*realin).code != 0) {
	  printf("Input realizations in FFTMA must be Gaussian white noises");
	  exit;
	  }*/
	
	/*covariance axis normalization*/
	axes(variogram.ap,variogram.scf,variogram.Nvario);
	
	
	/*pseudo-grid definition*/
	cgrid(variogram,grid,n);  
	
	
	/*constant definition*/
	NTOT = n[0]*n[1]*n[2];
	ntot = NTOT+1;
	NMAX = n[0];
	NDIM = 3;
	for (i=1;i<3;i++) {
		if (n[i] > NMAX)
			NMAX = n[i];
		if (n[i] == 1)
			NDIM--;
	}
	nmax = NMAX+1;
	NXYZ = grid.NX*grid.NY*grid.NZ;
	nxyz = NXYZ+1;
	
	
	/*array initialization*/
	covar = (double *) malloc(ntot * sizeof(double));
	if (covar == NULL) {
	  printf("FFTMA.c:  No memory available for covar");
	  exit;
	}

	table = (double *) malloc(ntot * sizeof(double));
	if (table == NULL) {
	  printf("FFTMA.c:  No memory available for table");
	  exit;
	}

	realization = (double *) malloc(ntot * sizeof(double));
	if (realization == NULL) {
	  printf("FFTMA.c:  No memory available for realization");
	  exit;
	}

	workr = (double *) malloc(nmax * sizeof(double));
	if (workr == NULL) {
	  printf("FFTMA.c:  No memory available for workr");
	  exit;
	}

	worki = (double *) malloc(nmax * sizeof(double));
	if (worki == NULL) {
	  printf("FFTMA.c:  No memory available for worki");
	  exit;
	}
	
	/*covariance function creation*/
	covariance(covar,variogram,grid,n);
	
	
	/*power spectrum*/
	fourt(covar,table,n,NDIM,1,0,workr,worki);
	
	
	/*organization of the input Gaussian white noise*/
	for ( k = 1; k <= n[2]; k++) {
	  for (j = 1; j <= n[1]; j++) {
	    for (i = 1; i <= n[0]; i++) {
	      maille1 = i+(j-1+(k-1)*n[1])*n[0];
	      if (i <= grid.NX && j <= grid.NY && k <= grid.NZ) {
		maille0 = i-1+(j-1+(k-1)*grid.NY)*grid.NX;
		realization[maille1] = (*realin).vector[maille0];
	      } else {
		realization[maille1] = 0.;
	      }
	    }
	  }
	}
	
	
	/*forward fourier transform of the GWN*/
	fourt(realization,table,n,NDIM,1,0,workr,worki);
	
	
	/*decomposition and multiplication in the spectral domain*/
	for ( k = 1; k <= n[2]; k++) {
	  for (j = 1; j <= n[1]; j++) {
	    for (i = 1; i <= n[0]; i++) {
	      maille1 = i+(j-1+(k-1)*n[1])*n[0];
	      temp = covar[maille1];
	      if (temp > 0.) {
		temp = sqrt(temp)/(double) NTOT;
	      } else if (temp < 0.) {
		temp = sqrt(-temp)/(double) NTOT;
	      }
	      realization[maille1] *= temp;
	      table[maille1] *= temp;
	    }
	  }
	}
	
	
	free(covar);
	
	
	/*backward fourier transform*/
	fourt(realization,table,n,NDIM,0,1,workr,worki);
	
	
	free(table);
	free(workr);
	free(worki);
	
	
	/*output realization*/
	
	/*is the output realization already allocated?*/
	/*if not, memory allocation*/
	if ((*realout).vector == NULL || (*realout).n != (*realin).n) {
	  (*realout).vector = (double *) malloc((*realin).n * sizeof(double));
	  if ((*realout).vector == NULL) {
	    printf("FFTMA.c:  No memory available");
	    exit;
	  }
	}
	(*realout).n = (*realin).n;
	(*realout).code = 1;
	for ( k = 1; k <= grid.NZ; k++) {
	  for (j = 1; j <= grid.NY; j++) {
	    for (i = 1; i <= grid.NX; i++) {
	      maille1 = i+(j-1+(k-1)*n[1])*n[0];
	      maille0 = i-1+(j-1+(k-1)*grid.NY)*grid.NX;
	      (*realout).vector[maille0] = realization[maille1];
	    }
	  }
	}
	
	
	free(realization);
	
	
	return;
}