// FFTseries.cc
// Reads a time series of a given variable from the qweak root tree Hel_Tree
// and computes the FFT.  The variable maxfrequency sets the frequency scale.
//
//
// Richard Jones, Qweak experiment
// November 20, 2010
//
//
// Usage:  root $QW_ROOTFILES/first100k_7359.root
//   root [1]  new TBrowser;  // browse down to Hel_Tree inside root file
//   root [2]  .L FFTseries.cc
//   root [3]  TH1* fft = FFTseries(Hel_Tree,"yield_qwk_bpm3h07cY","hw_sum");
//   root [4]  fft->SetAxisRange(1.,80.) // zoom in to see 60Hz peak
//   root [5]  fft->Draw();
// To notch out a single coefficient in the fft at bin 89, do
//   root [6]  quench(89);
// To notch out a range in the fft between bins 89 and 92, do
//   root [7]  for(int i=89; i<=92; ++i) { quench(i); }
// To scan a range in the fft and notch out anything above a certain height, do
//   root [8]  shave(height,89,92);
// To histogram the helicity sums use the project function
//   root [9]  TH1D *ts = gROOT->FindObject("tseries");
//   root [9]  project(ts,1,100,-1.e-3,1.e-3)


#include <iostream>
#include <algorithm>
#include <float.h>
#include "TTree.h"
#include "TBranch.h"
#include "TLeaf.h"
#include "TH1D.h"

TH1* FFTseries(TTree* tree, char* branchname, char* varname)
{
   double maxfrequency = 960.; // Hz
   if (strcmp(tree->GetName(),"Hel_Tree")==0) {
      maxfrequency /= 4;
   }
   else if (strcmp(tree->GetName(),"Mps_Tree")==0) {
      // continue
   }
   else {
      cout << "Error - sampling frequency unknown for tree "
           << tree->GetName() << endl;
   }
   int events_to_skip=3000;
   int numsamples = tree->GetEntries();
   numsamples -= events_to_skip;
   tree->ResetBranchAddresses();
   TH1D* ts = new TH1D("tseries",branchname,numsamples,0.,numsamples);
   TH1D* ta = new TH1D("tampl","tampl",numsamples,0,numsamples);

   double hel;
   tree->SetBranchAddress("actual_helicity",&hel);

   double var;
   if (strlen(branchname) > 0) {
      TBranch* br0 = tree->GetBranch(branchname);
      if (br0 == 0) {
         cout << "Error - tree contains no branch called " << branchname
              << endl;
         return ts;
      }
      TLeaf* br1 = br0->FindLeaf(varname);
      if (br1 == 0) {
         cout << "Error - branch " << branchname
              << " contains no leaf element called " << varname << endl;
         return ts;
      }
      br1->SetAddress(&var);
   }
   else {
      tree->SetBranchAddress(varname,&var);
   }

   int samp = events_to_skip;
   for (int i=0; i<numsamples; i++) {
      tree->GetEntry(samp++);
      hel = (hel == 1)? +1 : (hel == 0)? -1 : 0;
      ta->SetBinContent(i+1,hel);
      ts->SetBinContent(i+1,var);
   }

   char* title = Form("FFT of %s",branchname);
   TH1D* ffts = new TH1D("fftseries",title,numsamples,0.,maxfrequency);
   ts->FFT(ffts,"MAG");
   //ffts->Rebin();
   ffts->SetAxisRange(1.,maxfrequency/2);
   ffts->Draw();
   return ffts;
}

void quench(int bin)
{
   TH1* ts = (TH1*)gROOT->FindObject("tseries");
   TH1* ffts = (TH1*)gROOT->FindObject("fftseries");
   double maxtime = ts->GetXaxis()->GetXmax();
   double frequency = ffts->GetXaxis()->GetXmax();
   int N = ffts->GetNbinsX();

   TH1* fftr = new TH1D("fftreal","fftreal",N,0,frequency);
   TH1* ffti = new TH1D("fftimag","fftimag",N,0,frequency);
   ts->FFT(fftr,"RE");
   ts->FFT(ffti,"IM");
   double Cr = fftr->GetBinContent(bin);
   double Ci = ffti->GetBinContent(bin);

   cout << "FFT bin " << bin << " has amplitude ("
        << Cr << "," << Ci << ")" << endl;

   TH1* subtr = new TH1D("subtr","subtractr",N,0,N);
   TH1* subti = new TH1D("subti","subtracti",N,0,N);
   TH1* artif = new TH1D("artif","artifact",N,0,frequency);
   artif->SetBinContent(bin,Cr+Ci);
   artif->SetBinContent(N+2-bin,Cr-Ci);
   artif->FFT(subtr,"RE");
   artif->FFT(subti,"IM");
   ts->Add(subtr,-1./N);
   ts->Add(subti,-1./N);
   ffts->Reset();
   ts->FFT(ffts,"MAG");
   fftr->Reset();
   ts->FFT(fftr,"RE");
   ffti->Reset();
   ts->FFT(ffti,"IM");
   double Cpr = fftr->GetBinContent(bin);
   double Cpi = ffti->GetBinContent(bin);

   cout << "FFT bin " << bin << " has amplitude ("
        << Cpr << "," << Cpi << ")" << endl;

   ffts->SetAxisRange(1.,frequency/2);
   ffts->Draw();

   delete fftr;
   delete ffti;
   delete subtr;
   delete subti;
   delete artif;
}

void shave(double height, int minbin=1, int maxbin=99999999, TH1* ffts=0)
{
   if (ffts == 0) {
      ffts = (TH1*)gROOT->FindObject("fftseries");
   }
   int N = ffts->GetNbinsX();
   for (int i=minbin; i<N-minbin; ++i) {
      if (i > maxbin) {
         break;
      }
      if (ffts->GetBinContent(i+1) > height) {
         quench(i+1);
      }
   }
}

double signal()
{
   TH1* ta = (TH1*)gROOT->FindObject("tampl");
   int N = ta->GetNbinsX();
   double sbar=0;
   for (int i=0; i<N; ++i) {
      double val = ta->GetBinContent(i+1);
      sbar += fabs(val);
   }
   return sbar/N;
}

double error()
{
   TH1* ta = (TH1*)gROOT->FindObject("tampl");
   int N = ta->GetNbinsX();
   double sbar=0;
   for (int i=0; i<N; ++i) {
      double val = ta->GetBinContent(i+1);
      sbar += val*val;
   }
   return sqrt(sbar/N);
}

void project(TH1* hist, double normfact=1, int nbins=100, double x0=-FLT_MAX, double x1=+FLT_MAX)
{
   double xmax = hist->GetMaximum();
   double xmin = hist->GetMinimum();
   xmax = x0;
   xmin = x1;
   cout << xmax << "   " << xmin << "   " << normfact << endl;
   proj = new TH1D("proj",hist->GetTitle(),nbins,xmin,xmax);
   int N = hist->GetNbinsX();
   int rebin_factor=4;//100;
   double sum=0;
   double amp=0;
   TH1* amph = gROOT->FindObject("tampl");
   for (int i=0; i<N; ++i) {
      amp = amph->GetBinContent(i+1);
      sum += amp * hist->GetBinContent(i+1);
      if ((i+1)%rebin_factor == 0) {
         proj->Fill(sum/rebin_factor/normfact);
         sum=0;
      }
   }
   proj->Draw();
}