/************************************/
/* Spectral method to calculate the */
/* scalar thermal conductivity.     */
/*                                  */
/*       version 1.0 Apr.21, 1998   */
/*       developed by Ju Li (MIT)   */
/************************************/

#include <stdio.h>
#include <math.h>
#include <string.h>
#include "argon.h"

void main()
{
    FILE *output = fopen("smile.out", "w+");
    shear_smile (output);
    smile (output);
    fclose(output);
    printf ("Please get your final results in \"smile.out\".\n\n");
    return;
}

void smile(FILE *output)
{
    /* calculate the heat current correlation function averaged over */
    /* three orthogonal directions, and write results in OUTPUT */
    int i, j, k, l; 
    unsigned long size, n, ncorr, nfreq;
    double *data, *freq, *corr;
    double delta, total, first_dip;
    
    heat_current[0] = fopen (fn_heat_current[0], "r");
    fseek (heat_current[0], 0L, SEEK_END);
    size = ftell(heat_current[0]);
    
    /* can only use a chunk that is integer power of 2 */
    for (n=2; (size/sizeof(double)) > n; n<<=1); n>>=1;
    printf ("\nUsable heat current length = %ld,", n);
    if ((data=(double *)malloc(n*sizeof(double))) == NULL) 
    {
	fprintf (stdout, "\n **Error: not enough memory to load in %s\n.",
		 fn_heat_current[0]);
	exit(1);
    }

    /* the last double precision number is delta in ps */
    fseek (heat_current[0], size-sizeof(double), SEEK_SET);
    fread (&delta, sizeof(double), 1, heat_current[0]);
    fclose(heat_current[0]);
    printf ("  delta = %.4f ps\n\n", delta);
    
    ncorr = n/NPART;
    corr = (double *) malloc(ncorr*sizeof(double));
    for (i=0; i<ncorr; i++) corr[i] = 0.;
    nfreq = n/NPART;
    freq = (double *) malloc(nfreq*sizeof(double));
    for (i=0; i<nfreq; i++) freq[i] = 0.;
    /* conform to Fortran convention */
    data--; corr--; freq--;
    
    for (j=0; j<3; j++)
    {
	printf ("Loading in %s...\n", fn_heat_current[j]);
	heat_current[j] = fopen (fn_heat_current[j], "r");
	fseek (heat_current[j], size-(n+1)*sizeof(double), SEEK_SET);
	fread (data+1, sizeof(double), n, heat_current[j]);
	fclose (heat_current[j]);
	printf ("%s loaded in, FFT the heat current...\n", fn_heat_current[j]);

	realft (data, n, 1);
	printf ("FFT complete, getting the power spectrum...\n");
	    
	/* get rid of net drift, pure fluctuation */
	data[1] = 0.;
	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=1; i<=nfreq; i++) freq[i] += data[2*i-1]/3./2.;

	printf ("inverse FFT the power spectrum...\n");
	realft (data, n, -1);
	printf ("inverse FFT complete, accumulate to counters.\n\n");
	for (i=1; i<=ncorr; i++) corr[i] += data[i]/3.;
	
    }  /* end processing principle directions */
    
    /* look for point where the correlation function first turns negative */
    for (i=1; (corr[i]>0.) && (i<=ncorr); i++);
    for (total=0.,l=1; l<=i; l++) total += corr[l];
    
    fprintf (output, "\n   From a piece of %.1f ps heat current data,\n",
	     n*delta);
    fprintf (output, "   the first dip in correlation function occurs at %.3f ps\n",
	     i*delta);
    fprintf (output, "   and by then the thermal conductivity is %.3f mW/m/K.\n\n",
	     1000*total);
    
    /* save the correlation function */
    corr_output = fopen (fn_corr_output, "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_output, "w+");
    for (i=1; i<=nfreq; i++) fprintf (freq_output, "%d %f\n", i-1, freq[i]);
    fclose (freq_output);

    free(++data);
    free(++corr);
    free(++freq);
    
    return; 
} /* end smile() */

void shear_smile(FILE *output)
{
    /* calculate the shear stress correlation function averaged */
    /* over yz, xz, xy components, and write results in OUTPUT  */
    int i, j, k, l; 
    unsigned long size, n, ncorr, nfreq;
    double *data, *freq, *corr;
    double delta, total, first_dip;
    
    shear_stress[0] = fopen (fn_shear_stress[0], "r");
    fseek (shear_stress[0], 0L, SEEK_END);
    size = ftell(shear_stress[0]);
    
    /* can only use a chunk that is integer power of 2 */
    for (n=2; (size/sizeof(double)) > n; n<<=1); n>>=1;
    printf ("\nUsable shear stress length = %ld,", n);
    if ((data=(double *)malloc(n*sizeof(double))) == NULL) 
    {
	fprintf (stdout, "\n **Error: not enough memory to load in %s\n.",
		 fn_shear_stress[0]);
	exit(1);
    }

    /* the last double precision number is delta in ps */
    fseek (shear_stress[0], size-sizeof(double), SEEK_SET);
    fread (&delta, sizeof(double), 1, shear_stress[0]);
    fclose(shear_stress[0]);
    printf ("  delta = %.4f ps\n\n", delta);
    
    ncorr = n/NPART;
    corr = (double *) malloc(ncorr*sizeof(double));
    for (i=0; i<ncorr; i++) corr[i] = 0.;
    nfreq = n/NPART;
    freq = (double *) malloc(nfreq*sizeof(double));
    for (i=0; i<nfreq; i++) freq[i] = 0.;
    /* conform to Fortran convention */
    data--; corr--; freq--;
    
    for (j=0; j<3; j++)
    {
	printf ("Loading in %s...\n", fn_shear_stress[j]);
	shear_stress[j] = fopen (fn_shear_stress[j], "r");
	fseek (shear_stress[j], size-(n+1)*sizeof(double), SEEK_SET);
	fread (data+1, sizeof(double), n, shear_stress[j]);
	fclose (shear_stress[j]);
	printf ("%s loaded in, FFT the shear stress...\n", fn_shear_stress[j]);

	realft (data, n, 1);
	printf ("FFT complete, getting the power spectrum...\n");
	    
	/* get rid of net drift, pure fluctuation */
	data[1] = 0.;
	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=1; i<=nfreq; i++) freq[i] += data[2*i-1]/3./2.;

	printf ("inverse FFT the power spectrum...\n");
	realft (data, n, -1);
	printf ("inverse FFT complete, accumulate to counters.\n\n");
	for (i=1; i<=ncorr; i++) corr[i] += data[i]/3.;
	
    }  /* end processing principle directions */
    
    /* look for point where the correlation function first turns negative */
    for (i=1; (corr[i]>0.) && (i<=ncorr); i++);
    for (total=0.,l=1; l<=i; l++) total += corr[l];
    
    fprintf (output, "\n   From a piece of %.1f ps shear stress data,\n",
	     n*delta);
    fprintf (output, "   the first dip in correlation function occurs at %.3f ps\n",
	     i*delta);
    fprintf (output, "   and by then the shear viscosity is %.3f uPa x second.\n\n",
	     1000000*total);
    
    /* save the correlation function */
    shear_corr_output = fopen (fn_shear_corr_output, "w+");
    for (i=1; i<=ncorr; i++)
	fprintf (shear_corr_output, "%.4f %E\n", (i-1)*delta, corr[i]);
    fclose (shear_corr_output);
    
    /* save the power spectrum */
    shear_freq_output = fopen (fn_shear_freq_output, "w+");
    for (i=1; i<=nfreq; i++) fprintf (shear_freq_output, "%d %f\n", i-1, freq[i]);
    fclose (shear_freq_output);

    free(++data);
    free(++corr);
    free(++freq);
    
    return; 
} /* end shear_smile() */

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