#include <stdint.h>
#include <assert.h>	/* sheer abuse */
#include <math.h>

#include <usrp2/usrp2.h>
#include <usrp2/rx_nop_handler.h>
#include <usrp2/copiers.h>

#include <fftw3.h>

#include "SDL.h"
#include "SDL_gfxPrimitives.h"


//#define	N	1024
#define	XX	100
#define	N	(1024*XX)
#define	AVG	1


#define	SCALE	(2*XX)
#define	XRES	(N/SCALE)
#define	YRES	(768*XX/SCALE)
#define	DECIMATION	16

static fftw_plan plan;
static fftw_complex *in, *out;
static double *w;
static SDL_Surface *surf;


#define	SIZE	(surf->h*surf->pitch)


static int frames = 0;
static int paws = 0;
static int y_scale = 60;
static uint32_t *bg = NULL;
static long rate;
static int snapshot = 0;
static int avg_mode = 1;


static void draw(const double *s)
{
	int i, x, y, first, last;

	if (!bg)
		bg = (uint32_t *) calloc(SIZE, 1);
	if (paws)
		return;
	SDL_LockSurface(surf);
	memcpy(surf->pixels, bg, SIZE);
	for (y = y_scale; y < YRES; y += y_scale)
		hlineColor(surf, 0, XRES-1, y, 0x008040ff);
	uint32_t color = snapshot ? 0xff4000ff : 0xffffffff;
	last = YRES-1;
	for (i = 0; i != XRES; i++) {
		y = s[i]*y_scale;
		y = YRES-y+4*y_scale;
		if (y < 0)
			y = 0;
		if (y >= YRES)
			y = YRES-1;
		x = (i+XRES/2) % XRES;
		if (i)
			vlineColor(surf, x, last, y, color);
		if (!i)
			first = y;
		last = y;
	}
	vlineColor(surf, x+1, first, last, color);
	if (snapshot) {
		memcpy(bg, surf->pixels, SIZE);
		snapshot = 0;
	}
	SDL_UnlockSurface(surf);
	SDL_UpdateRect(surf, 0, 0, 0, 0);
	frames++;
}


static void collect(void)
{
	static double sum[N];
	static int cnt = 0;
	int i, x, j;

	i = 0;
	for (x = 0; x != XRES; x++)
		for (j = 0; j != SCALE; j++) {
			double s = pow(out[i][0], 2)+pow(out[i][1], 2);

			if (avg_mode)
				sum[x] += s;
			else {
				if (sum[x] < s)
					sum[x] = s;
			}
			i++;
		}
	if (++cnt != AVG)
		return;
	cnt = 0;
	for (i = 0; i != XRES; i++)
		if (avg_mode)
			sum[i] = log(sum[i]/AVG/SCALE)/log(10);
		else
			sum[i] = log(sum[i])/log(10);
	draw(sum);
	memset(sum, 0, sizeof(sum));
}


class rx : public usrp2::rx_nop_handler
{
public:
	rx(uint64_t max_samples) : usrp2::rx_nop_handler(max_samples)
	{
	}

	bool operator()(const uint32_t *items, size_t nitems,
	    const usrp2::rx_metadata *metadata);

	~rx(void)
	{
	}

};


bool rx::operator()(const uint32_t *items, size_t nitems,
	    const usrp2::rx_metadata *metadata)
{
	static std::complex<int16_t> host_items[N];
	static size_t have = 0, copy, over;
	bool ok;
	size_t i;

	ok = rx_nop_handler::operator()(items, nitems, metadata);
	copy = N-have;
	if (copy > nitems)
		copy = nitems;
	usrp2::copy_u2_16sc_to_host_16sc(copy, items, host_items+have);
//fprintf(stderr, "nitems %u have %u copy %u\n", (unsigned) nitems, (unsigned) have, (unsigned) copy);
	have += copy;
	if (have == N) {
		for (i = 0; i != nitems; i++) {
			in[i][0] = host_items[i].real()*w[i];
			in[i][1] = host_items[i].imag()*w[i];
		}
		fftw_execute(plan);
		collect();
		have = 0;
	}
	over = nitems-copy;
	if (over)
		usrp2::copy_u2_16sc_to_host_16sc(over, items+copy,
		    host_items+have);
	have += over;
	return ok;
}


static int key(SDLKey sym)
{
	switch (sym) {
	case SDLK_a:
		avg_mode = 1;
		break;
	case SDLK_c:
		memset(bg, 0, SIZE);
		break;
	case SDLK_f:
		fprintf(stderr, "%d\n", frames);
		break;
	case SDLK_m:
		avg_mode = 0;
		break;
	case SDLK_p:
		paws = !paws;
		break;
	case SDLK_q:
		return 1;
	case SDLK_s:
		memset(bg, 0, SIZE);
		snapshot = 1;
		break;
	default:
		break;
	}
	return 0;
}


static void make_window(void)
{
	int i;

	w = (double *) calloc(N, sizeof(double));
	for (i = 0; i != N; i++)
// hamming
		w[i] = 0.54-0.46*cos(M_PI*2*i/(N-1));
// hann		w[i] = 0.5-0.5*cos(M_PI*2*i/(N-1));
// blackman	w[i] = 0.42-0.5*cos(2*M_PI*i/(N-1))+0.08*cos(4*M_PI*i/(N-1));
}


int main(void)
{
	usrp2::usrp2::sptr u;
	usrp2::rx_nop_handler::sptr handler;
	usrp2::tune_result tr;

	handler = usrp2::rx_nop_handler::sptr(new rx(N));
	u = usrp2::usrp2::make("eth0", "");
	assert(u->stop_rx_streaming());
	assert(u->set_rx_gain(46));
	assert(u->set_rx_center_freq(2450000000ULL, &tr));
	assert(u->set_rx_decim(DECIMATION));
//	rate = u->adc_rate()/u->decim():
//	u->set_sample_rate(rate);

	u->adc_rate(&rate);

	in = (fftw_complex *) fftw_malloc(sizeof(fftw_complex)*N);
	out = (fftw_complex *) fftw_malloc(sizeof(fftw_complex)*N);
	plan = fftw_plan_dft_1d(N, in, out, FFTW_FORWARD, FFTW_ESTIMATE);

	make_window();

	if (SDL_Init(SDL_INIT_VIDEO) < 0) {
		fprintf(stderr, "SDL_init: %s\n", SDL_GetError());
		exit(1);
	}
	atexit(SDL_Quit);

	surf = SDL_SetVideoMode(XRES, YRES, 0, SDL_SWSURFACE);
	if (!surf) {
		fprintf(stderr, "SDL_SetVideoMode: %s\n", SDL_GetError());
		exit(1);
	}

	assert(u->start_rx_streaming(0));
	while (1||(!handler->has_errored_p() && !handler->has_finished_p())) {
//	while (!handler->has_finished_p()) {
		SDL_Event event;

		if (SDL_PollEvent(&event)) {
			if (event.type == SDL_KEYDOWN)
				if (key(event.key.keysym.sym))
					break;
		}
		assert(u->rx_samples(0, handler.get()));
		u->rx_samples(0, handler.get());
	}
	u->stop_rx_streaming(0);
	return 0;
}