#include <iostream>
#include <fstream>

gROOT->Reset();

// Setup file including parameters
#include "histMk_3D_setup.h"

void histMk_3D_B0B0_LMisID_D0KPi() {

  // Root macro to produce 3D histogram of B0B0 MC.
  // B0B0 MC is luminosity nomalized to OnPeak data.
  //###########################################################
  // Change these values to meet your condition
  //###########################################################
 
  TString RunNum("RunAll");
  if (run == 1) RunNum = "Run1";
  if (run == 2) RunNum = "Run2";
  if (run == 3) RunNum = "Run3";
  if (run == 4) RunNum = "Run4";

  TString MCCorr("noCorr");
  if (correctTrk) MCCorr ="effCorr";

  // Name base of output files
  TString outFileNameBase("histo_3D_"+RunNum+"_"+MCCorr);

  // Info of output files
  TString out3DHistoName(outFileNameBase+"_B0B0_LMisID_D0KPi.root");
  TString outInfoFileName(outFileNameBase+"_B0B0_LMisID_D0KPi.txt");
  TString histTitle_peak("B0B0 to LMisID (D0Kpi) D mass peak");
  TString histTitle_side("B0B0 to LMisID (D0Kpi) D mass sideband");
  TString histTitle_sub("B0B0 to LMisID (D0Kpi) sideband subtracted");

  // D decay mode selection
  int DDecMod_3D = DDecMod_D0KPi;

  // D mass range
  float DMass1 = DMass_D0KPi1; // lower edge of sideband
  float DMass2 = DMass_D0KPi2; // lower edge of D mass peak
  float DMass3 = DMass_D0KPi3; // upper edge of D mass peak
  float DMass4 = DMass_D0KPi4; // upper edge of sideband

  // Input files
  TString inFileName_Run1(baseDir_Run1+"/"+inFile_Run1_B0B0);
  int nStart_Run1 = nStart_Run1_B0B0;
  int nStop_Run1 = nStop_Run1_B0B0;
  TString inFileName_Run2(baseDir_Run2+"/"+inFile_Run2_B0B0);
  int nStart_Run2 = nStart_Run2_B0B0;
  int nStop_Run2 = nStop_Run2_B0B0;
  TString inFileName_Run3(baseDir_Run3+"/"+inFile_Run3_B0B0);
  int nStart_Run3 = nStart_Run3_B0B0;
  int nStop_Run3 = nStop_Run3_B0B0;
  TString inFileName_Run4(baseDir_Run4+"/"+inFile_Run4_B0B0);
  int nStart_Run4 = nStart_Run4_B0B0;
  int nStop_Run4 = nStop_Run4_B0B0;

  // Luminosity normalization factors
  double c_Run1 = c_Run1_B0B0;
  double c_Run2 = c_Run2_B0B0;
  double c_Run3 = c_Run3_B0B0;
  double c_Run4 = c_Run4_B0B0;
  double c_All = c_All_B0B0;


  //#######################################################################
  // Create D mass peak and sideband arrays, readin data and fill the arrays.

  cout << "starat B0B0ToLMisID\n";
  cout << "start chain\n";
  if (run == 1) {
    TChain * chain_3D = new TChain(ntpName);
    for (Int_t i=nStart_Run1; i<=nStop_Run1; i++) {
      chain_3D->Add(inFileName_Run1+"-"+i+".root");
    }
    TTree * tree_3D = chain_3D;
    double conv = c_Run1;
  } else if (run == 2) {
    TChain * chain_3D = new TChain(ntpName);
    for (Int_t i=nStart_Run2; i<=nStop_Run2; i++) {
      chain_3D->Add(inFileName_Run2+"-"+i+".root");
    }
    TTree * tree_3D = chain_3D;
    double conv = c_Run2;
  } else if (run == 3) {
    TChain * chain_3D = new TChain(ntpName);
    for (Int_t i=nStart_Run3; i<=nStop_Run3; i++) {
      chain_3D->Add(inFileName_Run3+"-"+i+".root");
    }
    TTree * tree_3D = chain_3D;
    double conv = c_Run3;
  } else if (run == 4) {
    TChain * chain_3D = new TChain(ntpName);
    for (Int_t i=nStart_Run4; i<=nStop_Run4; i++) {
      chain_3D->Add(inFileName_Run4+"-"+i+".root");
    }
    TTree * tree_3D = chain_3D;
    double conv = c_Run4;
  } else {
    TChain * chain_3D = new TChain(ntpName);
    for (Int_t i=nStart_Run1; i<=nStop_Run1; i++) {
      chain_3D->Add(inFileName_Run1+"-"+i+".root");
    }
    for (Int_t i=nStart_Run2; i<=nStop_Run2; i++) {
      chain_3D->Add(inFileName_Run2+"-"+i+".root");
    }
    for (Int_t i=nStart_Run3; i<=nStop_Run3; i++) {
      chain_3D->Add(inFileName_Run3+"-"+i+".root");
    }
    for (Int_t i=nStart_Run4; i<=nStop_Run4; i++) {
      chain_3D->Add(inFileName_Run4+"-"+i+".root");
    }
    TTree * tree_3D = chain_3D;
    double conv = c_All;
  }





  // Define and initialise arrays for 3D histo
  double num_peak[nbinsx][nbinsy][nbinsz];
  double num_side[nbinsx][nbinsy][nbinsz];
  for (int i=0; i<nbinsx; i++) {
    for (int j=0; j<nbinsy; j++) {
      for (int k=0; k<nbinsz; k++) {
	num_peak[i][j][k] = 0;
	num_side[i][j][k] = 0;
      }
    }
  }

  cout << "read tree\n";
  // variables to read out
  Int_t nDl_3D(0);
  Int_t DDecayMode_3D[200];
  Float_t DMass_3D[200];
  Int_t noExtraKorPi_3D[200], mcKaonMissID_3D[200];
  Int_t mcPiMissID_3D[200], mcKPifromSameD_3D[200];
  Int_t mcLepMissID_3D[200], mcDirectLepFromB_3D[200];
  Int_t trueLepFromSameB_3D[200];
  Int_t trueBDecMod_3D[200];
  double trkEffCorr_3D[200], PIDCorr_3D[200];
  double mcW_3D[200];
  Float_t varx_3D[200], vary_3D[200], varz_3D[200];
  
  // Set reference between the above variables and root-tuple entries
  tree_3D->SetBranchAddress("nDl", &nDl_3D);
  tree_3D->SetBranchAddress("DDecayMode", &DDecayMode_3D);
  tree_3D->SetBranchAddress("DMass", &DMass_3D);
  tree_3D->SetBranchAddress("noExtraKorPi", &noExtraKorPi_3D);
  tree_3D->SetBranchAddress("mcKaonMissID", &mcKaonMissID_3D);
  tree_3D->SetBranchAddress("mcPiMissID", &mcPiMissID_3D);
  tree_3D->SetBranchAddress("mcKPifromSameD", &mcKPifromSameD_3D);
  tree_3D->SetBranchAddress("mcLepMissID", &mcLepMissID_3D);
  tree_3D->SetBranchAddress("mcDirectLepFromB", &mcDirectLepFromB_3D);
  tree_3D->SetBranchAddress("trueBDecMod", &trueBDecMod_3D);
  tree_3D->SetBranchAddress("trueLepFromSameB", &trueLepFromSameB_3D);
  tree_3D->SetBranchAddress("trkEffCorr", &trkEffCorr_3D);
  tree_3D->SetBranchAddress("PIDCorr", &PIDCorr_3D);
  tree_3D->SetBranchAddress("mcW", &mcW_3D);
  tree_3D->SetBranchAddress(varx, &varx_3D);
  tree_3D->SetBranchAddress(vary, &vary_3D);
  tree_3D->SetBranchAddress(varz, &varz_3D);

  // Read out root-tuple entries and fill arrays
  Int_t nevent_3D =(Int_t)tree_3D->GetEntries(); // number of events
  cout << "number of 3D events = " << nevent_3D << endl;
  int numCand(0), numSelCand(0);
  int numCand_DDecMod(0), numCand_Index(0);
  for (Int_t i=0;i<nevent_3D;i++) { // loop over events
    if (i==1) cout << "3D event 1" << endl;
    if (i==100) cout << "3D event 100" << endl;
    if (i==1000) cout << "3D event 1000" << endl;
    if (i==10000) cout << "3D event 10000" << endl;
    if (i==100000) cout << "3D event 100000" << endl;
    if (i==1000000) cout << "3D event 1000000" << endl;
    if (i==2000000) cout << "3D event 2000000" << endl;
    if (i==3000000) cout << "3D event 3000000" << endl;
    if (i==4000000) cout << "3D event 4000000" << endl;
    
    tree_3D->GetEntry(i); //read complete event in memory

    for (Int_t j=0;j<nDl_3D;j++) { // loop over Dl candidates
      int DDecMod = DDecayMode_3D[j];
      float DhadMass = DMass_3D[j];
      int noExtKorPi = noExtraKorPi_3D[j];
      int kaonMissID = mcKaonMissID_3D[j];
      int piMissID = mcPiMissID_3D[j];
      int KPiFromSameD = mcKPifromSameD_3D[j];
      int lepMissID = mcLepMissID_3D[j];
      int directLepFromB = mcDirectLepFromB_3D[j];
      int lepFromSameB = trueLepFromSameB_3D[j];
      int mcBDecMod = trueBDecMod_3D[j];
      double trkCorr = trkEffCorr_3D[j] * trkCorrModif;
      double PidCorr = PIDCorr_3D[j] * PidCorrModif;
      double trueW = mcW_3D[j];
      float xvar = varx_3D[j];
      float yvar = vary_3D[j];
      float zvar = varz_3D[j];

      numCand++;
      if (DDecMod != DDecMod_3D) continue;
      numCand_DDecMod++;
      if (noExtKorPi == 0) continue;
      if (kaonMissID == 1) continue;
      if (piMissID == 1) continue;
      if (KPiFromSameD == 0) continue;
      if (lepMissID == 0) continue;

      // Get indeces of this candidate
      //cout << "getting indecies" << endl;
      int xIndex(-1);
      for (Int_t k=0;k<nbinsx;k++) {
	if (xbins[k]<  xvar &&  xvar <= xbins[k+1]) {
	  xIndex = k;
	}
      }
      int yIndex(-1);
      for (Int_t k=0;k<nbinsy;k++) {
	if (ybins[k]< yvar && yvar <= ybins[k+1]) {
	  yIndex = k;
	}
      }
      int zIndex(-1);
      for (Int_t k=0;k<nbinsz;k++) {
	if (zbins[k]< zvar && zvar <= zbins[k+1]) {
	  zIndex = k;
	}
      }

      numCand_Index++;

      if (xIndex>=0 && yIndex>=0 && zIndex>=0) {
	numSelCand++;

	// Get FF correction factor
	double W_D = 1; 

	if (DMass2<DhadMass && DhadMass<DMass3) {
	  if (correctTrk) {
	    num_peak[xIndex][yIndex][zIndex] 
	      += (trkCorr * PidCorr)*W_D;
	  } else {
	    num_peak[xIndex][yIndex][zIndex] += W_D;
	  }
	}
	if ((DMass1<DhadMass && DhadMass<DMass2) 
	    ||(DMass3<DhadMass && DhadMass<DMass4)){
	  if (correctTrk) {
	    num_side[xIndex][yIndex][zIndex] 
	      += (trkCorr * PidCorr)*W_D;
	  } else {
	    num_side[xIndex][yIndex][zIndex] += W_D;
	  }
	}
      }
      
    } //end of Dl loop
  } //end of event loop
  cout << "number of processed candidates = " << numCand << endl;
  cout << "number of candidates after D decay mode selection = " 
       << numCand_DDecMod << endl;
  cout << "number of candidates before index selection = " 
       << numCand_Index << endl;
  cout << "number of selected candidates = " << numSelCand << endl;
  

  //###################################

  // Define and fill histos
  TH3D *peak 
    = new TH3D("peak", histTitle_peak,
	       nbinsx, xbins, nbinsy, ybins, nbinsz, zbins);
  TH3D *side 
    = new TH3D("side", histTitle_side, 
	       nbinsx, xbins, nbinsy, ybins, nbinsz, zbins);
  TH3D *sub 
    = new TH3D("sub", histTitle_sub, 
	       nbinsx, xbins, nbinsy, ybins, nbinsz, zbins);
 
  // Now fill the histograms
  cout << "filling histograms\n";
  for (int i=0; i<nbinsx; i++) {
    for (int j=0; j<nbinsy; j++) {
      for (int k=0; k<nbinsz; k++) {
	peak->SetBinContent(i+1,j+1,k+1,num_peak[i][j][k]);
        side->SetBinContent(i+1,j+1,k+1,num_side[i][j][k]);

	double peak3D = conv * num_peak[i][j][k];
	double side3D = conv * num_side[i][j][k];
	double sub3D = peak3D - side3D;
	double var3D = num_peak[i][j][k] + num_side[i][j][k];
	double err3D = conv * sqrt(var3D);
        sub->SetBinContent(i+1,j+1,k+1,sub3D);
        sub->SetBinError(i+1,j+1,k+1,err3D);
      }
    }
  }

  //#############################################################
  // Open output file and write histograms in it. ######
  cout <<"saving histograms to " << out3DHistoName << endl;
  TFile *f = new TFile(out3DHistoName, "Recreate");

  peak->Write();
  side->Write();
  sub->Write();

  f->Write();
  f->Close();

  //#######################################################################
  // Open output file and save parameters. #####
  cout << "dumping parameters to " << outInfoFileName << endl;
  ofstream outFile(outInfoFileName); //open output file
  if (!outFile) {
    cout << "Cannot open file " << outFile << endl;
  }
  outFile.setf(ios::fixed);
  outFile.setf(ios::showpoint);
  outFile << "Input files" << endl;
  outFile << "  run = " << run << endl;
  if (run == 1) {
    outFile << "  File name base = " << inFileName_Run1 << endl;
    outFile << "  start number = " << nStart_Run1 << endl;
    outFile << "  stop number = " << nStop_Run1 << endl;
  } else if (run == 2) {
    outFile << "  File name base = " << inFileName_Run2 << endl;
    outFile << "  start number = " << nStart_Run2 << endl;
    outFile << "  stop number = " << nStop_Run2 << endl;
  } else if (run == 3) {
    outFile << "  File name base = " << inFileName_Run3 << endl;
    outFile << "  start number = " << nStart_Run3 << endl;
    outFile << "  stop number = " << nStop_Run3 << endl;
  } else if (run == 4) {
    outFile << "  File name base = " << inFileName_Run4 << endl;
    outFile << "  start number = " << nStart_Run4 << endl;
    outFile << "  stop number = " << nStop_Run4 << endl;
  } else {
    outFile << "  File name base = " << inFileName_Run1 << endl;
    outFile << "  start number = " << nStart_Run1 << endl;
    outFile << "  stop number = " << nStop_Run1 << endl;
    outFile << "  File name base = " << inFileName_Run2 << endl;
    outFile << "  start number = " << nStart_Run2 << endl;
    outFile << "  stop number = " << nStop_Run2 << endl;
    outFile << "  File name base = " << inFileName_Run3 << endl;
    outFile << "  start number = " << nStart_Run3 << endl;
    outFile << "  stop number = " << nStop_Run3 << endl;
    outFile << "  File name base = " << inFileName_Run4 << endl;
    outFile << "  start number = " << nStart_Run4 << endl;
    outFile << "  stop number = " << nStop_Run4 << endl;
  }
  outFile << "################" << endl;
  outFile << "Output files" << endl;
  outFile << "  histograms : " << out3DHistoName << endl;
  outFile << "  information : " << outInfoFileName << endl;
  outFile << "################" << endl;
  outFile << "Event number info" << endl;
  outFile << "  number of events processed = " << nevent_3D << endl;
  outFile << "  number of candidates processed = " << numCand << endl;
  outFile << "  number of candidates after D decay mode selection = " 
	  << numCand_DDecMod << endl;
  outFile << "  number of candidates before index selection = " 
	  << numCand_Index << endl;
  outFile << "  number of candidates selected = " << numSelCand << endl;
  outFile << "################" << endl;
  outFile << "Luminosity normalization factor" << endl;
  outFile << "  conv = " << conv << endl;
  outFile << "D decay mode selection" << endl;
  outFile << "  DDecMod_3D = " << DDecMod_3D << endl;
  outFile << "D mass cut" << endl;
  outFile << "  DMass1 = " << DMass1 << endl;
  outFile << "  DMass2 = " << DMass2 << endl;
  outFile << "  DMass3 = " << DMass3 << endl;
  outFile << "  DMass4 = " << DMass4 << endl;
  outFile << "3 variables" << endl;
  outFile << "  Variable1 = " << varx << endl;
  outFile << "  Variable2 = " << vary << endl;
  outFile << "  Variable3 = " << varz << endl;
  outFile << "################" << endl;
  outFile << "Do or do not tracking and PID correction" << endl;
  outFile << "  correctTrk = " << correctTrk << endl;
  outFile << "track and PID correction modificatin factors" << endl;
  outFile << "modified correction = Corr * Modif" << endl;
  outFile << "  trkCorrModif = " << trkCorrModif << endl;
  outFile << "  PidCorrModif = " << PidCorrModif << endl;
  outFile.close();

}