#include <iostream>
#include <Riostream.h>

#include <TROOT.h>
#include <TObject.h>

#include <TApplication.h>
#include <TBranch.h>
#include <TCanvas.h>
#include <TChain.h>
#include <TClass.h>
#include <TF1.h>
#include <TFile.h>
#include <TH1.h>
#include <TH2.h>
#include <TMath.h>
#include <TTree.h>
#include <TVirtualFFT.h>

#include "TEpicData.h"
#include "TInitialConditions.h"

#define DEBUG 0 
#define z_bin_length           22.0 // [um]
#define Electron_Vdrift       110.0 // [um/ns]

#define Convolution_Gaussian_Sigma_ns     10
#define Convolution_Gaussian_Range_nSigma  3

//#define NameFile "P_890mbar"
#define NameFile "P_1bar"

void run()
{

  // === ==========
  // === input data
  // === ==========

  TChain * t = new TChain("SimulatedTree");
  for (int file=1; file<=10; file++){
    if (file != 3)
    //t->Add(Form("/media/audrey/DATA1/EPIC/simulation/%s/K11_240Pu_nf_GEF_part%i.root",NameFile,file));
    t->Add(Form("/media/audrey/DATA1/EPIC/simulation/%s/K11_240Pu_nf_GEF_part%i.root",NameFile,file));
  }
  t->ls();

  // === ===================
  // === Initialize the data
  // === ===================
  TClass *epicClass = TClass::GetClass("TEpicData");
  if (!epicClass) {
      std::cerr << "ERROR: TEpicData Class is not loaded" << std::endl;
      return;
  }
  TEpicData *m_Data = nullptr;
  t->SetBranchAddress("Epic", &m_Data);
 
  TClass * icClass = TClass::GetClass("TInitialConditions");
  if (!icClass) {
      std::cerr << "ERROR: TInitialConditions Class is not loaded" << std::endl;
      return;
  }
  TInitialConditions * m_ic = nullptr;
  t->SetBranchAddress("InitialConditions", &m_ic);

  // === =========
  // === Variables
  // === =========
  double time_bin_width_ns = z_bin_length / Electron_Vdrift ; 
  double time_KtoA_ns      = 2500 / Electron_Vdrift ; 

  int nbins_QvsT = 114;  
  int nbins_QvsT_external = 114;  
  int nbins_convol;
  int nbins_convol_external;


  // --- Gaussian Distribution for convolution
  double range_gauss       = Convolution_Gaussian_Sigma_ns * Convolution_Gaussian_Range_nSigma * 1000; // n sigma in ps
  int nbins_gauss  = (2*range_gauss) / (1000*time_bin_width_ns);  
  cout << "nbins_gauss = " << nbins_gauss << " to cover +/- " << Convolution_Gaussian_Range_nSigma << " sigma" << endl;
  vector<double> gaussian_distribution(nbins_gauss,0.);
  double normalization_factor = 0;
  for (int i = 0; i < nbins_gauss; i++) {
    gaussian_distribution[i] = TMath::Gaus(-range_gauss+i*time_bin_width_ns*1000, 0, Convolution_Gaussian_Sigma_ns*1000, kTRUE);
    normalization_factor += gaussian_distribution[i];
  }
  for(int i=0; i<nbins_gauss; i++) gaussian_distribution[i] /= normalization_factor;

  // === ============
  // === TTree output
  // === ============



  
  // === ==========
  // === Histograms
  // === ==========

  TH1I * h1_AnodeNumber = new TH1I("AnodeNumber","AnodeNumber",120,-1.5, 119.5);
  TH1I * h1_MultInternal = new TH1I("MultInternal","MultInternal",6,-1.5,4.5);
  TH1I * h1_MultExternal = new TH1I("MultExternal","MultExternal",6,-1.5,4.5);
  TH1I * h1_MultTot      = new TH1I("MultTot","MultTot",7,-1.5,5.5);
  TH2I * h2_MultTot_vs_MultInternal = new TH2I("MultTotal_vs_MultInternal","MultTotal_vs_MultInternal",6,-1.5,4.5,7,-1.5,5.5);
  TH2I * h2_MultTot_vs_MultExternal = new TH2I("MultTotal_vs_MultExternal","MultTotal_vs_MultExternal",6,-1.5,4.5,7,-1.5,5.5);

  TH1D * h1_QmaxInternal = new TH1D("QmaxInternalAnode_per_GEFevent","QmaxInternalAnode_per_GEFevent",100000,0,100000);
  TH1D * h1_QmaxInternal_m1 = new TH1D("QmaxInternalAnode_per_GEFevent_m1","QmaxInternalAnode_per_GEFevent_m1",100000,0,100000);
  TH1D * h1_QmaxExternal = new TH1D("QmaxExternalAnode_per_GEFevent","QmaxExternalAnode_per_GEFevent",100000,0,100000);
  TH2D * h2_Qmax_Internal_vs_External = new TH2D("Qmax_Internal_vs_External","Qmax_Internal_vs_External",1800,0,36000,1800,0,36000);
  h1_QmaxInternal_m1->SetLineColor(kBlack);

  TH1D * h1_ic_Z    = new TH1D("IC_Zpos","IC_Zpos",10000,-2.509,-2.508);
  //TH2D * h2_ic_YvsX = new TH2D("IC_YvsX","IC_YvsX",);

  TH1D * h1_ic_Z_LowQmaxInt_m1 = new TH1D("IC_Zpos_LowQmaxInt_m1","IZ_Zpos_LowQmaxInt",10000,-2.509,-2.508);

  //// === DRAW
#if DEBUG
  TCanvas * can1D_QvsT = new TCanvas("1D_QvsT","One Event",0,0,2500,2500);
  can1D_QvsT->Divide(1,2);
#endif

  // === ============================================================================
  // === Loop over the entries
  unsigned int nentries = t->GetEntries();
  std::cout << "   nentries = " << nentries << std::endl;
#if DEBUG
  for (unsigned int Entry = 201; Entry < 1000; Entry++) {
#else
    for (unsigned int Entry = 0; Entry < nentries; Entry++) {
      if((Entry%100000)==0) cout << "\r ===  Entry = " << Entry << " ===" << flush;
#endif
      int bytesRead = t->GetEntry(Entry);
      if (bytesRead <= 0) {
        std::cerr << "Error : GetEntry failed at Entry #" << Entry << std::endl;
        continue;
      }
      if (!m_Data) {
        std::cerr << "Error: m_Data is NULL at Entry #" << Entry << std::endl;
        continue;
      }
      if (!m_ic) {
        std::cerr << "Error: m_ic is NULL at Entry #" << Entry << std::endl;
        continue;
      }

#if DEBUG
      cout << endl ; 
      cout << "***********************************************" << endl;
      cout << "@ Entry = " << Entry << ", mult = " << m_Data->GetMultiplicity() << endl;
#endif


      int mult_internal = 0;
      int mult_external = 0;

      h1_MultTot->Fill(m_Data->GetMultiplicity());
      h1_ic_Z->Fill(m_ic->GetIncidentPositionZ());

      if(m_Data->GetMultiplicity()==0) {
        h1_MultInternal->Fill(0);
        h1_MultExternal->Fill(0);
      }
      else if(m_Data->GetMultiplicity()!=1) {
        nbins_QvsT = 0;
        nbins_QvsT_external = 0;
        for(unsigned short pmult=0; pmult<m_Data->GetMultiplicity(); pmult++){
          h1_AnodeNumber->Fill(m_Data->GetAnodeNbr(pmult));
          int nsteps = m_Data->GetNumberOfSteps(pmult);
          if(m_Data->GetAnodeNbr(pmult)<100){
            mult_internal++;
            if(nsteps > 114) cout << "Attention !!! nsteps_internal = " << nsteps << " larger than maximum 114 steps " << endl;
            for(int smult=0; smult<nsteps; smult++){
              int nbins_current = ceil(m_Data->GetTimeCreationElectronsPerStep(pmult).at(smult) / time_bin_width_ns) + 114;
              if(nbins_QvsT < nbins_current) nbins_QvsT = nbins_current;
            }
          }
          else{
            mult_external++;
            if(nsteps > 114) cout << "Attention !!! nsteps_external = " << nsteps << " larger than maximum 114 steps " << endl;
            for(int smult=0; smult<nsteps; smult++){
              int nbins_current = ceil(m_Data->GetTimeCreationElectronsPerStep(pmult).at(smult) / time_bin_width_ns) + 114;
              if(nbins_QvsT_external < nbins_current) nbins_QvsT_external = nbins_current;
            }
          }
        }
        nbins_convol   = nbins_QvsT + 2*nbins_gauss;
        vector<double> influence(nbins_convol,0);
        nbins_convol_external   = nbins_QvsT_external + 2*nbins_gauss;
        vector<double> influence_external(nbins_convol_external,0);

#if DEBUG
        TH1D * h1int_QvsT = new TH1D("InternalAnode_influence_vs_time_1D","InternalAnode_influence_vs_time_1D",nbins_convol,-2*range_gauss,nbins_QvsT*1000*time_bin_width_ns+2*range_gauss);
        h1int_QvsT->GetXaxis()->SetTitle("Absolute time [ps] 200ps/bin"); 
        h1int_QvsT->SetLineColor(kBlue);
        h1int_QvsT->SetDirectory(0);
        TH1D * h1int_convolution_QvsT = new TH1D("InternalAnode_QvsT_convol","InternalAnode_QvsT_convol",nbins_convol,-2*range_gauss,nbins_QvsT*1000*time_bin_width_ns+2*range_gauss);
        h1int_convolution_QvsT->SetLineColor(kRed);

        TH1D * h1ext_QvsT = new TH1D("ExternalAnode_influence_vs_time_1D","ExternalAnode_influence_vs_time_1D",nbins_convol_external,-2*range_gauss,nbins_QvsT_external*1000*time_bin_width_ns+2*range_gauss);
        h1ext_QvsT->GetXaxis()->SetTitle("Absolute time [ps] 200ps/bin"); 
        h1ext_QvsT->SetLineColor(kBlue);
        h1ext_QvsT->SetDirectory(0);
        TH1D * h1ext_convolution_QvsT = new TH1D("ExternalAnode_QvsT_convol","ExternalAnode_QvsT_convol",nbins_convol_external,-2*range_gauss,nbins_QvsT_external*1000*time_bin_width_ns+2*range_gauss);
        h1ext_convolution_QvsT->SetLineColor(kRed);
#endif

        for(unsigned short pmult=0; pmult<m_Data->GetMultiplicity(); pmult++){
          int nsteps  = m_Data->GetNumberOfSteps(pmult);
          for(int smult=0; smult<nsteps; smult++){
            double  t_Step = m_Data->GetTimeCreationElectronsPerStep(pmult).at(smult) ; // [ns NB. Vff >> Vdrift]
            int     n_Step = m_Data->GetNumElectronsPerStep(pmult).at(smult);
            int     b_Step = t_Step/time_bin_width_ns; 
            for(int bin=smult; bin<114; bin++){
              if(m_Data->GetAnodeNbr(pmult)<100){
                if(nbins_QvsT < b_Step+bin-smult) {
                  cout << "=== nbins_QvsT = " << nbins_QvsT << " but need to put data at index " << b_Step+bin-smult << endl;
                }
                influence[b_Step+bin-smult+nbins_gauss] += n_Step * 22. / 2500. ;
              }
              else{
                if(nbins_QvsT_external < b_Step+bin-smult) {
                  cout << "=== nbins_QvsT_external = " << nbins_QvsT_external << " but need to put data at index " << b_Step+bin-smult << endl;
                }
                influence_external[b_Step+bin-smult+nbins_gauss] += n_Step * 22. / 2500. ;
              }
            }// end of loop over the steps up to anodes  
          }// end of loop over the steps
        }// end of loop over the particle multiplicity (if 2 FF)

        double Qmax_Internal = 0;
        for(int i=0; i<nbins_QvsT+nbins_gauss; i++){
          double convolution = 0;
          for (int j = 0; j < nbins_gauss; j++) convolution += influence[i+j] * gaussian_distribution[j]; 
#if DEBUG
          h1int_QvsT->SetBinContent(i+1,influence[i]);
          h1int_QvsT->SetBinError(i+1,0);
          h1int_convolution_QvsT->SetBinContent(i+1+0.5*nbins_gauss,convolution);
#endif
          if (convolution>Qmax_Internal) Qmax_Internal=convolution;
        }
        h1_QmaxInternal->Fill(Qmax_Internal);
        influence.clear();

        // if data @ external proceed
        if(nbins_QvsT_external>0){
          double Qmax = 0;
          for(int i=0; i<nbins_QvsT_external+nbins_gauss; i++){
            double convolution = 0;
            for (int j = 0; j < nbins_gauss; j++) convolution += influence_external[i+j] * gaussian_distribution[j]; 
#if DEBUG
            h1ext_QvsT->SetBinContent(i+1,influence_external[i]);
            h1ext_QvsT->SetBinError(i+1,0);
            h1ext_convolution_QvsT->SetBinContent(i+1+0.5*nbins_gauss,convolution);
#endif
            if (convolution>Qmax) Qmax=convolution;
          }
          h1_QmaxExternal->Fill(Qmax);
          h2_Qmax_Internal_vs_External->Fill(Qmax,Qmax_Internal);
        }
        influence_external.clear();
#if DEBUG
        can1D_QvsT->cd(1); h1int_QvsT->Draw(); h1int_convolution_QvsT->Draw("same");
        can1D_QvsT->cd(2); h1ext_QvsT->Draw(); h1ext_convolution_QvsT->Draw("same");

        can1D_QvsT->Update();
        can1D_QvsT->WaitPrimitive();

        h1int_QvsT->Reset();  h1int_convolution_QvsT->Reset(); delete h1int_QvsT;  delete h1int_convolution_QvsT;
        h1ext_QvsT->Reset();  h1ext_convolution_QvsT->Reset(); delete h1ext_QvsT;  delete h1ext_convolution_QvsT;
#endif
      }// end of if else mult>1

      else{
        // if mult==1, only internal anode
        int    nsteps  = m_Data->GetNumberOfSteps(0);
        if(nsteps > 114) cout << "Attention !!! nsteps = " << nsteps << " larger than maximum 114 steps " << endl;
        nbins_QvsT = 114;
        for(int smult=0; smult<nsteps; smult++){
          int nbins_current = ceil(m_Data->GetTimeCreationElectronsPerStep(0).at(smult) / time_bin_width_ns) + 114;
          if(nbins_QvsT < nbins_current) nbins_QvsT = nbins_current;
        }
        nbins_convol   = nbins_QvsT + 2*nbins_gauss;
        vector<double> influence(nbins_convol,0);
#if DEBUG
        TH1D * h1_QvsT = new TH1D("Mult1_InternalAnode_influence_vs_time_1D","Mult1_InternalAnode_influence_vs_time_1D",nbins_convol,-2*range_gauss,nbins_QvsT*1000*time_bin_width_ns+2*range_gauss);
        h1_QvsT->GetXaxis()->SetTitle("Absolute time [ps] 200ps/bin"); 
        h1_QvsT->SetLineColor(kBlue);
        h1_QvsT->SetDirectory(0);
        TH1D * h1_convolution_QvsT = new TH1D("Mult1_InternalAnode_QvsT_convol","Mult1_InternalAnode_QvsT_convol",nbins_convol,-2*range_gauss,nbins_QvsT*1000*time_bin_width_ns+2*range_gauss);
        h1_convolution_QvsT->SetLineColor(kRed);
#endif
        for(int smult=0; smult<nsteps; smult++){
          double  t_Step = m_Data->GetTimeCreationElectronsPerStep(0).at(smult) ; // [ns] NB. Vff >> Vdrift
          int     n_Step = m_Data->GetNumElectronsPerStep(0).at(smult);
          int     b_Step = ceil(t_Step/time_bin_width_ns);
          for(int bin=smult; bin<114; bin++){
            if(nbins_QvsT < b_Step+bin-smult) {
              cout << "--- nbins_QvsT = " << nbins_QvsT << " but need to put data at index " << b_Step+bin-smult << endl;
            }
            else
              influence[b_Step+bin-smult+nbins_gauss] += n_Step * 22. / 2500. ;
          }// end of loop over the steps up to anodes  
        }// end of loop over the steps
        double Qmax = 0;
        for(int i=0; i<nbins_QvsT+nbins_gauss; i++){
          double convolution = 0;
          for (int j = 0; j < nbins_gauss; j++) convolution += influence[i+j] * gaussian_distribution[j]; 
#if DEBUG
          h1_QvsT->SetBinContent(i+1,influence[i]);
          h1_QvsT->SetBinError(i+1,0);
          h1_convolution_QvsT->SetBinContent(i+1+0.5*nbins_gauss,convolution);
#endif  
          if (convolution>Qmax) Qmax=convolution;
        }
        h1_AnodeNumber->Fill(m_Data->GetAnodeNbr(0));
        if(m_Data->GetAnodeNbr(0)<100){
          mult_internal++;
          h1_QmaxInternal->Fill(Qmax);
          h1_QmaxInternal_m1->Fill(Qmax);
          if (Qmax<3000) h1_ic_Z_LowQmaxInt_m1->Fill(m_ic->GetIncidentPositionZ());
        }
        else{
          cout << "****** ******* ******** WARNING: MULT_TOT = 1 (" << m_Data->GetMultiplicity() << ") : INTERNAL ANODE ONLY BUT ANODE #" << m_Data->GetAnodeNbr(0) << endl;
          mult_external++;
        }
        influence.clear();
#if DEBUG
        can1D_QvsT->cd(1); h1_QvsT->Draw(); h1_convolution_QvsT->Draw("same");

        can1D_QvsT->Update();
        can1D_QvsT->WaitPrimitive();

        h1_QvsT->Reset();             delete h1_QvsT;
        h1_convolution_QvsT->Reset(); delete h1_convolution_QvsT;
#endif
      }// end of else mult=1
      h1_MultInternal->Fill(mult_internal);
      h1_MultExternal->Fill(mult_external);
      h2_MultTot_vs_MultInternal->Fill(mult_internal,m_Data->GetMultiplicity()); 
      h2_MultTot_vs_MultExternal->Fill(mult_external,m_Data->GetMultiplicity()); 
    }
    cout << endl;
    // === End of loop over the entries
    // === ============================================================================


    TCanvas * can_mult = new TCanvas("mult","mult",0,0,2000,2000);
    can_mult->Divide(3,2);
    can_mult->cd(1); gPad->SetLogy(); h1_AnodeNumber->Draw();
    can_mult->cd(2); gPad->SetLogy(); h1_MultInternal->Draw();
    can_mult->cd(3); gPad->SetLogy(); h1_MultExternal->Draw();
    can_mult->cd(4); gPad->SetLogy(); h1_MultTot->Draw();
  can_mult->cd(5); h2_MultTot_vs_MultInternal->Draw("COL");
  can_mult->cd(6); h2_MultTot_vs_MultExternal->Draw("COL");

  TCanvas * can_Qmax = new TCanvas("Qmax","Qmax",0,0,2000,1500);
  can_Qmax->Divide(3,1);
  can_Qmax->cd(1);  h1_QmaxInternal->Draw(); h1_QmaxInternal_m1->Draw("sames");
  can_Qmax->cd(2);  h1_QmaxExternal->Draw();
  can_Qmax->cd(3);  h2_Qmax_Internal_vs_External->Draw("COL");





  TFile * fsave = new TFile(Form("/media/audrey/DATA1/EPIC/simulation/%s/out_K11_GEF240Pu.root",NameFile),"recreate");
  can_Qmax->Write();
  h1_QmaxInternal->Write();
  h1_QmaxExternal->Write();
  h2_Qmax_Internal_vs_External->Write();

  can_mult->Write();
  h1_AnodeNumber->Write();
  h1_MultInternal->Write();
  h1_MultExternal->Write();
  h1_MultTot->Write();
  h2_MultTot_vs_MultInternal->Write();
  h2_MultTot_vs_MultExternal->Write();

  fsave->Close();


}