#define DecayFinder_cxx

#include "DecayFinder.h"
#include <TH2.h>
#include <TH1.h>
#include <TMath.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TCutG.h>
#include <TObjArray.h>
#include <TRandom.h>

//===================== Settings
Bool_t debug = false;

TString pidCutFileName = ""; //"PIDCuts.root";

double distThreshold = 0.1; 
double forwardTime = 100; //ms
double backwardTime = 50; //ms

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

TH2F * hPID;
TH1F ** hDecay ;
TH2F * hDist; 
TH1F ** hvetoF;

TH2F * hIonsPos;
TH2F * hBetaPos;

//===================== Parameters

int tick2ns = 4;
double tick2ms = tick2ns / 1e6;

TFile * cutFile;
TCutG * cut = NULL;
TObjArray * cutList;
int numCut = 0;

int numMatch = 0;

std::vector<std::vector<int>> isotopes = {
                            {19,  7}, {20,  7}, {21,  7},
                            {22,  8}, {23,  8}, {24,  8},
                  {24,  9}, {25,  9}, {26,  9}, {27,  9},           {29,  9},
                  {27, 10}, {28, 10}, {29, 10}, {30, 10}, {31, 10}, {32, 10},
        {29, 11}, {30, 11}, {31, 11}, {32, 11}, {33, 11}, {34, 11}, {35, 11},
        {32, 12}, {33, 12}, {34, 12}, {35, 12}, {36, 12}, {37, 12}, {38, 12},
        {35, 13}, {36, 13}, {37, 13}, {38, 13}, {39, 13}, {40, 13}, {41, 13},
                  {39, 14}, {40, 14}, {41, 14}, {42, 14}, 
                            {43, 15}, {44, 15}
                          }; 

std::string isoSym(int Z){
  switch (Z) {
    case  3 : return "Li"; break;
    case  4 : return "Be"; break;
    case  5 : return "B"; break;
    case  6 : return "C"; break;
    case  7 : return "N"; break;
    case  8 : return "O"; break;
    case  9 : return "F"; break;
    case 10 : return "Ne"; break;
    case 11 : return "Na"; break;
    case 12 : return "Mg"; break;
    case 13 : return "Al"; break;
    case 14 : return "Si"; break;
    case 15 : return "P"; break;
    case 16 : return "S"; break;
  }
  return "XX";
}

//===================== Begins
void DecayFinder::Begin(TTree * /*tree*/){

  TString option = GetOption();

  hPID = new TH2F("hPID", "PID (beam & Implant & back veto); A/Q; Z", 400, 2.2, 4.0, 400, 3, 16);
  
  hDist = new TH2F("hDist", "dist; dx; dy", 100, -distThreshold, distThreshold, 100, -distThreshold, distThreshold);
  
  hIonsPos = new TH2F("hIonsPos", "Ions Pos; x; y", 200, 0, 1, 200, 0, 1);
  hBetaPos = new TH2F("hBetaPos", "Beta Pos; x; y", 200, 0, 1, 200, 0, 1);
  
  //-------------- GetCut;
  cutFile = new TFile(pidCutFileName, "READ");
  bool listExist = cutFile->GetListOfKeys()->Contains("cutList");
  if( !listExist ) {
    cutList = new TObjArray();
    numCut = (int) isotopes.size();
    
    printf("============= no Cut found, creating new default Cuts\n");
    
    for( int i = 0 ; i < numCut ; i++){
      
      cut = new TCutG();
      
      printf("%2d | A: %2d, Z: %d \n", i,  isotopes[i][0], isotopes[i][1]); 
        
      TString name; name.Form("cut%02d%s", isotopes[i][0], isoSym(isotopes[i][1]).c_str());
      cut->SetName(name);
      cut->SetVarX("AoQ");
      cut->SetVarY("Z");
      cut->SetTitle(Form("cut%2d", i));
      cut->SetLineColor(i+1);
      
      cut->SetPoint(0, (isotopes[i][0] - 0.45)/isotopes[i][1], isotopes[i][1] - 0.45);
      cut->SetPoint(1, (isotopes[i][0] - 0.45)/isotopes[i][1], isotopes[i][1] + 0.45);
      cut->SetPoint(2, (isotopes[i][0] + 0.45)/isotopes[i][1], isotopes[i][1] + 0.45);
      cut->SetPoint(3, (isotopes[i][0] + 0.45)/isotopes[i][1], isotopes[i][1] - 0.45);
      cut->SetPoint(4, (isotopes[i][0] - 0.45)/isotopes[i][1], isotopes[i][1] - 0.45);

      cutList->AddAtAndExpand(cut, i);
    }
  }else{
    cutList = (TObjArray*) cutFile->FindObjectAny("cutList");
    numCut = cutList->GetLast()+1;
    printf("------------ found %d cuts  in %s \n", numCut, pidCutFileName.Data());
    for( int k = 0; k < numCut; k++){
      if( cutList->At(k) != NULL ){
        printf("found a cut at %2d \n", k);
      }else{
        printf("     No cut at %2d \n", k);
      }
    }
  }

  hDecay = new TH1F *[numCut];
  hvetoF = new TH1F *[numCut];
  
  for( int i = 0; i < numCut; i++){
    hDecay[i] = new TH1F(Form("hDecay%02d%s",isotopes[i][0], isoSym(isotopes[i][1]).c_str()), Form("Decay cut-%02d ; [ms]; count", i), 200, -backwardTime, forwardTime); 
    hvetoF[i] = new TH1F(Form("hvetoF%02d%s",isotopes[i][0], isoSym(isotopes[i][1]).c_str()), Form("veto-F cut-%02d ; [ms]; count", i), 100, 0, 6000); 
  }
  
  printf("=====================================\n");
  
}

double dist(double x1, double x2, double y1, double y2){
  return sqrt((x1-x2)*(x1-x2) + (y1-y2)*(y1-y2));
}

//===================== Process
Bool_t DecayFinder::Process(Long64_t entry){
  
  //if( entry > (Long64_t)1000 ) return kTRUE;

  b_Z->GetEntry(entry);
  b_AoQ->GetEntry(entry);
  b_crossTime->GetEntry(entry);
  b_crossEnergy->GetEntry(entry);
  b_flag->GetEntry(entry);
  b_vetoFlag->GetEntry(entry);
  b_xIons->GetEntry(entry);
  b_yIons->GetEntry(entry);
  b_dyIonsTime->GetEntry(entry);
  b_veto_r->GetEntry(entry);
  b_veto_f->GetEntry(entry);

  ULong64_t ImplantEntry = 0;
  ULong64_t ImplantTime = 0;
  double ImplantX = TMath::QuietNaN();
  double ImplantY = TMath::QuietNaN();
  
  if( debug && entry % 10 == 0 ) printf("------------- %llu\n", entry);
  
  if( (flag == 3) && (vetoFlag & 2) == 0 ) hPID->Fill(AoQ, Z); // with Beam + Ions, and has veto_rear
  
  int cutID = -1;
  for(int i = 0; i < numCut; i++){
    cut = (TCutG*) cutList->At(i);
    if( cut->IsInside(AoQ, Z) ) cutID = i;
  }
  if( cutID == -1 ) return kTRUE;

  if( flag >= 4 ) return kTRUE; /// has beta event
  ///if( flag != 3 ) return kTRUE; ///

  //if( (flag & 1) == 0 ) return kTRUE; /// no beam
  //if( (flag & 2) == 0 ) return kTRUE; /// no Ions
  //if( (flag & 4) == 4 ) return kTRUE; /// has only beta

  //if( veto_f > 0 ) return kTRUE;
  //if( veto_r > 0 ) return kTRUE;
  //if( dyIonsTime[0] == 0 ) return kTRUE;

  ImplantEntry = entry;
  ImplantTime = dyIonsTime[0];
  ImplantX = xIons;
  ImplantY = yIons;
  
  if( TMath::IsNaN(xIons) ) return kTRUE;
  if( TMath::IsNaN(yIons) ) return kTRUE;
  
  hIonsPos->Fill( xIons, yIons);
    
  if( debug ) printf("===========%8lld, %13llu, %f, %f\n", ImplantEntry, ImplantTime, ImplantX, ImplantY); 
    
  //searching from past 10k events to all.
  for( ULong64_t k = ImplantEntry - 10000; k < totNumEntry ; k++){ 
    b_flag->GetEntry(k);
    b_vetoFlag->GetEntry(k);
    b_dyBetaTime->GetEntry(k);
    b_crossTime->GetEntry(k);
    b_crossEnergy->GetEntry(k);
    b_veto_f->GetEntry(k);
    b_veto_r->GetEntry(k);
    b_xBeta->GetEntry(k);
    b_yBeta->GetEntry(k);
    
    if ( k == ImplantEntry ) continue; // skip itself
    
    //if( crossTime > 0 ) continue;
    //if( crossEnergy > 0 ) continue;
    if( (flag & 2) == 2 ) continue; /// has only Ions
    //if( flag  != 4 ) continue; /// no only Beta
    if( (flag & 4) == 0 ) continue; /// no any Beta
    if( !TMath::IsNaN(veto_f) ) continue; /// no veto_front
    if( !TMath::IsNaN(veto_r) ) continue; /// no veto_rear
    //if( dyBetaTime[0] == 0) continue;
    if( dist(ImplantX, xBeta, ImplantY, yBeta) > distThreshold ) continue; /// when dist between ions and beta > distThread
    
    if( k < ImplantEntry && ( ImplantTime - dyBetaTime[0] ) * tick2ms > backwardTime) continue; /// when past event, more than backwardTime
    if( k > ImplantEntry && ( dyBetaTime[0] - ImplantTime ) * tick2ms > forwardTime ) break; /// when future event, more then forwardTime
    
    //===== print status
    clock.Stop("timer");
    Double_t time = clock.GetRealTime("timer");
    clock.Start("timer");
    printf(" %llu[%4.2f%%], searched next %lld events | match case %d | expect : %5.2f min\r", 
         ImplantEntry, ImplantEntry*100./totNumEntry, k - ImplantEntry, 
         numMatch, totNumEntry*time/(ImplantEntry+1.)/60.);
      
    if( debug ) printf("           %8lld, %13llu, %f, %f\n", k, dyBetaTime[0], xBeta, yBeta);
    if( debug ) printf("+++++++++++%8s, %13llu | %.f ms\n", "+", dyBetaTime[0], (dyBetaTime[0] - ImplantTime)*tick2ms);
    
    hDist->Fill(ImplantX - xBeta, ImplantY - yBeta);
      
    numMatch ++;           
    double dT = 0;
    if( dyBetaTime[0] > ImplantTime ){
      dT = (dyBetaTime[0] - ImplantTime)*tick2ms;
    }else{
      dT = (ImplantTime - dyBetaTime[0])*tick2ms;
      dT = -dT;
    }
  
    hDecay[cutID]->Fill(dT);
    //hvetoF[cutID]->Fill(crossEnergy);
    hBetaPos->Fill(xBeta, yBeta);
    
  }
  
  return kTRUE;
}


void DecayFinder::Terminate(){
  
  //============== Save the results
  TFile * haha = new TFile("results.root", "recreate");
  haha->cd();
  for( int i = 0; i < numCut; i++){
    hDecay[i]->Write();
  }
  hPID->Write();
  haha->Close();
  
  //============== Plot result
  gStyle->SetOptStat("neiou");
  TCanvas * cDecay = new TCanvas("cDecay", "Decay", 1000, 1000);
  cDecay->Divide(2,2);
  
  int padID = 1; cDecay->cd(padID); 
  hPID->Draw("colz");
  for( int i = 0; i < numCut; i++){
    cut = (TCutG*) cutList->At(i);
    cut->SetLineColor(i+2);
    cut->Draw("same");
  }
  
  padID++; cDecay->cd(padID); cDecay->cd(padID)->SetLogy();
  double yMax = 0;
  for( int i = 0; i < numCut; i++){
    if( hDecay[i]->GetMaximum() > yMax ) yMax = hDecay[i]->GetMaximum();
  }
  for( int i = 0; i < numCut; i++){
    hDecay[i]->SetLineColor(i+2);
    hDecay[i]->SetMaximum(yMax *1.2);
    hDecay[i]->Draw(i == 0 ? "" : "same");
  }
  
  //padID++; cDecay->cd(padID);
  //hDist->Draw("colz");
  //
  //padID++; cDecay->cd(padID);
  //for( int i = 0; i < numCut; i++){
  //  hvetoF[i]->SetLineColor(i+2);
  //	hvetoF[i]->Draw(i == 0 ? "" : "same");
  //}

  padID++; cDecay->cd(padID);
  hIonsPos->Draw("colz");

  padID++; cDecay->cd(padID);
  hBetaPos->Draw("colz");
  
  //================ Save historgram 
  TString fHistRootName = "zzz_all_decay.root";
  if( true ){
    TFile * fHist = new TFile(fHistRootName, "recreate");
    fHist->cd();
    
    for( int i = 0 ; i < numCut; i++){
      hDecay[i]->Write();
    }
    
    fHist->Close();
    
    printf("---- Save PID histogram as %s \n", fHistRootName.Data());
  }

}