#include "tdrstyle.C"
#include <vector>
#include <pair>

void retriggering() {
  float etamin=0.;
  float etamax=999.;
  float p1=20.; 
  float p2=50.; 
  float ecut=2.;
  int NOMmin=10;


  //  run(NOMmin,etamin,etamax,p1,p2,ecut,"NPHarm2_future","crab_run69522_dbx/res","run69522");
  //  run(NOMmin,etamin,etamax,p1,p2,ecut,"NPTru40_future","crab_run69522_dbx/res","run69522");
  //  run(NOMmin,etamin,etamax,p1,p2,ecut,"NPMed_future","crab_run69522_dbx/res","run69522");

  run(NOMmin,etamin,etamax,p1,p2,ecut,"NPHarm2","crab_run69522_dbxfromevaux/res","run69522");
}


run(int NOMmin = 10, float etamin=0., float etamax=999., float p1=50., float p2=80., float ecut=4.5, TString var = "NPMed", TString inputdir = "crab_muenriched_gm/res", TString outputdir = "gif") {
  // NPMed, NPTru40, NPHarm2
  setTDRStyle();

  TString treename("dedxDiscrimPower");
  treename += var;
  treename += "/MyTree";
  cout << "Opening " << treename << endl;

  TString filename(inputdir);
  filename += "/dedx_*.root";

  TChain *tree = new TChain(treename);
  tree->Add(filename);
  

  float estimator,momentum,eta;
  int NOM;
  int dbx;
  int eventNumber;
  tree->SetBranchAddress("estimator",&estimator);
  tree->SetBranchAddress("momentum",&momentum);
  tree->SetBranchAddress("eta",&eta);
  tree->SetBranchAddress("NOM",&NOM);
  tree->SetBranchAddress("dbx",&dbx);
  tree->SetBranchAddress("eventNumber",&eventNumber);

  TH2F * histo = new TH2F("histo","",50,0.,200., 100, 0. ,10.);
  TH1F * hP = new TH1F("hP","",50,0.,200.);

  TH1F * hE = new TH1F("hE","",50,0.,5.);
  TH1F * hE0 = new TH1F("hE0","",50,0.,5.);
  TH1F * hEplus3 = new TH1F("hEplus3","",50,0.,5.);
  TH1F * hEminus3 = new TH1F("hEminus3","",50,0.,5.);
  TH1F * hEplus4 = new TH1F("hEplus4","",50,0.,5.);
  TH1F * hEminus4 = new TH1F("hEminus4","",50,0.,5.);

  TH1F * hDBX = new TH1F("hDBX","",20,-10.,10.);
  TH2F * hEdbx = new TH2F("EvsDBX","",20,-10.,10., 50, 0.,5.);

  TH2F * hEE3plus = new TH2F("EE3plus","",50,0.,5., 50, 0.,5.);
  TH2F * hEE3minus = new TH2F("EE3minus","",50,0.,5., 50, 0.,5.);
  TH2F * hEE4plus = new TH2F("EE4plus","",50,0.,5., 50, 0.,5.);
  TH2F * hEE4minus = new TH2F("EE4minus","",50,0.,5., 50, 0.,5.);

  TH2F * hPP3plus = new TH2F("PP3plus","",50,0.,200., 50, 0.,200.);
  TH2F * hPP3minus = new TH2F("PP3minus","",50,0.,200., 50, 0.,200.);
  TH2F * hPP4plus = new TH2F("PP4plus","",50,0.,200., 50, 0.,200.);
  TH2F * hPP4minus = new TH2F("PP4minus","",50,0.,200., 50, 0.,200.);

  int nA=0;
  int nB=0;
  int nC=0;
  int nD=0;

  TTree* tree3plus = new TTree("tree3plus","tree3plus");
  tree3plus->Branch("eventNumber",&eventNumber,"eventNumber/I");
  tree3plus->Branch("estimator",&estimator,"estimator/F");
  tree3plus->Branch("momentum",&momentum,"momentum/F");
  TTree* tree3minus = new TTree("tree3minus","tree3minus");
  tree3minus->Branch("eventNumber",&eventNumber,"eventNumber/I");
  tree3minus->Branch("estimator",&estimator,"estimator/F");
  tree3minus->Branch("momentum",&momentum,"momentum/F");

  TTree* tree4plus = new TTree("tree4plus","tree4plus");
  tree4plus->Branch("eventNumber",&eventNumber,"eventNumber/I");
  tree4plus->Branch("estimator",&estimator,"estimator/F");
  tree4plus->Branch("momentum",&momentum,"momentum/F");
  TTree* tree4minus = new TTree("tree4minus","tree4minus");
  tree4minus->Branch("eventNumber",&eventNumber,"eventNumber/I");
  tree4minus->Branch("estimator",&estimator,"estimator/F");
  tree4minus->Branch("momentum",&momentum,"momentum/F");


  for (int i = 0; i < tree->GetEntries(); i++) {
    tree->GetEntry(i);
    if (fabs(eta) > etamin && fabs(eta) < etamax && NOM > NOMmin) {
      histo->Fill(momentum,estimator);
      hP->Fill(momentum);
      hE->Fill(estimator);
      if (dbx==0) hE0->Fill(estimator);
      if (dbx==3) {
	hEplus3->Fill(estimator);
	tree3plus->Fill();
      }
      if (dbx==-3) {
	hEminus3->Fill(estimator);
	tree3minus->Fill();
      }
      if (dbx>3) {
	hEplus4->Fill(estimator);
	tree4plus->Fill();
      }
      if (dbx<-3) {
	hEminus4->Fill(estimator);
	tree4minus->Fill();
      }

      if (momentum > p2 && estimator > ecut) nA++;
      if (momentum > p2 && estimator < ecut) nB++;
      if (momentum > p1 && momentum < p2 && estimator > ecut) nC++;
      if (momentum > p1 && momentum < p2 && estimator < ecut) nD++;

      if (dbx != 0) hEdbx->Fill(dbx,estimator);
      hDBX->Fill(dbx);

      // safety check
      if (estimator > 12) cout << "!!!!! dE/dx = " << estimator << ", NOM = " << NOM << endl;
    }
  }

  cout << "---------------------------" << endl;
  cout << "entries = " << tree->GetEntries() << endl;
  cout << "---------------------------" << endl;

  // matching:
  vector<int> alreadyMatched;
  alreadyMatched.clear();
  cout << "tree3plus contains " << tree3plus->GetEntries() << " entries" << endl;
  for (int i = 0; i < tree3plus->GetEntries(); i++) {
    tree3plus->GetEntry(i);
    float est1 = estimator;
    float est2 = matchedValue(eventNumber,tree3minus,alreadyMatched).first;
    if (est2 > 0.) alreadyMatched.push_back(eventNumber);
    float p1 = momentum;
    float p2 = matchedValue(eventNumber,tree3minus,alreadyMatched).second;
    hEE3plus->Fill(est1,est2);
    hPP3plus->Fill(p1,p2);
  }
  cout << "tree3minus contains " << tree3minus->GetEntries() << " entries" << endl;
  for (int i = 0; i < tree3minus->GetEntries(); i++) {
    tree3minus->GetEntry(i);
    float est2 = estimator;
    float est1 = matchedValue(eventNumber,tree3plus,alreadyMatched).first;
    if (est1 > 0.) alreadyMatched.push_back(eventNumber);
    hEE3minus->Fill(est1,est2);
  }
  cout << "tree4plus contains " << tree4plus->GetEntries() << " entries" << endl;
  for (int i = 0; i < tree4plus->GetEntries(); i++) {
    tree4plus->GetEntry(i);
    float est1 = estimator;
    float est2 = matchedValue(eventNumber,tree4minus,alreadyMatched).first;
    if (est2 > 0.) alreadyMatched.push_back(eventNumber);
    hEE4plus->Fill(est1,est2);
  }
  cout << "tree4minus contains " << tree4minus->GetEntries() << " entries" << endl;
  for (int i = 0; i < tree4minus->GetEntries(); i++) {
    tree4minus->GetEntry(i);
    float est2 = estimator;
    float est1 = matchedValue(eventNumber,tree4plus,alreadyMatched).first;
    if (est1 > 0.) alreadyMatched.push_back(eventNumber);
    hEE4minus->Fill(est1,est2);
  }

  // plots:
  TCanvas* myCanvas = new TCanvas("myCanvas","2d",1);
  myCanvas->Divide(2,2);
  gStyle->SetPalette(1);

  char* xlabel = "DBX";
  char* ylabel = "dE/dx [MeV/cm]";;

  myCanvas->cd(1);
  hEdbx->GetXaxis()->SetTitle(xlabel);
  hEdbx->GetYaxis()->SetTitle(ylabel);
  hEdbx->DrawCopy("colz");

  myCanvas->cd(2);
  gPad->SetLogy();
  hDBX->GetXaxis()->SetTitle(xlabel);
  hDBX->GetYaxis()->SetTitle("tracks");
  hDBX->DrawCopy();

  myCanvas->cd(3);
  TH2F * hEE3 = (TH2F*)hEE3plus->Clone("sum3");
  hEE3->Add(hEE3minus);
  hEE3->GetXaxis()->SetTitle("dE/dx, first event");
  hEE3->GetYaxis()->SetTitle("dE/dx, second event");
  hEE3->DrawCopy("colz");

  myCanvas->cd(4);
  TH2F * hEE4 = (TH2F*)hEE4plus->Clone("sum4");
  hEE4->Add(hEE4minus);
  hEE4->GetXaxis()->SetTitle("dE/dx, first event");
  hEE4->GetYaxis()->SetTitle("dE/dx, second event");
  hEE4->DrawCopy("colz");


  TString gif(outputdir);
  gif += "/dbx_";
  gif += var;
  gif += Form("_NOM_%2d.gif",NOMmin);
  myCanvas->Print(gif);

  TCanvas* myCanvas2 = new TCanvas("myCanvas2","2d",1);
  myCanvas2->Divide(2,2);
  gStyle->SetPalette(1);

  char* xlabel = "dE/dx, first event";
  char* ylabel = "dE/dx, second event";

  myCanvas2->cd(1);
  hEE3plus->GetXaxis()->SetTitle(xlabel);
  hEE3plus->GetYaxis()->SetTitle(ylabel);
  hEE3plus->DrawCopy("colz");

  myCanvas2->cd(2);
  hEE3minus->GetXaxis()->SetTitle(xlabel);
  hEE3minus->GetYaxis()->SetTitle(ylabel);
  hEE3minus->DrawCopy("colz");

  myCanvas2->cd(3);
  hEE4plus->GetXaxis()->SetTitle(xlabel);
  hEE4plus->GetYaxis()->SetTitle(ylabel);
  hEE4plus->DrawCopy("colz");

  myCanvas2->cd(4);
  hEE4minus->GetXaxis()->SetTitle(xlabel);
  hEE4minus->GetYaxis()->SetTitle(ylabel);
  hEE4minus->DrawCopy("colz");


  TString gif2(outputdir);
  gif2 += "/dbx2_";
  gif2 += var;
  gif2 += Form("_NOM_%2d.gif",NOMmin);
  myCanvas2->Print(gif2);

  TCanvas* myCanvas3 = new TCanvas("myCanvas3","2d",1);
  myCanvas3->cd(1);
  //  hEE3plus->GetXaxis()->SetTitle(xlabel);
  //  hEE3plus->GetYaxis()->SetTitle(ylabel);
  hPP3plus->DrawCopy("colz");


  TCanvas* myCanvasFit = new TCanvas("myCanvasFit","fit",2);
  myCanvasFit->Divide(2,3);

  myCanvasFit->cd(1);
  hE->DrawCopy();

  myCanvasFit->cd(2);
  hE0->DrawCopy();
  
  myCanvasFit->cd(3);
  hEplus3->DrawCopy();
  
  myCanvasFit->cd(4);
  hEminus3->DrawCopy();
  doubleFit(hEminus3);
  
  myCanvasFit->cd(5);
  hEplus4->DrawCopy();
  
  myCanvasFit->cd(6);
  hEminus4->DrawCopy();
  
  TString fitgif(outputdir);
  fitgif += "/fit_";
  fitgif += var;
  fitgif += ".gif";
  myCanvasFit->Print(fitgif);


  delete histo;
  delete hP;
  delete hE;
  delete hEdbx;
}


void normalize(TH1* histo, double area=1.)
{
  double scale = area/(histo->Integral());
  histo->Scale(scale);
  //  histo->Sumw2();
}

/* round number n to d decimal points */
double fround(double n, unsigned d)
{
  return floor(n * pow(10., d) + .5) / pow(10., d);
}

pair<float,float> matchedValue(int ev, TTree* inputtree, vector<int>& alreadyMatched) {
  float est=0.;
  float p=0.;

  for (int i = 0; i < alreadyMatched.size(); i++) {
    if (ev<=alreadyMatched.at(i)+1 && ev>=alreadyMatched.at(i)-1) {pair<float,float> pp = make_pair(est,p); return pp;};
  }

  TTree* tree = inputtree->CloneTree();
  float estimator;
  float momentum;
  int eventNumber;
  tree->SetBranchAddress("estimator",&estimator);
  tree->SetBranchAddress("momentum",&momentum);
  tree->SetBranchAddress("eventNumber",&eventNumber);

  for (int i = 0; i < tree->GetEntries(); i++) {
    tree->GetEntry(i);
    if (eventNumber == ev+1 || eventNumber == ev-1) {
      est=estimator; 
      p=momentum;
      //      cout << "found matching: event " << ev << " is matched to " << eventNumber << endl;
    }
  }

  pair<float,float> pp =  make_pair(est,p);
  return pp;
}

void doubleFit(TH1F* hE) {

  Double_t par[6];
  TF1 *g1    = new TF1("g1","gaus",0.,2.5);
  TF1 *g2    = new TF1("g2","gaus",2.5,4.);
  TF1 *total = new TF1("total","gaus(0)+gaus(3)",0.,4.);
  total->SetLineColor(2);
  hE->Fit(g1,"R");
  hE->Fit(g2,"R+");
  g1->GetParameters(&par[0]);
  g2->GetParameters(&par[3]);
  total->SetParameters(par);
  hE->Fit(total,"R+");
  hE->DrawCopy();

}