/*
 *  Code for frequency analysis
 *
 *  Copyright (C) 2015  Nest Labs. All rights reserved.
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include "fftw3.h"
#include "wav_header.h"

#define SAMPLE_BYTES			2
#define MAX_NUM_CHANNELS		2

#define RECTANGULAR			1
#define HAMMING				2
#define BLACKMAN_NUTTALL		3
#define HANNING				4
#define BLACKMAN_HARRIS			5
#define WINDOW_TYPE			HANNING

#define SQUARE(a)			( (a) * (a) )

void apply_window(int16_t *input, double *windowed, int32_t num_samples);

#if WINDOW_TYPE == RECTANGULAR
void apply_window(int16_t *input, double *windowed, int32_t num_samples)
{
	int32_t sample_index;
	for(sample_index = 0; sample_index < num_samples; sample_index++)
		windowed[sample_index] = (double)input[sample_index]/32767.0;
}
#elif WINDOW_TYPE == HAMMING
void apply_window(int16_t *input, double *windowed, int32_t num_samples)
{
	int32_t sample_index;
	for(sample_index = 0; sample_index < num_samples; sample_index++) { 
		windowed[sample_index] = 0.54 - 0.46 * cos(2 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] *= (double)input[sample_index]/32767.0;
	}
}
#elif WINDOW_TYPE == BLACKMAN_NUTTALL 
void apply_window(int16_t *input, double *windowed, int32_t num_samples)
{
	int32_t sample_index;
	for(sample_index = 0; sample_index < num_samples; sample_index++) { 
		windowed[sample_index] = 0.3635819;
		windowed[sample_index] -= 0.4891775 * cos(2 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] += 0.1365995 * cos(4 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] -= 0.0106411 * cos(6 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] *= (double)input[sample_index]/32767.0;
	}
}
#elif WINDOW_TYPE == HANNING
void apply_window(int16_t *input, double *windowed, int32_t num_samples)
{
	int32_t sample_index;
	for(sample_index = 0; sample_index < num_samples; sample_index++) { 
		windowed[sample_index] = 0.5 - 0.5 * cos(2 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] *= (double)input[sample_index]/32767.0;
	}
}
#elif WINDOW_TYPE == BLACKMAN_HARRIS
void apply_window(int16_t *input, double *windowed, int32_t num_samples)
{
	int32_t sample_index;
	for(sample_index = 0; sample_index < num_samples; sample_index++) { 
		windowed[sample_index] = 0.35875;
		windowed[sample_index] -= 0.48829 * cos(2 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] += 0.14128 * cos(4 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] -= 0.01168 * cos(6 * M_PI * sample_index/(num_samples -1));
		windowed[sample_index] *= (double)input[sample_index]/32767.0;
	}
}
#endif

typedef struct {
	double *data_r;
	double *mag;
	fftw_complex *data_c;
} scratch_t;

typedef struct {
	double max_freq;
	double max;
} output_t;

void execute_fft(int16_t * input_buffer,
		 int32_t nfft,
		 int32_t offset,
		 int32_t num_input_samples,
		 int32_t sample_rate,
		 scratch_t *scratch,
		 output_t *fft_output)
{
	fftw_plan p;
	int32_t max_freq_index, sample_index, block_index;
	double *data_r = scratch->data_r;
	double *mag = scratch->mag;
	double max = -32768.0; /* Initialize to a very low dB */
	fftw_complex *data_c = scratch->data_c;

	/* This array will contain the average over multiple blocks*/
	for(sample_index = 0; sample_index < nfft/2+1; sample_index++) {
		mag[sample_index] = 0;
	}
	/* Setup the plan to compute FFT */
	p = fftw_plan_dft_r2c_1d(nfft, data_r, data_c, FFTW_ESTIMATE);
	/* Execute FFT */
	for(block_index = 0; block_index * offset + nfft < num_input_samples; block_index++) {
		/* Windowing */
		apply_window(input_buffer + block_index * offset, data_r, nfft);
		/* FFT */
		fftw_execute(p);
		/* Accumulate for averaging. nfft is the normalization factor to
		 * account for the difference in forward and reverse DFT
		 */
		for(sample_index = 0; sample_index < nfft/2+1; sample_index++) {
			mag[sample_index] += SQUARE(data_c[sample_index][0]/nfft) + SQUARE(data_c[sample_index][1]/nfft);
		}
	}
	/* Find the tone with the highest frequency, ignore the DC bin */
	for(sample_index = 1; sample_index < nfft/2+1; sample_index++) {
		mag[sample_index] = 10 * log10(mag[sample_index]/block_index);
		if (mag[sample_index] > max) {
			max = mag[sample_index];
			max_freq_index = sample_index;
		}
	}
	fft_output->max = max;
	fft_output->max_freq = (float)max_freq_index/nfft * sample_rate;
	printf("Number of blocks for FFT %d\n", block_index);
#if 0
	/* Print the final output */
	for(sample_index = 0; sample_index < nfft/2+1; sample_index++) {
		mag[sample_index] = 10 * log10(mag[sample_index]/block_index);
		printf("%d: %lf(dB)\n", sample_index, mag[sample_index]);
	}
#endif
	fftw_destroy_plan(p);
}

int check_boundary(output_t *fft_output, double input_freq, double freq_dev,
					 double input_vol, double vol_dev)
{
	double lower_limit, upper_limit;

	/* check if freq is within deviation */
	lower_limit = input_freq * (1 - freq_dev/100);
	upper_limit = input_freq * (1 + freq_dev/100);
	if ( (fft_output->max_freq < lower_limit) || (fft_output->max_freq > upper_limit) )
		return -1;

	/* Check if volume is within deviation */
	lower_limit = input_vol * (1 - vol_dev/100);
	lower_limit = 20 * log10(lower_limit);
	upper_limit = input_vol * (1 + vol_dev/100);
	upper_limit = 20 * log10(upper_limit);
	if ( (fft_output->max < lower_limit) || (fft_output->max > upper_limit) )
		return -1;

	return 0;
}

int main(int argc, char **argv)
{
	FILE *fin = NULL;
	int16_t *input_buffer_i = NULL; /* contains interleaved data */
	int16_t *input_buffer_n = NULL; /* contains one channel of non-interleaved data */
	wave_header_t hdr;
	int ret = 0;
	int32_t nfft, offset, channel_index, num_input_samples, sample_index;
	scratch_t scratch = {NULL, NULL, NULL};
	output_t fft_output;
	double input_freq[MAX_NUM_CHANNELS], freq_dev[MAX_NUM_CHANNELS];
	double input_vol[MAX_NUM_CHANNELS], vol_dev[MAX_NUM_CHANNELS];

	/* Verify all required arguments are present */
	if (argc < 8) {
		printf("Usage: freq_analyzer <num FFT points> <overlap percent> <input file> <input freq> <%% freq deviation> <volume (dB)> <\%% volume deviation>\n");
		printf("                                                       channel 2:    <input freq> <%% freq deviation> <volume (dB)> <\%% volume deviation>\n");
		printf("                                                       ...........                                                                        \n");
		ret = -1;
		goto _done;
	}
	/* Number of FFT points */
	nfft = atoi(argv[1]);
	/* overlap */
	offset = ( (100 - atoi(argv[2])) * nfft) / 100;
	/* Read the samples from the input file into an array */
	fin = fopen(argv[3], "rb");
	if (!fin) {
		printf("Unable to open input file %s\n", argv[3]);
		ret = -1;
		goto _done;
	}
	/* input frequency */
	input_freq[0] = atof(argv[4]);
	/* Deviation in input frequency */
	freq_dev[0] = atof(argv[5]);
	/* Input volume */
	input_vol[0] = pow(10,atof(argv[6])/20);
	/* Deviation from input volume */
	vol_dev[0] = atof(argv[7]);
	/* Read the header & input samples */
	fread(&hdr, sizeof(wave_header_t), 1, fin);
	if (hdr.num_channels > MAX_NUM_CHANNELS) {
		printf("Only upto stereo supported\n");
		ret = -1;
		goto _done;
	}
	if ((hdr.num_channels == 2) && (argc < 12) ) {
		printf("Not enough arguments\n");
		ret = -1;
		goto _done;
	}
	if (hdr.num_channels == 2) {
		/* input frequency */
		input_freq[1] = atof(argv[8]);
		/* Deviation in input frequency */
		freq_dev[1] = atof(argv[9]);
		/* Input volume */
		input_vol[1] = pow(10,atof(argv[10])/20);
		/* Deviation from input volume */
		vol_dev[1] = atof(argv[11]);
	}
	num_input_samples = hdr.subchunk2_size/(hdr.num_channels * SAMPLE_BYTES);
	input_buffer_i = malloc(hdr.subchunk2_size);
	if (NULL == input_buffer_i) {
		printf("Unable to allocate space for input buffer\n");
		ret = -1;
		goto _done;
	}
	fread(input_buffer_i, 1, hdr.subchunk2_size, fin);
	/* Make necessary memory allocations */
	scratch.data_r = fftw_alloc_real(nfft);
	if (NULL == scratch.data_r) {
		printf("Unable to allocate space for input buffer\n");
		ret = -1;
		goto _done;
	}
	scratch.data_c = fftw_alloc_complex(nfft/2 + 1);
	if (NULL == scratch.data_c) {
		printf("Unable to allocate space for output buffer\n");
		ret = -1;
		goto _done;
	}
	scratch.mag = malloc(sizeof(double) * (nfft/2 + 1));
	if (NULL == scratch.mag) {
		printf("Unable to allocate space for output averaging buffer\n");
		ret = -1;
		goto _done;
	}

	/* FFT */
	if (1 == hdr.num_channels) {
		printf("Channel 0: ");
		/* Mono */
		execute_fft(input_buffer_i, nfft, offset, num_input_samples, hdr.sample_rate, &scratch, &fft_output);
		/* Print the output data */
		printf("%f (hz): %lf(dB)\n", fft_output.max_freq, fft_output.max);
		if (check_boundary(&fft_output, input_freq[0], freq_dev[0], input_vol[0], vol_dev[0]) == 0)
			printf("PASS\n");
		else
			printf("FAIL\n");
	}
	else {
		/* Processing for more than 1 channel */
		input_buffer_n = malloc(num_input_samples * SAMPLE_BYTES);
		if (NULL == input_buffer_n) {
			printf("Unable to allocate space for deinterleaving buffer\n");
			ret = -1;
			goto _done;
		}
		for(channel_index =  0; channel_index < hdr.num_channels; channel_index++) {
			/* Deinterleave and fft for one channel at a time */
			for(sample_index = 0; sample_index < num_input_samples; sample_index++)
				input_buffer_n[sample_index] = input_buffer_i[hdr.num_channels * sample_index + channel_index];
			/* FFT */
			printf("Channel %d: ", channel_index);
			execute_fft(input_buffer_n, nfft, offset, num_input_samples, hdr.sample_rate, &scratch, &fft_output);
			/* Print the output data */
			printf("%f (hz): %lf(dB)\n", fft_output.max_freq, fft_output.max);
			if (check_boundary(&fft_output, input_freq[channel_index], freq_dev[channel_index], input_vol[channel_index], vol_dev[channel_index]) == 0)
				printf("PASS\n");
			else
				printf("FAIL\n");
		}
	}
_done:
	if (fin) fclose(fin);
	if (input_buffer_i) free(input_buffer_i);
	if (input_buffer_n) free(input_buffer_n);
	if (scratch.data_r) fftw_free(scratch.data_r);
	if (scratch.data_c) fftw_free(scratch.data_c);
	if (scratch.mag) free(scratch.mag);
	return ret;
}