#define peachCake_cxx

#include "peachCake.h"
#include <TH2.h>
#include <TH1.h>
#include <TMath.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TCutG.h>

#include "./armory/AnalysisLibrary.h"

//============== User settings

int tofRange[3] = {500, -280, -210}; 
int dERange[3] = {500, 0, 5500};


//============== histograms

TH2F * hCloverID;
TH2F * hPID;
TH1F * hZ;
TH2F * hPID0; // raw
TH2F * hPID2; // z vs A/Q

TH1F * hTOF;
TH1F * hdE;

TH2F * hImplantIons;
TH2F * hImplantBeta; //high gain

TH2F * hImplantX;
TH2F * hImplantY;

TH2F * hImplantDyIons;
TH2F * hImplantDyBeta;

TH1F * hDecay;
TH1F * hDecay_veto;

TH1F * hFlag;
TH1F * hvetoFlag;

TH2F * heVrun;
TH2F * htVrun;

//============ parameters
//TString cutFileName="PIDCuts.root";
//TFile * cutFile;
//TCutG * cut = NULL;
//TObjArray * cutList;
//int numCut = 0;

vector<vector<double>> eCorr;

double TOFCorrection(double x){
   
  //for run 238 
  //return (-225.8 - 0.040017 * energy 
  //             +0.0000710839 * energy*energy 
  //             -6.28402e-8 * TMath::Power(energy,3) 
  //             +2.90563e-11 * TMath::Power(energy,4) 
  //             -6.70137e-15 * TMath::Power(energy,5)  
  //             +6.08418e-19 * TMath::Power(energy,6) );
   
  // for run 250
  double par[6] = {10.9179, 
                   0.00416034,
                   -233.306,
                   3607.16,
                   1538.37,
                   0.000430609};
  return par[0] * exp( - par[1] * x ) + par[2] - par[5] * (par[3] - x) * exp( - pow( (par[3]-x)/par[4],2) ); 
  
}

TH2F * createTH2F(const char* name, const char* title, Int_t nbinsx, Double_t xlow, Double_t xup, Int_t nbinsy, Double_t ylow, Double_t yup){
  
  TH2F * hist2 = (TH2F *) gROOT->FindObjectAny( name );
  if ( hist2 == NULL ) hist2 = new TH2F( name , title , nbinsx, xlow, xup, nbinsy, ylow, yup);
  hist2->Reset();
  
  return hist2;
}

TH1F * createTH1F(const char* name, const char* title, Int_t nbinsx, Double_t xlow, Double_t xup){
  
  TH1F * hist1 = (TH1F *) gROOT->FindObjectAny( name );
  if ( hist1 == NULL ) hist1 = new TH1F( name , title , nbinsx, xlow, xup);
  hist1->Reset();
  
  return hist1;
}

void peachCake::Begin(TTree * /*tree*/){
  
  TString option = GetOption();

  hCloverID = createTH2F("hCloverID", "Clover; ID; energy [keV]", 52, 0, 52, 400, 0, 2000);

  hPID = createTH2F("hPID", "PID; ns; msx100", tofRange[0], tofRange[1], tofRange[2], dERange[0], dERange[1], dERange[2]);   
  hPID0 = createTH2F("hPID0", "PID raw; ns; msx100", tofRange[0], tofRange[1], tofRange[2], dERange[0], dERange[1], dERange[2]);
  hPID2 = createTH2F("hPID2", "PID2; A/Q; Z", tofRange[0], 2.4, 3.6, dERange[0], 0, 16.5);

  hTOF = createTH1F("hTOF", "TOF", tofRange[0], tofRange[1], tofRange[2]);
  hdE = createTH1F("hdE", "dE", dERange[0], dERange[1], dERange[2]);

  hZ = createTH1F("hZ", "Z; Z; count", 500, 0, 16);

  hImplantIons = createTH2F("hImplantIons", "Implat low gain; x ; y", 400, 0, 1, 400, 0, 1);
  hImplantBeta = createTH2F("hImplantBeta", "Implat high gain; x ; y", 400, 0, 1, 400, 0, 1);   

  hImplantX = createTH2F("hImplantX", "Implat X; low ; high", 400, 0, 1, 400, 0, 1);
  hImplantY = createTH2F("hImplantY", "Implat Y; low ; high", 400, 0, 1, 400, 0, 1);   


  hImplantDyIons = createTH2F("hImplantDyIonsow", "Implat low; sum_a; dy", 400, 0, 80000, 400, 0, 80000);
  hImplantDyBeta = createTH2F("hImplantDyBetaigh", "Implat high; sum_a; dy", 400, 0, 8000, 400, 0, 8000);

  hDecay = createTH1F("hDecay", "Decay (beta.time - ion.time) ; [ticks]; counts", 100, 0, 2000);
  hDecay_veto = createTH1F("hDecay_veto", "Decay (beta.time - ion.time) [veto]; [ticks]; counts", 100, 0, 2000);

  hFlag = createTH1F("hFlag", "Flag ( 1 = beam, 2 = Ions, 4 = Beta)", 8, 0, 8);
  hvetoFlag = createTH1F("hvetoFlag", "vetoFlag ( 1 = front, 4 = rear)", 8, 0, 8);

  heVrun = createTH2F("heVrun", "energy vs run; run; energy", 200, 237, 257, 400, 0, 2000);
  htVrun = createTH2F("hrVrun", "TOF vs run; run; energy", 200, 237, 257, 400, -265, -220);

  eCorr = LoadCorrectionParameters("correction_e.dat");
  
  if( pidCorrFileName != "" ){
    pidCorr = LoadCorrectionParameters(pidCorrFileName, 1);
  }

  /**
  cutFile = new TFile(cutFileName, "READ");
  bool listExist = cutFile->GetListOfKeys()->Contains("cutList");
  if( listExist ) {
    cutList = (TObjArray*) cutFile->FindObjectAny("cutList");
    numCut = cutList->GetLast()+1;
    printf("----- found %d cuts \n", numCut);
    for( int k = 0; k < numCut; k++){
       if( cutList->At(k) != NULL ){
          printf("found a cut at %2d \n", k);
       }
    }
  }
  * */
   
  printf("============ start \n");

}


Bool_t peachCake::Process(Long64_t entry){

  nProcessed++;

  b_multi->GetEntry(entry);
  b_detID->GetEntry(entry);
  b_e->GetEntry(entry);
  b_e_t->GetEntry(entry);
  b_cfd->GetEntry(entry);
  b_runID->GetEntry(entry);

  energy = TMath::QuietNaN();
  Long64_t startTimeL = 0;
  Long64_t startTimeR = 0;
  Long64_t stopTime = 0;
  double cfdL = TMath::QuietNaN();
  double cfdR = TMath::QuietNaN();
  double cfd2 = TMath::QuietNaN();

  Long64_t startTime40 = 0;
  double cfd40 = TMath::QuietNaN();

  double a_L[5] = {TMath::QuietNaN()};  ///0 = dy, 1-4 = anode
  double a_H[5] = {TMath::QuietNaN()};

  veto_f = TMath::QuietNaN();
  veto_r = TMath::QuietNaN();

  veto_f_time = 0;
  veto_r_time = 0;

  crossEnergy = 0; /// for analog signal, cross T2
  crossTime = 0;

  for( int i = 0; i < 4 ; i++) {
   dyIonsTime[i] = 0;
   dyBetaTime[i] = 0;
   dyIonsEnergy[i] = 0;
   dyBetaEnergy[i] = 0;
  }

  int multiDyBeta = 0;
  int multiDyIons = 0;

  flag = 0;
  vetoFlag = 0;

  //======= Scanning the event
  for( int i = 0; i < multi; i++){
    //----- clover
    if( 0 <= detID[i] && detID[i] < 100 ){ 
      if( e[i] > 0 ) {
        int id = detID[i];
        double eCal = 0;
        for( int k = 0; k < int(eCorr[id].size()); k++){
          eCal += eCorr[id][k] * TMath::Power(e[i], k);
        }
        hCloverID->Fill(detID[i], eCal);
      }
    }
    //----- dy Beta
    if( detID[i] == 200 ) { 
      if ( multiDyBeta < 4 ){
        dyBetaTime[multiDyBeta] = e_t[i]; 
        dyBetaEnergy[multiDyBeta] = e[i];
        multiDyBeta ++;
      }
    }
    if( 201 <= detID[i] && detID[i] < 250){ 
      a_H[detID[i] - 200] = e[i];
    }
    //----- dy Ions
    if( detID[i] == 250 ) { 
      if ( multiDyIons < 4 ){
        dyIonsTime[multiDyIons] = e_t[i]; 
        dyIonsEnergy[multiDyIons] = e[i]; 
        multiDyIons ++;
      }
    }
    if( 251 <= detID[i] && detID[i] < 260){
      a_L[detID[i] - 250] = e[i];
    }
    //----- veto_front 
    if( detID[i] == 290 ) { 
      veto_f = e[i];
      veto_f_time = e_t[i];
      if( veto_f > 0 && (vetoFlag & 1) == 0) vetoFlag += 1;
    }
    //----- veto_rear
    if( detID[i] == 291 ) { 
      veto_r = e[i];
      veto_r_time = e_t[i];
      if( veto_r > 0 && (vetoFlag & 2) == 0) vetoFlag += 2;
    }
    //----- image scint L
    if( detID[i] == 300 ) { 
      if( e[i] > 1000 ) {
        startTimeL = e_t[i];
        cfdL = cfd[i]/16384.;
      }
    }
    //----- image scint R
    if( detID[i] == 301 ) { 
      startTimeR = e_t[i];
      cfdR = cfd[i]/16384.;
    }
    //----- cross T2 analog
    if( detID[i] == 306 ) { 
      if( (vetoFlag & 4) == 0) vetoFlag += 4;
      crossEnergy = e[i];
      crossTime = e_t[i];
    }
    //----- cross B2 logic
    if( detID[i] == 307 ) { 
      stopTime = e_t[i];
      cfd2 = cfd[i]/16384.;
      if( (vetoFlag & 4) == 0) vetoFlag += 4;
    }
    //----- MSX40
    ///if( detID[i] == 308 ) energy = e[i]; 
    //----- MSX100
    if( detID[i] == 309 ) energy = e[i]; 
    //----- MSX40 logic
    if( detID[i] == 310 ) {  
      startTime40 = e_t[i];
      cfd40 = cfd[i]/16384.;
    }
  }

  //================================== PID
  Long64_t startTime = startTimeL;
  double startCFD = cfdL;

  TOF = TMath::QuietNaN();
  if( startTime > 0 && stopTime > 0 && energy > 0 ){
    if( stopTime > startTime ){
      TOF = double(stopTime - startTime) - startCFD + cfd2;
    }else{
      TOF = -double(startTime - stopTime) - startCFD + cfd2 ;
    }
    TOF = TOF * 8; // to ns
    
    ///============== PID correction
    if( pidCorrFileName != "" ){
      for( int k = 0; k < (int) pidCorr.size() ; k++){
        if( int(pidCorr[k][0]) != runID/100 ) continue;
        
        TOF = pidCorr[k][1] + pidCorr[k][2] * TOF;
        energy = pidCorr[k][3] + pidCorr[k][4] * energy;
        
      }
    }
    
  
    /// TOF correction, A/Q = 3 to have a fixed TOF.
    
    if( energy > 160./44.*(TOF + 260) +180 ) {
      hPID0->Fill(TOF, energy);
      hTOF->Fill(TOF);
      hdE->Fill(energy);
    }

    double newTOF = TOF - TOFCorrection(energy) - 240.;

    hPID->Fill(newTOF, energy);
    
    heVrun->Fill(runID/100., energy);
    htVrun->Fill(runID/100., TOF);

    flag += 1;

    ///======== Z vs A/Q
    double c = 299.792458; /// mm/ns
    double FL = 37656 ;  /// mm  
    double Brho = 4.90370; /// T.mm   
    double ma = 931.5;
    double beta0 = Brho * c / TMath::Sqrt( 9*ma*ma + Brho*Brho*c*c); // for AoQ = 3
    double tofOffset = 509.79; /// ns

    double beta = FL/c/(newTOF+tofOffset);
    double gamma = 1./TMath::Sqrt(1.-beta*beta);

    double Z = sqrt((energy + 11.9473) / 23.097) * TMath::Power(beta / beta0,1.3)  ;
    double AoQ = c * Brho / gamma/ beta / ma;

    //printf("tof : %f, beta: %f (%f), Z : %f, A/Q : %f \n", TOF + tofOffset, beta, beta0, Z, AoQ);

    //Z = -0.15938 + 1.01736 *Z + 0.000203316 * Z*Z; 
    //Z = -0.254632 + 1.06285 * Z - 0.00539634 * Z * Z + 0.000169443 * Z * Z * Z; ///2022-06-09

    hPID2->Fill(AoQ, Z);

    hZ->Fill(Z);
   
  }

  //================================== Implant
  double sumA_L = 0;
  double sumA_H = 0;
  for( int i = 1; i <= 4 ; i++){
    if( a_L[i] > 0 ) sumA_L += a_L[i];
    if( a_H[i] > 0 ) sumA_H += a_H[i];
  }

  xIons = (a_L[3]+a_L[2])/sumA_L;
  yIons = (a_L[1]+a_L[2])/sumA_L;
  xBeta = (a_H[3]+a_H[2])/sumA_H;
  yBeta = (a_H[1]+a_H[2])/sumA_H;


  if( (flag & 1) == 0 ) {
    hImplantIons->Fill(xIons, yIons);
    hImplantBeta->Fill(xBeta, yBeta);
    hImplantDyIons->Fill(sumA_L, dyIonsEnergy[0]);
    hImplantDyBeta->Fill(sumA_H, dyBetaEnergy[0]);

    hImplantX->Fill(xIons, xBeta);
    hImplantY->Fill(yIons, yBeta);
  }

  if( !TMath::IsNaN(veto_f) && !TMath::IsNaN(veto_r) && crossTime == 0 ){
    //hImplantIons->Fill(xIons, yIons);
    //hImplantBeta->Fill(xBeta, yBeta);
    //hImplantDyIons->Fill(sumA_L, dyIonsEnergy[0]);
    //hImplantDyBeta->Fill(sumA_H, dyBetaEnergy[0]);
    //
    //hImplantX->Fill(xIons, xBeta);
    //hImplantY->Fill(yIons, yBeta);
  }

  if( 0 < xIons && xIons < 1 && 0 < yIons && yIons < 1 ) flag += 2;
  if( 0 < xBeta && xBeta < 1 && 0 < yBeta && yBeta < 1 ) flag += 4;

  hFlag->Fill(flag);
  hvetoFlag->Fill(vetoFlag);


  //============================== debug
  if ( false  ) {

    printf("################# %lld, multi:%d\n", entry, multi);

    printf("----- beam Line \n");
    printf(" MSX100 eneergy : %f \n", energy);
    printf(" startTime : %llu, CFD : %f \n", startTime, startCFD);
    printf("  stopTime : %llu, CFD : %f \n", stopTime, cfd2);
    printf("       TOF : %f ns \n", TOF);

    printf("----- Ions, low gain \n");
    for( int i = 1; i <= 4; i++ )printf("a%d : %f\n", i, a_L[i]);
    printf("sum A : %f \n", sumA_L);
    printf("x : %f , y : %f \n", xIons, yIons);
    for ( int i = 0; i < 4; i++ )printf("dy : %u,  %llu \n", dyIonsEnergy[i], dyIonsTime[i]);

    printf("----- Beta, high gain \n");
    for( int i = 1; i <= 4; i++ )printf("a%d : %f\n", i, a_H[i]);
    printf("sum A : %f \n", sumA_H);
    printf("x : %f , y : %f \n", xBeta, yBeta);
    for ( int i = 0; i < 4; i++ )printf("dy : %u, %llu \n", dyBetaEnergy[i], dyBetaTime[i]);

    printf("----- Veto \n");
    printf("   cross Time : %llu \n", crossTime);
    printf(" cross energy : %d \n", crossEnergy);
    printf("veto_f energy : %f \n", veto_f);
    printf(  "veto_f Time : %llu \n", veto_f_time);
    printf("veto_r energy : %f \n", veto_r);
    printf(  "veto_r Time : %llu \n", veto_r_time);

    printf("============ flag : %d,  vetoFlag : %d\n", flag, vetoFlag);

  }


  //============================= Progress
  clock.Stop("timer");
  Double_t time = clock.GetRealTime("timer");
  clock.Start("timer");

  if ( !shown ) {
    if (fmod(time, 10) < 1 ){
      printf( "%12llu[%2d%%]|%3.0f min %5.2f sec | expect:%5.2f min \n", 
             nProcessed, 
             TMath::Nint((nProcessed+1)*100./totnumEntry),
             TMath::Floor(time/60.), time - TMath::Floor(time/60.)*60.,
             totnumEntry*time/(nProcessed+1.)/60.);
             shown = 1;}
  }else{
    if (fmod(time, 10) > 9 ){
       shown = 0;
    }
  }

  //============================= Rejection
  if ( flag == 0 ) return kTRUE;

  //============================= Save Tree
  if( saveNewTree ){
    saveFile->cd();
    newTree->Fill();
  }

  return kTRUE;
}


void peachCake::Terminate(){

   printf("\n===============\n");

  if( saveNewTree ){
   saveFile->cd();
   newTree->Write();
   saveFile->Close();
  }

  if( plotHists ){

    gStyle->SetOptStat(111111);
    int div[2] = {3, 2} ; ///x, y
    TCanvas * c1 = new TCanvas("c1", "c1", 700 * div[0], 700 * div[1]);
    c1->Divide(div[0], div[1]);


    int padID=1;

    c1->cd(padID); //c1->cd(padID)->SetLogz();
    hPID0->Draw("colz");

    //padID ++;
    //c1->cd(padID); //c1->cd(padID)->SetLogz();
    //hTOF->Draw();
    //
    //padID ++;
    //c1->cd(padID); //c1->cd(padID)->SetLogz();
    //hdE->Draw();

    padID ++;
    c1->cd(padID); //c1->cd(padID)->SetLogz();
    hPID->Draw("colz");

    padID ++;
    c1->cd(padID); //c1->cd(padID)->SetLogz();
    hPID2->Draw("colz");

    padID ++;
    c1->cd(padID); //c1->cd(padID)->SetLogz();
    hZ->Draw("colz");

    padID ++;
    c1->cd(padID); //c1->cd(padID)->SetLogz();
    heVrun->Draw("colz");

    padID ++;
    c1->cd(padID); //c1->cd(padID)->SetLogz();
    htVrun->Draw("colz");

    //padID ++;
    //c1->cd(padID); c1->cd(padID)->SetGrid(); c1->cd(padID)->SetLogz();
    //hCloverID->Draw("colz"); hCloverID->SetNdivisions(-13);

/*
   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantIons->Draw("colz");

   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantBeta->Draw("colz");

   //padID ++;
   //c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantBeta_veto->Draw("colz");

   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantDyIons->Draw("colz");

   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantDyBeta->Draw("colz");
   
   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantX->Draw("colz");
   
   padID ++;
   c1->cd(padID); c1->cd(padID)->SetLogz(); hImplantY->Draw("colz");
   
   padID ++;
   c1->cd(padID); hFlag->Draw(""); 
   hvetoFlag->SetLineColor(2); hvetoFlag->Draw("same");
      
   */
   
  }
   
  //================ Save historgram 
  if( fHistRootName != "" ){
    TFile * fHist = new TFile(fHistRootName, "recreate");
    fHist->cd();
    
    hTOF->Write();
    hdE->Write();
    hPID0->Write();

    fHist->Close();
    
    printf("---- Save PID histogram as %s \n", fHistRootName.Data());
  }
  
}