#include <iostream>
#include <fstream>

gROOT->Reset();

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

void histMk_3D_BpBm_D1_DchKPiPi() {

  // Root macro to produce 3D histogram of BpBm MC.
  // BpBm 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);

  // FF parameters
  // SP5/SP6
  double R1_old = 1.18;
  double R2_old = 0.72;
  double rhoSq_old = 0.92;
  // New values
  double R1_new = 1.18; //1.396;
  double R2_new = 0.72; //0.885;
  double rhoSq_new = 0.92; //1.145;

  // mass
  double m_D1p = 2010.0;
  double m_D10 = 2006.7;
  double m_Bp = 5279.0;
  double m_B0 = 5279.4;

  // Info of output files
  TString out3DHistoName(outFileNameBase+"_BpBm_D1_DchKPiPi.root");
  TString outInfoFileName(outFileNameBase+"_BpBm_D1_DchKPiPi.txt");
  TString histTitle_peak("BpBm to D1 (DchKPipi) D mass peak");
  TString histTitle_side("BpBm to D1 (DchKPipi) D mass sideband");
  TString histTitle_sub("BpBm to D1 (DchKPipi) sideband subtracted");

  // B decay mode selection
  int selBDecMod1 = 1540;
  int selBDecMod2 = 2530;
  int selBDecMod3 = -1;
  int selBDecMod4 = -1;

  // D decay mode selection
  int DDecMod_3D = DDecMod_DchKPiPi;

  // D mass range
  float DMass1 = DMass_DchKPiPi1; // lower edge of sideband
  float DMass2 = DMass_DchKPiPi2; // lower edge of D mass peak
  float DMass3 = DMass_DchKPiPi3; // upper edge of D mass peak
  float DMass4 = DMass_DchKPiPi4; // upper edge of sideband

  // Input files
  TString inFileName_Run1(baseDir_Run1+"/"+inFile_Run1_BpBm);
  int nStart_Run1 = nStart_Run1_BpBm;
  int nStop_Run1 = nStop_Run1_BpBm;
  TString inFileName_Run2(baseDir_Run2+"/"+inFile_Run2_BpBm);
  int nStart_Run2 = nStart_Run2_BpBm;
  int nStop_Run2 = nStop_Run2_BpBm;
  TString inFileName_Run3(baseDir_Run3+"/"+inFile_Run3_BpBm);
  int nStart_Run3 = nStart_Run3_BpBm;
  int nStop_Run3 = nStop_Run3_BpBm;
  TString inFileName_Run4(baseDir_Run4+"/"+inFile_Run4_BpBm);
  int nStart_Run4 = nStart_Run4_BpBm;
  int nStop_Run4 = nStop_Run4_BpBm;

  // Luminosity normalization factors
  double c_Run1 = c_Run1_BpBm;
  double c_Run2 = c_Run2_BpBm;
  double c_Run3 = c_Run3_BpBm;
  double c_Run4 = c_Run4_BpBm;
  double c_All = c_All_BpBm;


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

  cout << "starat BpBmToD1\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 == 1) continue;
      if (directLepFromB == 0) continue;
      if (lepFromSameB == 0) continue;
      int goodBDecMod(0);
      if (mcBDecMod == selBDecMod1 )  goodBDecMod=1;
      if (mcBDecMod == selBDecMod2 )  goodBDecMod=1;
      if (mcBDecMod == selBDecMod3 )  goodBDecMod=1;
      if (mcBDecMod == selBDecMod4 )  goodBDecMod=1;
      if (goodBDecMod != 1) 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 << "B decay mode selection" << endl;
  outFile << "  selBDecMod1 = " << selBDecMod1 << endl;
  outFile << "  selBDecMod2 = " << selBDecMod2 << endl;
  outFile << "  selBDecMod3 = " << selBDecMod3 << endl;
  outFile << "  selBDecMod4 = " << selBDecMod4 << 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 << "FF parameters" << endl;
  outFile << "SP5/SP6" << endl;
  outFile << "  R1_old = " << R1_old << endl;
  outFile << "  R2_old = " << R2_old << endl;
  outFile << "  rhoSq_old = " << rhoSq_old << endl;
  outFile << "New values" << endl;
  outFile << "  R1_new = " << R1_new << endl;
  outFile << "  R2_new = " << R2_new << endl;
  outFile << "  rhoSq_new = " << rhoSq_new << endl;
  outFile << "Masses" << endl;
  outFile << "  m_D1p = " << m_D1p << endl;
  outFile << "  m_D10 = " << m_D10 << endl;
  outFile << "  m_Bp = " << m_Bp << endl;
  outFile << "  m_B0 = " << m_B0 << endl;
  outFile.close();

}