// Put 	Font

//	8c -w fft.c &&  8l fft.8  && window '8.out -n 16  | plot '

//   8c fft.c &&  8l fft.8  && 8.out          		

//	;	mv 8.out /usr/maht/bin/386/fft 


#include <u.h>
#include <libc.h>

double
f(int n, double x) {
	x *= 2 * PI / n;
//	x += PI;
	return sin(x); // + sin(2 * x) + 0.5 * sin(3 * (PI + x));

}

/* routines from "The Scientist and Engineer's Guide to Digital Signal Processing"
	see BASIC at the end for source
*/
 
typedef struct Complex Complex;
typedef struct Freq Freq;
struct Complex
{
	double r, i;
};

struct Freq
{
	double f, r, a;
};

void
C_to_F(Complex *c, Freq *f) {
	f->r = sqrt(c->r * c->r + c->i * c->i);
	f->a = (c->r > 0) ? atan2(c->i, c->r) : 0;
}

void
ri_to_F(double r, double i, Freq *f) {
	f->r = sqrt(r * r + i * i);
	f->a = atan2(i, r);
	if(f->r > 0)
		f->a += PI;
}

void
ris_to_fs(int n, double rex[], double imx[], Freq f[]) {
	int i;

	for(i = 0; i < n; i++) {
		ri_to_F(rex[i], imx[i], &f[i]);
	}
}

void
plot(int n, double rex[], char *c, double rad, int r) {
	int i;
	double f;
	if(r)
		print("ra 0 -1.1 1 1.1\n");
	if(c)
		print("co %s\n", c);
	for(i = 0; i < n; i++) {
		f = (double)i / (double)n;
		print("ci %0.3f %0.3f %0.3f\n", f, rex[i], rad);
	}
}

void
range(int n, double rex[], Freq f[]) {
	int i;
	double minf = 0;
	double maxf = 0;

//	print("# %d mf %0.3f\n", n, maxf);
	for(i = 0; i < n; i++) {
		if(rex[i] < minf)
			minf = rex[i];
		else if (rex[i] > maxf)
			maxf = rex[i];
//		print("# f %d, %0.3f\n", i, f[i].r);
		if (f[i].r/n > maxf)
			maxf = f[i].r/n;
		if(f[i].a/PI < minf)
			minf = f[i].a/PI;
		else if (f[i].a/PI > maxf)
			maxf = f[i].a/PI;
	}
//	print("# mf %0.3f\n", maxf);

	print("ra 0 %0.3f 1 %0.3f\n", minf, maxf);
}

void
plot_freqs(int n, Freq f[], double rad) {
	int i;
	double x;
	for(i = 0; i < n; i++) {
		x = (double)i / (double)n;
		print("ci %0.3f %0.3f %0.3f\n", x, f[i].r/n, rad);
		print("ci %0.3f %0.3f %0.3f\n", x, f[i].a/PI, rad/2);
	}
		
}

void
fft_bit_reversal(int n, double rex[], double imx[]) {
	int nd2 = n >> 1;
	int i, j, k;
	double tr, ti;
	for(i = 1, j = nd2; i < n-1; i++) { 
		if(i < j) { 
			tr = rex[j];
			ti = imx[j];
			rex[j] = rex[i];
			imx[j] = imx[i];
			rex[i] = tr;
			imx[i] = ti;
		}
		for(k = nd2; k <= j ; j -= k, k >>= 1);
		j += k;
	}
}

void
fft_butterfly(int n, double rex[], double imx[]) {
	int i, j, k, m, ie, ie2, kp;
	double ur, ui, tr, ti, sr, si, pie2;
	m = ceil(log(n) / log(2));
	for(i = 1, ie = 2, ie2 = 1; i <= m; i++, ie <<= 1, ie2 <<= 1) {
		ur = 1;
		ui = 0;
		pie2 = PI / ie2;
		sr = cos(pie2);
		si = -sin(pie2);
		for(j = 0; j < ie2; j++) {
			for(k = j; k < n; k += ie) { 
				kp = k + ie2;
				tr = rex[kp] * ur - imx[kp] * ui;
				ti = rex[kp] * ui + imx[kp] * ur;
				rex[kp] = rex[k] - tr;
				imx[kp] = imx[k] - ti;
				rex[k] += tr;
				imx[k] += ti;
			}
			tr = ur;
			ur = tr * sr - ui * si;
			ui = tr * si + ui * sr;
		}
	}
}

void
fft(int n, double rex[], double imx[]) { // 1000
	fft_bit_reversal(n, rex, imx);
	fft_butterfly(n, rex, imx);
}

void
ifft(int n, double rex[], double imx[]) { // 2000
	int i;

	for(i = 0; i <= n-1; i++)
		imx[i] = -imx[i];

	fft(n, rex, imx);

	for(i = 0; i <= n-1; i++) {
		rex[i] /= n;
		imx[i] /= -n;
	}
}

void
dft(int n, double rex[], double imx[]) { // 5000
	int i, k;
	double c, s, np, knp, p;

	double *xr = (double*)calloc(n, sizeof(double));
	double *xi =  (double*)calloc(n, sizeof(double));
	
	memcpy(xr, rex, n * sizeof(double));
	memcpy(xi, imx, n * sizeof(double));
	
	np = 2 * PI / n;
	for(k = 0, knp = 0; k < n; k++, knp += np) {
		rex[k] = 0;
		imx[k] = 0;
		for(i = 0, p = 0; i < n; i++, p += knp) {
			c = cos(p);
			s = -sin(p);
			rex[k] += xr[i] * c - xi[i] * s;
			imx[k] += xr[i] * s + xi[i] * c;
		}
	}
}

void
idftx(int n, double rex[], double imx[], Complex *cx, double a) {
	int i;
	double c, s;
	cx->r = 0;
	cx->i = 0;
	for(i = 0; i < n; i++) {
		cx->r += xr[i] 
	}
}



void
idft(int n, double rex[], double imx[]) { // adapted from 2000
	int i;

	for(i = 0; i < n; i++)
		imx[i] = -imx[i];

	dft(n, rex, imx);

	for(i = 0; i < n; i++) {
		rex[i] /= n;
		imx[i] /= -n;
	}
}


void
separate_odd_even(int n, double rex[], double imx[]) {
	int i, k;
	for(i = 0, k = 0; i <= n; i++, k = i * 2) {
		rex[i] = rex[k];
		imx[i] = rex[k + 1];
	}
}

void
odd_even_freq_decomp(int n, double rex[], double imx[]) {
	int nd2 = n / 2;
	int n4 = n / 4;
	int i, im, ip2, ipm;

	for(i = 1; i < n4; i++) {
		im = nd2 - i;
		ip2 = i + nd2;
		ipm = im + nd2;
		rex[ip2] = (imx[i] + imx[im]) / 2;
		rex[ipm] = rex[ip2];
		imx[ip2] = (rex[i] - rex[im]) / -2;
		imx[ipm] = -imx[ip2];
		rex[i] = (rex[i] + rex[im]) / 2;
		rex[im] = rex[i];
		imx[i] = (imx[i] - imx[im]) / 2;
		imx[im] = -imx[i];
	}
	rex[n * 3/4] = imx[n4];
	rex[nd2] = imx[0];
	imx[n * 3/4] = 0;
	imx[nd2] = 0;
	imx[n4] = 0;
	imx[0] = 0;
}

void
fftr(int n, double rex[], double imx[]) { // 3000

	separate_odd_even(n / 2 - 1, rex, imx);
	fft(n / 2, rex, imx);
	odd_even_freq_decomp(n, rex, imx);

	int i, ip, j, jm1, k, ke, ke2, nm1;
	double ur, ui, sr, si, tr, ti;

	nm1 = n - 1;
	k = ceil(log(n) / log(2));
	ke = pow(2, k);
	ke2 = ke / 2;
	ur = 1;
	ui = 0;
	sr = cos(PI / ke2);
	si = -sin(PI / ke2);
	for(j = 1; j <= ke2; j++) {
		jm1 = j - 1;
		for(i = jm1; i <= nm1; i += ke) {
			ip = i + ke2;
			tr = rex[ip] * ur - imx[ip] * ui;
			ti = rex[ip] * ui - imx[ip] * ur;
			rex[ip] = rex[i] - tr;
			imx[ip] = imx[i] - ti;
			rex[i] = rex[i] + tr;
			imx[i] = imx[i] + ti;
		}
		tr = ur;
		ur = tr * sr - ui * si;
		ui = tr * si + ui * sr;
	}	
}

void
ifftr(int n, double rex[], double imx[]) { // 4000
	int k;
	for(k = n / 2 + 1; k < n; k++) {
		rex[k] = rex[n - k];
		imx[k] = -imx[n - k];
	}

	for(k = 0; k < n; k++)
		rex[k] += imx[k];

	fftr(n, rex, imx);

	for(k = 0; k < n; k++) {
		rex[k] += imx[k];
		rex[k] /= n;
		imx[k] = 0;
	}
}

void
negative_freq_generation(int n, double rex[], double imx[]) { // 6000
	int k;
	for(k = n / 2 + 1; k < n; k++) {
		rex[k] = rex[n - k];
		imx[k] = -imx[n - k];
	}
}


int
test_fft(int n) {
	int i;
	double *rex = (double*)calloc(n, sizeof(double));
	double *imx =  (double*)calloc(n, sizeof(double));
	srand(1);
	for(i = 0; i < n; i++) {
		rex[i] = f(n, i);
		imx[i] = 0;
	}
	fft(n, rex, imx);
	ifft(n, rex, imx);
	srand(1);
	for(i = 0; i < n; i++) {
		if(fabs(rex[i] - f(n, i)) > 0.002)
			break;
	}
	free(rex);
	free(imx);

	return i == n;
}

int
test_dft(int n) {
	double *rex = (double*)calloc(n, sizeof(double));
	double *imx = (double*)calloc(n, sizeof(double));
	int i;
	for(i = 0; i < n; i++) {
		rex[i] = f(n, i);
		imx[i] = 0;
	}
	dft(n, rex, imx);
	idft(n, rex, imx);
	for(i = 0; i < n; i++) {
		if(fabs(rex[i] - f(n, i)) > 0.002)
			break;
	}
	free(rex);
	free(imx);

	return i == n;
}

void
test_algs(void) {
	int i, n;
	for(i = 3; i < 6; i++) {
		n = pow(2, i);
		print("i %d %d samples\n", i, n);
		if(test_dft(n))
			print("ok\n");
		else
			print("fail\n");
	}
}

void
main(int argc, char **argv)
{
	int n = 8;
	ARGBEGIN {
	case 'n' :
		n = atoi(ARGF());
		break;
	} ARGEND


	print("atan(2/1) %0.3f  atan2(2,1) %0.3f	%0.3f	%0.3f\n", atan(2/1), atan2(2, 1),  2 * PI / atan2(2, 1) , 2 * PI * 64.435 / 360);

	exits(nil);


	double *rex = (double*)calloc(n, sizeof(double));
	double *imx = (double*)calloc(n, sizeof(double));
	Freq *freq = (Freq*)calloc(n>>1, sizeof(Freq));
	int i;
	for(i = 0; i < n; i++) {
		rex[i] = f(n, i);
		imx[i] = 0;
	}
	fft(n, rex, imx);
	ris_to_fs(n, rex, imx, freq);
	ifft(n, rex, imx);
	range(n, rex, freq);
	plot_freqs(n, freq, 0.02);
	plot(n, rex, "red", 0.01, 0);

	free(rex);
	free(imx);
	free(freq);
}