#include <u.h>
#include <libc.h>
#include <draw.h>
#include <event.h>
#include <keyboard.h>

#include "fft.h"

enum {
	/* HZ */
	RATE		= 44100,

	/* also HZ */
	CUTOFF		= 20000,

	/* 44100 * 0.4 = 400ms of samples */
	NSAMPLES	= 11025,

	/* 44100 * 0.05 = 50 ms of samples */
	WINDOW		= 2205,

	/* nearest power of 2 */
	NFFT		= 1<<13,
};

double
rms(short *v, int n)
{
	int i;
	double tmp, in, out, sum;

	const double tau = 2 * 3.141592653 * CUTOFF / RATE;

	sum = 0;

	for(i = 0; i < n; i++){
		tmp = v[i]/32767.0;
		in = tmp * tmp;

		//out = (1 - tau/2) * in - tau * sum;
		//out = (1-tau/2) * in - tau*sum;

		out = in;
		sum += out;
	}

	return sqrt(sum/(double)n);
}

double rex[NFFT];
double imx[NFFT];
Freq freqs[NFFT];

void
dofft(int n, short *v, double *rex, double *imx, Freq *freqs)
{
	int i;

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

	fft(n, rex, imx);
	ris_to_fs(n, rex, imx, freqs);
}

short hist[NSAMPLES];
short buf[NFFT];
double rmsv, db;
Image *green, *red;

void
ddot(Point p)
{
	Rectangle r;

	r = screen->r;
	r.min = addpt(r.min, p);
	r.max = addpt(r.min, Pt(1,1));

	draw(screen, r, red, nil, ZP);
}

double
itofreq(int idx, int freq, int nsamples)
{
	double didx, dfreq, dnsamples;

	didx = idx;
	dfreq = freq;
	dnsamples = nsamples;

	return didx * dfreq / dnsamples;
}

double
getdb(int k)
{
	double mag;

	mag = 2 * freqs[k].r / NFFT*2;
	if(mag == 0.0)
		return 0.0;

	return 20 * log10(mag);
}

void
fftplot(void)
{
	int x, xf, i, wid, skip;
	double v;
	Point p;

	wid = Dx(screen->clipr);

	if(wid > NFFT/2)
		wid = NFFT/2;

	skip = NFFT/2 / wid;
	if(skip < 1)
		skip = 1;

	xf = 0;

	for(x = 0; x < wid; x++){
		v = 0;
		for(i = 0; i < skip; i++){
			v += getdb(xf+i);
		}

		v /= skip;
		xf += skip;

		p.x = x;
		p.y = abs(v - 10);
		p.y += 30;

		ddot(p);
	}
}

void
redraw(void)
{
	int x, max;
	double pct;
	Rectangle r;
	char buf[128];

	static int t = 0;

	t++;

	//if(t++ % 2 == 0)
		//return;

	draw(screen, screen->r, display->black, nil, ZP);

	pct = (fabs(db))/100;
	r = screen->r;
	r.min.y += (pct * Dy(r));
	//draw(screen, r, green, nil, ZP);


	fftplot();

	//max = (Dx(screen->r) > WINDOW) ? WINDOW : Dx(screen->r);
	max = NFFT/2;

	double big = 0;
	int bigidx = 0;
	for(x = 0; x < max; x++){
		//fprint(2, "%5.2f ", freqs[x].r);
		if(freqs[x].r > big){
			big = freqs[x].r;
			bigidx = x;
		}
	}

	r = screen->r;
	snprint(buf, sizeof(buf), "%d RMS %5.2f dB %5.2f", t, rmsv, db);
	string(screen, r.min, red, ZP, font, buf);
	r.min.y += font->height + 2;

	snprint(buf, sizeof(buf), "big = %4d (%8.2f HZ) %6.2f DB = %5.2f", bigidx, itofreq(bigidx, RATE, NFFT/2), big, getdb(bigidx));
	string(screen, r.min, red, ZP, font, buf);

	flushimage(display, 1);
}

void
main(int argc, char *argv[])
{
	long n;

	ARGBEGIN{
	default:
		sysfatal("no arguments");
	}ARGEND

	memset(hist, 0, sizeof(hist));

	if(initdraw(nil, nil, argv0) < 0)
		sysfatal("initdraw: %r");

	green = allocimage(display, Rect(0,0,1,1), screen->chan, 1, DGreen);
	if(green == nil)
		sysfatal("allocimage: %r");

	red = allocimage(display, Rect(0,0,1,1), screen->chan, 1, DRed);
	if(red == nil)
		sysfatal("allocimage: %r");

	einit(Emouse);

	eresized(0);

	while((n = readn(0, buf, sizeof(buf))) > 0){
		if(n == -1)
			break;

		if(n != sizeof(buf))
			sysfatal("short read: %r");

		//memmove(hist+sizeof(buf), hist, sizeof(hist)-sizeof(buf));
		//memcpy(hist, buf, sizeof(buf));

		dofft(nelem(buf), buf, rex, imx, freqs);

		rmsv = rms(buf, nelem(buf));
		if(rmsv == 0.0)
			rmsv = 0.01;
		db = 20.0 * log10(rmsv);

		//fprint(2, "RMS %5.2f dB %5.2f\n", rmsv, db);

		while(ecanmouse()){
			emouse();
		}

		redraw();

		write(1, buf, n);
	}

	exits(nil);
}

void
eresized(int new)
{
	if(new && getwindow(display, Refnone) < 0)
		sysfatal("getwindow: %r");

	redraw();
}