#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <afm_fft2_ff.h>
#include <gr_io_signature.h>
#include <math.h>

#ifndef min
#define min(X,Y) ((X) < (Y) ? : (X) , (Y))
#endif
#ifndef max
#define max(X,Y) ((X) > (Y) ? : (X) , (Y))
#endif


afm_fft2_ff_sptr afm_make_fft2_ff(int nfft, int naverages, float sampling_freq, float fundamental_freq, const std::vector<float> &freqs,const std::vector<float> &window)
{
    return afm_fft2_ff_sptr (new afm_fft2_ff(nfft,naverages,sampling_freq,fundamental_freq,freqs,window));
}

afm_fft2_ff::afm_fft2_ff(int nfft, int naverages, float sampling_freq, float fundamental_freq, const std::vector<float> &freqs, const std::vector<float> &window)
    : gr_sync_decimator("fft2_ff", 
	    gr_make_io_signature(3, 3, sizeof(float)),
	    gr_make_io_signature(1, 1, sizeof(float)),
	    (int) ceil(naverages*sampling_freq/(fundamental_freq*(2+freqs.size() ))) )
{
    d_nfft = nfft;
    d_fft_input = (double *) malloc(sizeof(double)*nfft);
    d_fft_out = (fftw_complex *) fftw_malloc(sizeof(fftw_complex) * (nfft/2 +1));
    d_plan = fftw_plan_dft_r2c_1d(nfft,d_fft_input, d_fft_out,FFTW_MEASURE); 

    d_naverages = naverages;
    d_window = window;
    d_freqs = freqs;
    DPRINTF(("Freq vector size = %d\n", d_freqs.size())); 
    d_sampling_freq = sampling_freq;
    d_fundamental_freq = fundamental_freq;
    d_bufferlen = (int) ceil(naverages*sampling_freq/fundamental_freq);
    d_buffer.reserve(d_bufferlen);
    d_buffer.clear();
    d_runonce = false;

    set_output_multiple((2+freqs.size()));
    //X,Y,freqs

    unsigned int i, index;
    for (i=0;i<d_freqs.size();i++){
	index = (int)round(d_freqs[i]*d_nfft/d_sampling_freq);
	DPRINTF(("Index for frequency%d = %d\n",i,index));
	try {
	    if (index > d_nfft/2)
		throw "Frequency must be less that 1/2 Sampling Frequency";
	}
	catch( char * str ) {
	    printf("Exception raised: %s\n",str);
	    exit(1);
	}
	d_freqs_index.push_back(index);
    }

    try {
	if(window.size() != nfft)
	    throw "Window size and nfft should be the same";
    }
    catch( char * str) {
	printf("Exception raised: %s\n",str);
	exit(1);
    }

    if( d_bufferlen > nfft)
	printf("\nWarning: Buffer length = %d, FFT length = %d. Signal will be truncated",d_bufferlen,nfft);

}

afm_fft2_ff::~afm_fft2_ff()
{
    fftw_destroy_plan(d_plan);
    fftw_free(d_fft_out);
    free(d_fft_input);
}

int afm_fft2_ff:: work(int noutput_items,
	gr_vector_const_void_star &input_items,
	gr_vector_void_star &output_items)

{
    /* chan[0] is signal
     * chan[1] X imaging signal
     * chan[2] Y imaging signal
     */
    int i, j, noutput_samples = 0; 
    const float * chan[3]; 
    float ** out = (float **) &output_items[0];
    for(i=0;i<3;i++)
        chan[i] = (const float *) input_items[i]; 

    if (d_runonce) {
	d_upordown = (chan[2][0] > 0) ? AFM_UP : AFM_DOWN;
	d_leftorright = (chan[1][0] > 0) ? AFM_LEFT : AFM_RIGHT;
	d_runonce = true;
    }

    for(i=0;i< noutput_items/(2+d_freqs.size());i++) {
	//stuff the buffer
	for(j=0;j<d_bufferlen;j++) 
	    d_buffer.push_back(chan[0][i*d_bufferlen + j]);
	
	//X imaging voltage
	if(d_leftorright == AFM_LEFT) {
	    if (chan[1][i*d_bufferlen] < 0)
		d_leftorright = AFM_RIGHT;
	}
	else {
	    if (chan[1][i*d_bufferlen] > 0.02)
		d_leftorright = AFM_LEFT;
	}
	out[0][noutput_samples] = (d_leftorright == AFM_LEFT) ? 1.0 : -1.0;
	noutput_samples++;

	//Y imaging voltage
	if(d_upordown == AFM_UP) {
	    if (chan[2][i*d_bufferlen] < 0)
		d_upordown = AFM_DOWN;
	}
	else {
	    if (chan[2][i*d_bufferlen] > 0.05)
		d_upordown = AFM_UP;
	}
	out[0][noutput_samples] = (d_upordown == AFM_DOWN) ? 1.0 : -1.0;
	noutput_samples++;

	noutput_samples = compute_params(out,noutput_samples);
	
	d_buffer.clear();
    }
    
    DPRINTF(("Outputs requested: %d, Output produced: %d, Input consumption(expected): %d, Input consumption(actual): %d\n",noutput_items,noutput_samples,
    (int) ceil(d_naverages*d_sampling_freq/(d_fundamental_freq*(2+d_freqs.size()))) * noutput_samples,
    noutput_items*d_bufferlen/(2+d_freqs.size())
    ));
    
    return noutput_samples;
}


int afm_fft2_ff:: compute_params(float ** out,int noutput_samples) {
    unsigned int i,index;
    float temp;
    if (d_bufferlen < d_nfft) {
	for (i=0;i<d_bufferlen;i++) 
	    d_fft_input[i] = (double) (d_buffer[i]*d_window[i]);
	for (;i<d_nfft;i++) 
	    d_fft_input[i] = 0; //zero-pad
   } 
    else {
	for (i=0;i<d_nfft;i++) 
	    d_fft_input[i] = (double) (d_buffer[i]*d_window[i]); //truncate after d_nfft
    }

    fftw_execute(d_plan);
    for(i=0;i<d_freqs.size();i++) {
	index = d_freqs_index[i];
	temp = (float) sqrt(d_fft_out[index][0]*d_fft_out[index][0] + d_fft_out[index][1]*d_fft_out[index][1]);
	out[0][noutput_samples] = temp;
	noutput_samples++;
    }
    return noutput_samples;
}