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

#include "fft.h"

/* routines from "The Scientist and Engineer's Guide to Digital Signal Processing"
	see BASIC at the end for source
*/

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
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;
	}
}