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

#include "symbol.h"
#include "code.h"
#ifdef VMS
#include "math_tab.h"
#else
#include "math.tab.h"
#endif
#include "fudgit.h"
#include "setshow.h"
#include "dalloca.h"
#include "head.h"

static void fft(double *data, int nn, int isign);
static void merge(double *real, double *ima, double *comp, int nn);
static void split(double *comp, double *real, double *ima, int nn);
static int checkbit(unsigned int data);
static int fftsmooth(double *fromvec, double *smvec, int ndata, float sigma);

/* called with 4 arguments (argc=5) */
int Ft_run_fft(int argc, char **argv, int dir, int ndata)
{
	double *vectors[4], *complex;
	Symbol *sym;
	int i;
	extern Symbol *Ft_lookup(char *s), *Ft_install(char *s, int t, int size);
	int Ft_varcpy(char *to, char *from);

	if (ndata == 0) {
		fputs("(inv)fft: No data!\n", stderr);
		return(ERRR);
	}
	if (checkbit(ndata) != 1) {
		fputs("(inv)fft: Fast fourier transform requires 2^n data points.\n",
		stderr);
		return(ERRR);
	}
	ADVECTOR(complex, 1, 2*ndata, "(inv)fft", return);
	for (i=1; i<=4; i++) {
		if ((sym = Ft_lookup(argv[i])) == 0) { /* get all vectors */
			if (i<=2) {
				fprintf(stderr, "(inv)fft: %s: Unknown vector.\n", argv[i]);
				return(ERRR);
			}
			if (Ft_varcpy(NULL, argv[i]) != VEC) {
				fprintf(stderr,
				"(inv)fft: %s: Not a legal vector name.\n", argv[i]);
				return(ERRR);
			}
			sym = Ft_install(argv[i], VEC, Ft_Samples);
		}
		vectors[i-1] = sym->u.vec;
	}
	/* make it complex */
	merge(vectors[0], vectors[1], complex, ndata);
	fft(complex, ndata, dir);  /* take the fourier transform  */
	/* split in two vectors  */
	split(complex, vectors[2], vectors[3], ndata);
	return(0);
}

/* called with 3 arguments (argc=4) */
int Ft_run_smooth(int argc, char **argv, int ndata)
{
	double *vectors[2];
	Symbol *sym;
	int i;
	float sigma;
	extern Symbol *Ft_lookup(char *s), *Ft_install(char *s, int t, int size);
	int Ft_varcpy(char *to, char *from);

	if (ndata == 0) {
		fputs("smooth: No data!\n", stderr);
		return(ERRR);
	}
	if (sscanf(argv[1], "%f", &sigma) != 1) {
		fputs("Error: smooth: Cannot read smoothing factor.\n", stderr);
		return(ERRR);
	}
	for (i=2; i<=3; i++) {
		if ((sym = Ft_lookup(argv[i])) == 0) { /* get all vectors */
			if (i==2) {
				fprintf(stderr,
				"Error: smooth: %s: Unknown vector.\n", argv[i]);
				return(ERRR);
			}
			if (Ft_varcpy(NULL, argv[i]) != VEC) {
				fprintf(stderr,
				"Error: smooth: %s: Not a legal vector name.\n", argv[i]);
				return(ERRR);
			}
			sym = Ft_install(argv[i], VEC, Ft_Samples);
		}
		vectors[i-2] = sym->u.vec;
	}
	return(fftsmooth(vectors[0], vectors[1], ndata, sigma));
}

static int checkbit(unsigned int data)
{
	int bitsum = 0;
	int i;

	for (i=0;i<8*sizeof(int);i++) {
		bitsum += data & 01;
		data >>= 1;
	}
	return(bitsum);
}

static void split(double *comp, double *real, double *ima, int nn)
{
	int i, j;
	double norm, sqrt(double);

	norm = sqrt((double)nn);
	for (i=1,j=1;j<=nn;i+=2,j++) {
	   	real[j] = comp[i]/norm;
	   	ima[j] = comp[i+1]/norm;
	}
	return;
}

static void merge(double *real, double *ima, double *comp, int nn)
{
	int i, j;

	if (ima) {
		for (i=j=1;j<=nn;i+=2,j++) { /* merge, i always odd */
			comp[i] = real[j];  /* real  */
			comp[i+1] = ima[j];  /* imaginary  */
		}
	}
	else {
		for (i=j=1;j<=nn;i+=2,j++) { /* merge, i always odd */
			comp[i] = real[j];  /* real  */
			comp[i+1] = 0.0;  /* imaginary  */
		}
	}
	return;
}

#define TWO_PI  6.2831853071795864769252868

static void fft(double *data, int nn, int isign)
{
	int n, mmax, m, j, istep, i;
	double wtemp, wr, wpr, wpi, wi, theta;
	double rtemp, itemp;

	n = nn<<1;
	j = 1;
	for (i=1;i<n;i+=2) {
		if (j > i) {
			rtemp = data[j];
			data[j] = data[i];
			data[i] = rtemp;
			itemp = data[j+1];
			data[j+1] = data[i+1];
			data[i+1] = itemp;
		}
		m = n>>1;
		while (m >= 2 && j > m) {
			j -= m;
			m >>= 1;
		}
		j += m;
	}
	mmax = 2;
	while (n > mmax) {
		istep = mmax<<1;
		theta = TWO_PI/(isign*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;
				rtemp = wr*data[j] - wi*data[j+1];
				itemp = wr*data[j+1] + wi*data[j];
				data[j] = data[i]-rtemp;
				data[j+1] = data[i+1]-itemp;
				data[i] += rtemp;
				data[i+1] += itemp;
			}
			wr = (wtemp=wr) * wpr - wi * wpi + wr;
			wi = wi * wpr + wtemp * wpi+ wi;
		}
		mmax = istep;
	}
	return;
}

static int fftsmooth(double *fromvec, double *smvec, int ndata, float sigma)
{
	int fmax, num, k, kk, f, last, vecsize=2;
	double yyn, yy1, rn1, fac;
	double *comp;

	k = ndata << 1;
	while (vecsize < k)
		vecsize <<= 1;
	ADVECTOR(comp, 1, vecsize, "smooth", return);
	merge(fromvec, NULL, comp, ndata);
	num = vecsize >> 1;  	/* num = number of complex elements */
	fmax = vecsize >> 2;	/* maximum frequency  */
	sigma *= fmax;
	sigma *= -sigma;
	last = 2*ndata-1;
	yy1 = comp[1];
	yyn = comp[last];
	rn1 = 1.0/(ndata-1);
	for (kk=f=1;f<=ndata;f++,kk+=2) {
		comp[kk] += (-rn1*(yy1*(ndata-f) + yyn*(f-1)));
	}
	for (f=last+2;f<=vecsize;f++) {
		comp[f] = 0.0;   /* zero pad both real and ima */
	}
	fft(comp, num, 1);
	comp[1] /= num;  /* normalize: should be real (comp[2] == 0) */
	for (f=1;f<=fmax;f++) {
		k = 2*f+1;					/* positive frequencies */
		kk = vecsize-k+2;  			  /* negative frequencies */
		fac = exp(f*f/sigma)/num;	 /* Gaussian window */
		comp[kk] = (comp[k] *= fac);  /* real */
		comp[kk+1] = -(comp[k+1] *= fac);  /* ima */
	}
	fft(comp, num, -1);
	for (kk=f=1;f<=ndata;f++,kk+=2) {
		smvec[f] = comp[kk] + (rn1*(yy1*(ndata-f) + yyn*(f-1)));
	}
	return(0);
}
