#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>
#include <err.h>
#include <math.h>
#include <fftw3.h>

int
usage(const char* progname)
{
	fprintf(stderr, "usage: %s [-n count] [-m] [-w]\n", progname);
	return 1;
}

void
apply_window(float* samples, size_t n)
{
	for (size_t i = 0; i < n; ++i) {
		float a = (1.0 - cosf((2.0f * M_PI * i) / n)) / 2.0;
		samples[i] *= a;
	}
}

int
main(int argc, char** argv)
{
	int optch;
	int nflag = 0;
	const char* narg = NULL;
	int mflag = 0;
	int wflag = 0;

	while ((optch = getopt(argc, argv, "n:mwh")) != -1) {
		switch (optch) {
		case 'n':
			nflag = 1;
			narg = optarg;
			break;
		case 'm':
			mflag = 1;
			break;
		case 'w':
			wflag = 1;
			break;
		case 'h':
		default:
			return usage(*argv);
		}
	}

	if (!nflag)
		return usage(*argv);

	unsigned nsamples = strtoull(narg, NULL, 10);
	float* samples = aligned_alloc(0x40, nsamples * sizeof(float));

	if (samples == NULL)
		err(1, "malloc()");

	fftwf_complex* output = fftwf_malloc(nsamples * sizeof(fftwf_complex)); 
	fftwf_plan plan = fftwf_plan_dft_r2c_1d(nsamples, samples, output, 0);

	while (fread(samples, sizeof(float), nsamples, stdin) == nsamples) {
		if (wflag) {
			apply_window(samples, nsamples);
		}

		fftwf_execute(plan);

		if (!mflag) {
			fwrite(output, sizeof(fftwf_complex), nsamples / 2, stdout);
		} else {
			size_t i;
			
			for (i = 0; i < nsamples / 2; ++i) {
				float mag = sqrtf(
					output[i][0] * output[i][0] +
					output[i][1] * output[i][1]
				);

				output[i / 2][i & 1] = mag;
			}

			memset(&output[i], 0, (nsamples - i) * sizeof(fftwf_complex));
			fwrite(output, sizeof(float), nsamples / 2, stdout);
		}
	}

	return 0;
}