/*

  FakePhysicist.com
  admin@fakephysicist.com
  28 May 2017

  compile with:
  g++ fft.cc -o fft `root-config --cflags --libs`

*/



#include <iostream>
#include <cmath>
#include <new>
#include <vector>
#include <complex>

#include "TPad.h"
#include "TAxis.h"
#include "TCanvas.h"
#include "TGraph.h"
#include "TMath.h"

using namespace std;
using namespace TMath;

typedef std::complex<double> Complex;
typedef std::vector<Complex> CNArray;


/********** Cooley–Tukey FFT **********/

void fft(CNArray &x) {
  
  const Long64_t N = x.size();
  if (N <= 1) return;

  // divide
  CNArray even, odd;
  even.clear(); odd.clear();
  for (Long64_t k = 0; k < N; k += 2) {
    if(k+1 == N) {break;}
    even.push_back(x[k]);
    odd.push_back(x[k+1]);
  }
  
  // compute
  fft(even);
  fft(odd);

  // combine
  for (Long64_t k = 0; k < N/2; k++) {
    Complex t = std::polar(1.0, - 2. * Pi() * k / N) * odd[k];
    x[k    ] = even[k] + t;
    x[k+N/2] = even[k] - t;
  }  
};


/********* inverse fft *********/

void ifft(CNArray& x) {

  const Long64_t N = x.size();

  // conjugate the complex numbers
  for (Long64_t k=0;k<N;k++) {
    x[k] = std::conj(x[k]);
  }

  // forward fft
  fft(x);

  // conjugate the complex numbers again
  for (Long64_t k=0;k<N;k++) {
    x[k] = std::conj(x[k]);
  }

  // scaling
  for (Long64_t k=0;k<N;k++) {
    x[k] /= N;
  }
};



const double tPD = 1.;	// Total Path Difference in sec
const Long64_t maxGenPt = 1000;	// Number of Data Points to Generate
const Long64_t maxDataPt = pow(2,18); // Number of Padded Data Points
const double maxFrequency = 5000.; // Maximum Frequency to Calculate


int main() {

  const Long64_t minPt = 2.*(maxFrequency*tPD+1); // N_min
  const Long64_t inputDataPt = max(maxGenPt,minPt);
  
  cout << " Generating Data...." << endl;

  vector<double> pos0,amp0;	// Input Data Vector before Padding
  pos0.clear();
  amp0.clear();
  
  /****** Generating Data Points ********/
  for(int ij=0;ij<inputDataPt;ij++) {
    double tt = ij*tPD/inputDataPt;
    double amp = 5.*tt		// Addition of 1000, 1500, 2000,
      + sin(2.*Pi()*1000.*tt)	// 3000 and 4000 Hz Sine Wave
      + sin(2.*Pi()*1500.*tt)
      + sin(2.*Pi()*2000.*tt)
      + sin(2.*Pi()*3000.*tt)
      + sin(2.*Pi()*4000.*tt);
    pos0.push_back(tt);
    amp0.push_back(amp);
  }
  
  cout << " Input Data Size " << inputDataPt << endl;

  const Long64_t maxDesiredPt = max(inputDataPt,maxDataPt);

  /******** Finding the suitable 2^p value *********/
  int pFactor = 0;
  double breakPt = 1.*inputDataPt/2.;
  while(1) {
    if(1.*maxDesiredPt/pow(2.,pFactor)>=1.) {
      pFactor++;
    } else {break;}
  }

  const Long64_t nn = pow(2,pFactor); // Number for Padding
  cout << " Size of FFT Input " << nn << endl;
  
  vector<double> outPos,outAmp;	// Vector for Padded Data
  outPos.clear();outAmp.clear();

  /******* Padding the Data *********/
  double startPos,endPos,
    startAmp,endAmp;
  int ampCnt = 0;
  double interVal = (pos0.back()-pos0.front())/(nn-1);
  
  for(Long64_t ij=0;ij<nn;) {
    startPos = pos0[ampCnt];
    endPos = pos0[ampCnt+1];
    startAmp = amp0[ampCnt];
    endAmp = amp0[ampCnt+1];
      
    double newPos = pos0.front() + ij*interVal;
    
    if(newPos > endPos) {
      ampCnt++;
    } else {
      outPos.push_back(newPos);
      outAmp.push_back(startAmp + (newPos - startPos)*(endAmp - startAmp)/(endPos - startPos));
      ij++;
    }
  }

  CNArray data_C;		// vector for Complex number
  data_C.clear();
  for(Long64_t ij=0;ij<nn;ij++) {
    Complex tempComp(outAmp[ij],0.); // Setting the Real Part
    data_C.push_back(tempComp);
  }

  /****** d-FFT *********/
  cout << " DFFT Starts...." << endl;
  fft(data_C);
  cout << " DFFT Ends...." << endl;
  
  vector<double> cBin,cAmp;
  cBin.clear();cAmp.clear();
  for(Long64_t ij=0;ij<nn;ij++) {
    cBin.push_back(ij/tPD);
    cAmp.push_back(abs(data_C[ij]));
  }

  /****** Inverse d-FFT *********/
  CNArray data_I = data_C;
  cout << " iDFFT Starts...." << endl;
  ifft(data_I);
  cout << " iDFFT Ends...." << endl;

  vector<double> iBin,iAmp;
  iBin.clear();iAmp.clear();
  for(Long64_t ij=0;ij<nn;ij++) {
    iBin.push_back(ij*tPD/(nn-1.));
    iAmp.push_back(real(data_I[ij]));
  }
    
  TCanvas *c1 = new TCanvas("c1","c1",2400,2400);
  c1->Divide(1,4);

  TGraph* h3 = new TGraph(inputDataPt,&pos0[0],&amp0[0]); // Gen
  TGraph* h0 = new TGraph(nn,&outPos[0],&outAmp[0]);	  // Pad
  TGraph* h1 = new TGraph(int(nn/2),&cBin[0],&cAmp[0]);	  // d-FFT
  TGraph* h2 = new TGraph(nn,&iBin[0],&iAmp[0]);	  // id-FFT
  

  c1->cd(1);
  h3->Draw("AL*");
  c1->cd(2);
  h0->Draw("AL");
  c1->cd(3);
  h1->Draw("AL*");
  gPad->SetLogy(1);
  h1->GetXaxis()->SetRangeUser(500., 5000.);
  // h1->GetXaxis()->SetRangeUser(1995., 2005.);
  c1->cd(4);
  h2->Draw("AL");
  
  c1->SaveAs("fft.jpg");
  
  return 0;
  
}