#include <iomanip>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include "TAxis.h"
#include <stdlib.h>
#include "TInterpreter.h"
#include "TSystem.h"
#include "TSystemDirectory.h" 
#include "TSystemFile.h"
#include "TCanvas.h"
#include "TH1.h"
#include "TVirtualFFT.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"

int main();

std::vector<TString> m_fnames;
TString m_dir = "/eos/home-e/egkougko/Root/Boron/CNM_10478_W4S1067_1e14_p/";
TString m_file;

std::vector<TFile *> input;
std::vector<TTree *> imputtree;
std::vector<TH1D *> fftS;
std::vector<TH1D *> fftN;
std::vector<TH1D *> fftSM;
std::vector<TH1D *> fftNM;
std::vector<TCanvas *> c1;
std::vector<TAxis *> xAxisS;
std::vector<TAxis *> xAxisN;

const int nbinsS = 100000;
const int nbinsN = 100000;

int main()
{
    // Count the number of files in the idrectory
    unsigned int nspl = 0;
    std::vector <TString> dirs;
    dirs.push_back(m_dir);
    unsigned int hep = dirs.size();
    for (unsigned int i = 0; i < hep; i++)
        {
         TSystemDirectory *directory = new TSystemDirectory("", dirs.at(i));
         TList *files = directory->GetListOfFiles();
         if (files)
            {
             TSystemFile *sfile;
             TIter next(files);
             TString fname;
             while ((sfile=(TSystemFile*)next()))
                   {
                    fname = sfile->GetName();
                    if (!sfile->IsDirectory())
                       {
                        if (fname.EndsWith(".root")) 
                           {
                            std::cout << "Root file: " << dirs.at(i) + fname << std::endl;
                            nspl++;
                            m_fnames.push_back(dirs.at(i) + fname);
                           }
                       }
                    else {
                          if (fname != "." && fname !="..")
                             {
                              dirs.push_back(dirs.at(i) + fname + "/");
                              // std::cout << "Direcotry: " << dirs.at(i) + fname + "/" << std::endl;
                              hep = hep + 1;
                             }
                         }
                   }
            }
        }

   std::cout << nspl << std::endl;
   input.resize(nspl);
   imputtree.resize(nspl);

   fftS.resize(nspl);
   fftN.resize(nspl);
   fftSM.resize(nspl);
   fftNM.resize(nspl);
   c1.resize(nspl);
   xAxisS.resize(nspl);
   xAxisN.resize(nspl);

   std::vector<double> xbinsS;
   std::vector<double> ybinsS;
   std::vector<double> xbinsN;
   std::vector<double> ybinsN;

   for (unsigned int z = 0; z < nspl; z++)
       {
        // file does not exist in local directory
        if (gSystem->AccessPathName(m_fnames[z])) std::cout << "--> Cannot Access input file!!!" << std::endl;
        input[z] = TFile::Open(m_fnames[z]);
        m_file = m_fnames[z];
        m_file = TString( m_file (m_dir.Length() + 1, ((m_fnames[z]).Length() - m_dir.Length() + 1)));
        m_file = TString( m_file (0, m_file.Length() - 5));
        TBranch        *b_SignalFFT02;   //!
        TBranch        *b_NoiseFFT02;    //!
        imputtree[z] = (TTree*)input[z]->Get("wfm");
        double treevars[2];
        imputtree[z]->SetBranchAddress("SignalFFT02", &(treevars[0]), &b_SignalFFT02);
        imputtree[z]->SetBranchAddress("NoiseFFT02", &(treevars[1]), &b_NoiseFFT02);
        fftS[z] = new TH1D((TString("fftS_") + m_file).Data(), (TString("fftS_") + m_file).Data(), nbinsS + 1, 0, 1.0e+9);
        fftN[z] = new TH1D((TString("fftN_") + m_file).Data(), (TString("fftN_") + m_file).Data(), nbinsN + 1, 0, 1.0e+9);
        for (Long64_t ievt = 0; ievt<imputtree[z]->GetEntries(); ievt++)
            {
             b_SignalFFT02->GetEntry(ievt);
             b_NoiseFFT02->GetEntry(ievt);
             if (treevars[0] > 0) fftS[z]->Fill(abs(treevars[0]));
             if (treevars[1] > 0) fftN[z]->Fill(abs(treevars[1]));
            }
        std::cout << m_fnames[z] << " " << fftS[z]->GetEntries() << std::endl;

        // Signal Generation
        xbinsS.clear();
        ybinsS.clear();
        xAxisS[z] = fftS[z]->GetXaxis();
        xbinsS.push_back(xAxisS[z]->GetBinLowEdge(1));
        unsigned int i = 1;
        unsigned int j = 1;
        for (i = 1; i < nbinsS; i++)
            {
             // std::cout << "----------------------------------------------------------" << std::endl;
             // std::cout << "Processing data Bin: " << i << " " << fftS[z]->GetBinContent(i) << std::endl;
             for (j = i + 1; j < nbinsS+1; j++)
                 {
                  // std::cout << " j testing: " << j << " entries: " << fftS[z]->GetBinContent(j) << " lower edge: " << xAxis[z]->GetBinLowEdge(j) << std::endl;
                  if (fftS[z]->GetBinContent(j) != 0)
                     {
                      xbinsS.push_back(xAxisS[z]->GetBinLowEdge(j));
                      ybinsS.push_back(fftS[z]->GetBinContent(i));
                      // std::cout << "----------------------------------------------------------" << std::endl;
                      // std::cout << "Wrote something " << i << " " << xAxisS[z]->GetBinLowEdge(j) << " " << fftS[z]->GetBinContent(i) << std::endl;
                      i = j-1;
                      break;
                     }
                 else {
                       if (j < nbinsS) continue;
                       else {
                             xbinsS.push_back(xAxisS[z]->GetBinLowEdge(i+1));
                             ybinsS.push_back(0);
                             // std::cout << "END WRITING " << i << " " << xAxisS[z]->GetBinLowEdge(j) << " " << fftS[z]->GetBinContent(i) << std::endl;
                             i = nbinsS+1;
                            }
                       }
                 }
            }

        std::cout <<"Signal: " <<  xbinsS.size() << " " << ybinsS.size() << std::endl;
        const unsigned int gz = xbinsS.size();
        double *binsS = new double[gz];
        for (unsigned int k = 0; k < gz; k++) binsS[k] = xbinsS.at(k);
        fftSM[z] = new TH1D((TString("fftSM_") + m_file).Data(), (TString("fftSM_") + m_file).Data(), gz - 1, binsS);
        for (unsigned int e = 0; e < gz - 1; e++) fftSM[z]->SetBinContent(e, ybinsS.at(e));
        delete[] binsS;

        // Noise generation
        xbinsN.clear();
        ybinsN.clear();
        xAxisN[z] = fftN[z]->GetXaxis();
        xbinsN.push_back(xAxisN[z]->GetBinLowEdge(1));
        i = 1;
        j = 1;
        for (i = 1; i < nbinsN; i++)
            {
             // std::cout << "----------------------------------------------------------" << std::endl;
             // std::cout << "Processing data Bin: " << i << " " << fftS[z]->GetBinContent(i) << std::endl;
             for (j = i + 1; j < nbinsN+1; j++)
                 {
                  // std::cout << " j testing: " << j << " entries: " << fftS[z]->GetBinContent(j) << " lower edge: " << xAxis[z]->GetBinLowEdge(j) << std::endl;
                  if (fftN[z]->GetBinContent(j) != 0)
                     {
                      xbinsN.push_back(xAxisN[z]->GetBinLowEdge(j));
                      ybinsN.push_back(fftN[z]->GetBinContent(i));
                      i = j-1;
                      // std::cout << "----------------------------------------------------------" << std::endl;
                      // std::cout << "Wrote something " << i << " " << xAxis[z]->GetBinLowEdge(j) << " " << fftS[z]->GetBinContent(i) << std::endl;
                     break;
                    }
                 else {
                       if (j < nbinsN) continue;
                       else {
                             xbinsN.push_back(xAxisN[z]->GetBinLowEdge(i+1));
                             ybinsN.push_back(0);
                             // std::cout << "END WRITING " << i << " " << xAxis[z]->GetBinLowEdge(j) << " " << fftS[z]->GetBinContent(i) << std::endl;
                             i = nbinsN+1;
                            }
                      }
                 }
            }

        std::cout << "Noise: " << xbinsN.size() << " " << ybinsN.size() << std::endl;
        const unsigned int sz = xbinsN.size();
        double *binsN = new double[gz];
        for (unsigned int a = 0; a < sz; a++) binsN[a] = xbinsN.at(a);
        fftNM[z] = new TH1D((TString("fftNM_") + m_file).Data(), (TString("fftNM_") + m_file).Data(), sz - 1, binsN);
        for (unsigned int w = 0; w < sz - 1; w++) fftNM[z]->SetBinContent(w, ybinsN.at(w));
        delete[] binsN;

        c1[z] = new TCanvas((TString("c_") + m_file).Data(), (TString("FFT_") + m_file).Data());
        c1[z]->Divide(2, 2, 0.004, 0.004, 0);
        c1[z]->cd(1);
        fftS[z]->SetLineColor(4);
        fftS[z]->SetLineWidth(2);
        fftS[z]->SetStats(0);
        fftS[z]->Draw("");
        xAxisS[z]->SetTitle("Frequency (Hz)");
        fftS[z]->GetYaxis()->SetTitle("N. of events");
        fftS[z]->GetYaxis()->SetTitleOffset(1.5);
        fftS[z]->GetYaxis()->SetLabelSize(0.03);
        fftS[z]->SetTitle((TString("Signal FFT Original - ") + m_file).Data());
        c1[z]->cd(2);
        fftN[z]->SetLineColor(4);
        fftN[z]->SetLineWidth(2);
        fftN[z]->SetStats(0);
        fftN[z]->Draw("");
        xAxisN[z]->SetTitle("Frequency (Hz)");
        fftN[z]->GetYaxis()->SetTitle("N. of events");
        fftN[z]->GetYaxis()->SetTitleOffset(1.5);
        fftN[z]->GetYaxis()->SetLabelSize(0.03);
        fftN[z]->SetTitle((TString("Noise FFT Original - ") + m_file).Data());
        c1[z]->cd(3);
        fftSM[z]->SetLineColor(8);
        fftSM[z]->SetLineWidth(2);
        fftSM[z]->SetStats(0);
        fftSM[z]->Draw("C");
        fftSM[z]->GetXaxis()->SetTitle("Frequency (Hz)");
        fftSM[z]->GetYaxis()->SetTitle("N. of events");
        fftSM[z]->GetYaxis()->SetTitleOffset(1.5);
        fftSM[z]->GetYaxis()->SetLabelSize(0.03);
        fftSM[z]->SetTitle((TString("Signal FFT Variable Binning - ") + m_file).Data());
        c1[z]->cd(4);
        fftNM[z]->SetLineColor(8);
        fftNM[z]->SetLineWidth(2);
        fftNM[z]->SetStats(0);
        fftNM[z]->Draw("");
        fftNM[z]->GetXaxis()->SetTitle("Frequency (Hz)");
        fftNM[z]->GetYaxis()->SetTitle("N. of events");
        fftNM[z]->GetYaxis()->SetTitleOffset(1.5);
        fftNM[z]->GetYaxis()->SetLabelSize(0.03);
        fftNM[z]->SetTitle((TString("Signal FFT Variable Binning - ") + m_file).Data());
        c1[z]->Update();
       }

   return 0;
}