/**************************************/
/*  Spectral method to calculate the  */
/*  scalar thermal conductivity and   */
/*  shear viscosity.                  */
/*                                    */
/*          Version 1.1 Jan.22, 1999  */
/*                       Ju Li (MIT)  */
/**************************************/

#include "smile.h"

static void realft(double data[], unsigned long n, int isign);
static void four1(double data[], unsigned long nn, int isign);

void realft (double data[], unsigned long n, int isign)
{
    void four1(double data[], unsigned long nn, int isign);
    unsigned long i,i1,i2,i3,i4,np3;
    double c1=0.5,c2,h1r,h1i,h2r,h2i;
    double wr,wi,wpr,wpi,wtemp,theta;

    theta=3.141592653589793/(n>>1);
    if (isign == 1) {
	c2 = -0.5;
	four1(data,n>>1,1);
    } else {
	c2=0.5;
	theta = -theta;
    }
    wtemp=sin(0.5*theta);
    wpr = -2.0*wtemp*wtemp;
    wpi=sin(theta);
    wr=1.0+wpr;
    wi=wpi;
    np3=n+3;
    for (i=2;i<=(n>>2);i++) {
	i4=1+(i3=np3-(i2=1+(i1=i+i-1)));
	h1r=c1*(data[i1]+data[i3]);
	h1i=c1*(data[i2]-data[i4]);
	h2r = -c2*(data[i2]+data[i4]);
	h2i=c2*(data[i1]-data[i3]);
	data[i1]=h1r+wr*h2r-wi*h2i;
	data[i2]=h1i+wr*h2i+wi*h2r;
	data[i3]=h1r-wr*h2r+wi*h2i;
	data[i4] = -h1i+wr*h2i+wi*h2r;
	wr=(wtemp=wr)*wpr-wi*wpi+wr;
	wi=wi*wpr+wtemp*wpi+wi;
    }
    if (isign == 1) {
	data[1] = (h1r=data[1])+data[2];
	data[2] = h1r-data[2];
    } else {
	data[1]=c1*((h1r=data[1])+data[2]);
	data[2]=c1*(h1r-data[2]);
	four1(data,n>>1,-1);
    }
    /* Added by Ju Li: factor of 2/n is */
    /* multiplied for inverse realft(); */
    if (isign == -1)
	for (i=1; i<=n; i++)
	    data[i] *= 2./n;
    return;
} /* end realft() */


#define SWAP(a,b) tempr=(a);(a)=(b);(b)=tempr
void four1(double data[], unsigned long nn, int isign)
{
    unsigned long n,mmax,m,j,istep,i;
    double wtemp,wr,wpr,wpi,wi,theta;
    double tempr,tempi;

    n=nn << 1;
    j=1;
    for (i=1;i<n;i+=2) {
	if (j > i) {
	    SWAP(data[j],data[i]);
	    SWAP(data[j+1],data[i+1]);
	}
	m=n >> 1;
	while (m >= 2 && j > m) {
	    j -= m;
	    m >>= 1;
	}
	j += m;
    }
    mmax=2;
        while (n > mmax) {
	    istep=mmax << 1;
	    theta=isign*(6.28318530717959/mmax);
	    wtemp=sin(0.5*theta);
	    wpr = -2.0*wtemp*wtemp;
	    wpi=sin(theta);
	    wr=1.0;
	    wi=0.0;
                for (m=1;m<mmax;m+=2) {
		    for (i=m;i<=n;i+=istep) {
			j=i+mmax;
			tempr=wr*data[j]-wi*data[j+1];
			tempi=wr*data[j+1]+wi*data[j];
			data[j]=data[i]-tempr;
			data[j+1]=data[i+1]-tempi;
			data[i] += tempr;
			data[i+1] += tempi;
		    }
		    wr=(wtemp=wr)*wpr-wi*wpi+wr;
		    wi=wi*wpr+wtemp*wpi+wi;
                }
                mmax=istep;
        }
} /* end four1() */
#undef SWAP


/*****************************************************/
/* Calculate the power spectrum and auto-correlation */
/* function of a "current" averaged over three       */
/* orthogonal components and write result to OUTPUT. */
/*****************************************************/
double smile (int w, FILE *output)
{
    int i, j, k, l; 
    unsigned long size, n, ncorr, nfreq;
    double delta, total, first_dip, *data, *corr,
	freq[SMILE_FREQ_CHANNELS] = {0};
    FILE *current[3], *freq_output, *corr_output;
    
    if ( (w!=SMILE_HEAT) && (w!=SMILE_STRESS) )
    {
	fprintf (stderr, "smile: illegal option w = %d.\n", w);
	exit(1);
    }

    for (i=0; i<3; i++)
	if ((current[i]=fopen(fn_flux[w][i],"r")) == NULL)
	{
	    fprintf (stderr, "smile: cannot open file \"%s\".\n",
		     fn_flux[w][i]);
	    exit(1);
	}
    
    /* the first double precision number is delta in ps */
    fread (&delta, sizeof(double), 1, current[0]);
    printf ("Found delta = %.4f ps\n", delta);
    fseek (current[0], 0L, SEEK_END);
    size = ftell(current[0]);
    if (size/sizeof(double)*sizeof(double)!=size)
    {
	fprintf (stderr, "smile: size of \"%s\" = %ld.\n",
		 fn_flux[w][0], size);
	exit(1);
    }
    size = size/sizeof(double)-1;
    
    /* FFT only uses the chunk that is integer power of 2 */
    for (n=2; n<=size; n*=2); n/=2;
    printf ("\nUsable %s current length = %ld,\n",
	    fn_flux[w][0], n);
    if ( (data=(double*)malloc(n*sizeof(double))) == NULL )
    {
	fprintf (stderr, "smile: not enough memory to load in %s\n.",
		 fn_flux[w][0]);
	exit(1);
    }
    /* conform to Fortran array convention */
    data--; 

    nfreq = (SMILE_FREQ_CHANNELS>n/2)?n/2:SMILE_FREQ_CHANNELS;
    
    for (j=0; j<3; j++)
    {
	printf ("Loading in \"%s\"...\n", fn_flux[w][j]);
	fseek (current[j], -n*sizeof(double), SEEK_END);
	fread (&data[1], sizeof(double), n, current[j]);
	fclose (current[j]);
	printf ("raw data loaded, FFT this current...\n");

	if (n>=2) realft (data, n, 1);
	printf ("FFT complete, getting the power spectrum...\n");
	    
	/* get rid of net drift */
	data[1] = 0.;
	if (n>=2) data[2] = data[2]*data[2]/n;
	for (i=3; i<=n; i+=2)
	{
	    data[i] = (data[i]*data[i]+data[i+1]*data[i+1])/n;
	    data[i+1] = 0.;
	}
	/* divide by 2 to be the single sided integral */
	for (i=0; i<nfreq; i++)
	    freq[i] += data[2*i+1]/3./2.;

	printf ("Inverse FFT the power spectrum...\n");
	if (n>=2) realft (data, n, -1);
	printf ("complete, accumulate to counters.\n\n");

	if (j==0)
	{  /* search approx. location corr. turns negative */
	    for (ncorr=0;
		 (ncorr<n)&&(data[1+ncorr/SMILE_CORR_EXTENSION]>0.);
		 ncorr++);
	    corr = (double *) malloc((ncorr+1)*sizeof(double));
	    corr--;
	    for (i=1; i<=ncorr; i++) corr[i] = data[i]/3.;
    	}
	else for (i=1; i<=ncorr; i++) corr[i] += data[i]/3.;
    } /* loop over principle directions */
    
    /* look for the first point where the */
    /* correlation function turns negative */
    for (i=1; (i<=ncorr)&&(corr[i]>0.); i++);
    for (total=0.,l=1; l<i; l++) total += corr[l];
    
    fprintf (output, "\n   Based on a piece of %.1f ps current data (x 3),\n",
	     n * delta);
    fprintf (output, (w==SMILE_HEAT)?
	     "   <J(0)^2>/kT^2V = %f W/M/K/ps,\n":
	     "   <s_xy(0)^2>V/kT = %e Pa,\n", corr[1]/delta);
    fprintf (output,
	     "   the first dip in correlation function occurs at %.3f ps\n",
	     i*delta);
    if (w == SMILE_HEAT)
	fprintf (output,
		 "   and by then the thermal conductivity is %.3f mW/m/K.\n\n",
		 1000*total);
    else
	fprintf (output,
		 "   and by then the shear viscosity is %.3f uPa.second.\n\n",
		 1000000*total);
    
    /* save the correlation function */
    corr_output = fopen (fn_corr[w], "w+");
    for (i=1; i<=ncorr; i++)
	fprintf (corr_output, "%.4f %e\n", (i-1)*delta, corr[i]);
    fclose (corr_output);
    
    /* save the power spectrum */
    freq_output = fopen (fn_freq[w], "w+");
    for (i=0; i<nfreq; i++)
	fprintf (freq_output, "%d %f\n", i, freq[i]);
    fclose (freq_output);

    /* free the allocations */
    free(data+1);
    free(corr+1);
    return (total); 
} /* end smile() */


#ifdef SMILE_TEST
int main (int argc, char *argv[])
{
    if (argc!=2)
    {
	printf ("Usage: smile 0 (conductivity), or 1 (viscosity)\n");
	return(1);
    }
    smile (atoi(argv[1]), stdout);
    return(0);
} /* end main() */
#endif