#include "nr3.h"
#include "fourier.h"
#include <fstream>

const int NSEC=128;
const int FSAMPLE=4096;
const int SIZE=NSEC*FSAMPLE;
const double lambda=-5.0;
const double timstep=1.0/FSAMPLE;
const double PI=4*atan(1);

void mult(VecDoub &in1, VecDoub &in2, VecDoub &out) {
// return A times B*
  for (int i=0; i<in1.size()/2; i++) {
     out[2*i]=in1[2*i]*in2[2*i]+in1[2*i+1]*in2[2*i+1];
     out[2*i+1]=in1[2*i+1]*in2[2*i]-in1[2*i]*in2[2*i+1];
  }
}

double Makefilter(VecDoub &filt, double f0) {
   double norm=0;
   for (int i=0; i<SIZE; i++) {
      double t=timstep*i;
      double sig=0;
      if (lambda*t>-30) sig=exp(lambda*t)*sin(2*PI*f0*t);
      filt[2*i]=sig;
      filt[2*i+1]=0;
      norm+=sig*sig;
   }
   norm=1.0/sqrt(norm);
   for (int i=0; i<SIZE; i++) filt[2*i]*=norm; // normalize filter to total power of 1
   return norm; // needed for amplitude signal.
}
   

int main() {
// define constants
   VecDoub data(2*SIZE); // sum of signal and noise
   VecDoub testdata(2*SIZE);
   VecDoub filter(2*SIZE); // the filters which which we will compare
   VecDoub prod(2*SIZE); // for the correlation

   ifstream datafile("./data.txt");
   ofstream logfile("./log");
   ofstream correlations("./correl");
   ofstream allcors("./allcors"); // give the overlap of the filter with the data for all lags
   ofstream filtercorrelation("./filtercorrelation"); // plot the overlap of the filters with the signal.

// read in the data
   double powdat=0;
   double powav=0;
   double Normdata=(SIZE*1.0);
   for (int i=0; i<SIZE;i++) {
      datafile >> data[2*i]; // real part sample i
      data[2*i+1]=0; // imag part sample i
      powdat+=data[2*i]*data[2*i];
      powav+=data[2*i];
   }
   logfile << " data read : average signal " << powav/SIZE << " RMS " << powdat/SIZE << endl;
   
// test normalization Fourier routine : is the FFT of the data and its inverse
// normalized?
   for (int i=0; i<2*SIZE; i++) testdata[i]=data[i];
   four1(testdata,1);
   for (int i=0; i<2*SIZE; i++) testdata[i]/=Normdata;
   four1(testdata,-1);
   for (int i=0; i<2*SIZE;i++) {
      if (fabs(data[i]-testdata[i])>1e-7) {
          cerr << "Normalization error! " << i << " " << data[i] << " " << testdata[i] << endl;
	  break;
      }
   }
// to test the procedure, make a test filter outside the region of the fiducial
// limits
   double f0=400;
   logfile << "Norm filter equals " << Makefilter(filter,f0); 
   logfile << " should be " << sqrt(lambda*2/FSAMPLE) << endl;
   for (int i=0; i<SIZE; i++) {
      double t=i*timstep;
      testdata[2*i]=data[2*i];
      if (t>25&&t<35) testdata[2*i]+=80*filter[2*(i-25*FSAMPLE)]; //add filter of 400 Hz with an offset of T0 = 25 seconds
   }

   four1(data,1);
   four1(testdata,1);
   four1(filter,1);
   for (int i=0; i<2*SIZE; i++) data[i]/=Normdata;
   for (int i=0; i<2*SIZE; i++) testdata[i]/=Normdata;
   mult(testdata,filter,prod);
   four1(prod,-1);
   logfile << " test of correlation : " << prod[50*FSAMPLE] << " should be 80 + ";
   mult(data,filter,prod);
   four1(prod,-1);
   logfile << prod[50*FSAMPLE] << endl;
   // second test: write out all correlations
   // should be Gaussian-distributed with RMS 1
   for (int i=0; i<SIZE; i++) allcors << prod[2*i] << endl;
   
// now start random search over a filter bank. The amplitudes for all
// correlations larger than 4 will be written to file
// also, the overlaps of the filters with the signal will be written to file
// for 1 filter, plot all correlations to show the Gaussian distribution

   double fstep=-lambda*0.1; // a resolution of 0.1 lambda in frequency
   double f=100.0;
   double maxcor=0;
   double maxt;
   double maxf=f;
   while (f<300+fstep) {
     // make a filter with this step
     double norm=Makefilter(filter,f);
     four1(filter,1);
// check correlation with data
     mult(data,filter,prod);
     four1(prod,-1);
     for (int i=0; i<SIZE; i++) {
        double t=timstep*i;
        if (prod[2*i]>4) correlations << t << " " << f << " " << prod[2*i] << " " << prod[2*i]*norm << endl;
        if (prod[2*i]*norm>maxcor) {
		maxcor=prod[2*i]*norm;
		maxt = t;
                maxf=f;
        }
     }
     f+=fstep;
  }
  logfile << "Highest signal found : A=  " << maxcor << endl;
  logfile << " at frequency " << maxf << " and";
  logfile << " time " << maxt << " " << endl;
  return 0;
}