#include <cmath>
#include <iostream>
#include <string>
#include <sstream>
#include <fstream>
#include <cstdlib>
#include <vector>
#include <map>
#include <unistd.h>

#include <TCanvas.h>
#include <TChain.h>
#include <TFile.h>
#include <TGaxis.h>
#include <TH1D.h>
#include <TROOT.h>
#include <TStyle.h>

#include "../Paths.hpp"
#include "../GlobalFunctions.hh"


//Adds q2 and b mass branches into the evtGen tuple
bool addQ2ToEvtGenMC(bool ReferenceChannel, bool PHSP){

  gStyle -> SetOptStat(0);
  LHCbStyle();
  gROOT->SetBatch(kTRUE);
  TH1::SetDefaultSumw2(kTRUE);

  TChain * ch = new TChain("MCDecayTreeTuple/MCDecayTree");
  string file = GetGenLevelFile(ReferenceChannel,PHSP);
  replace(file,".root","_orig.root");
  ch->Add(file.c_str());

  //Reading the initial tree
  unsigned int N = ch->GetEntries();
  if(N == 0){
    coutERROR("No entries found in TFiles");
    return 1;
  }
  else coutInfo("Processing " +to_string(N) + " events!");

  //Create new tree
  TTree * newTree = ch->CloneTree(0);
  newTree->SetName("DecayTree");
  double J_psi_M = 0;
  newTree->Branch("J_psi_M", &J_psi_M, "J_psi_M/D");
  double J_psi_PX = 0;
  newTree->Branch("J_psi_PX",&J_psi_PX,"J_psi_PX/D");
  double J_psi_PY = 0;
  newTree->Branch("J_psi_PY",&J_psi_PY,"J_psi_PY/D");
  double J_psi_PZ = 0;
  newTree->Branch("J_psi_PZ",&J_psi_PZ,"J_psi_PZ/D");
  double J_psi_PE = 0;
  newTree->Branch("J_psi_PE",&J_psi_PE,"J_psi_PE/D");
  double Q2 = 0;
  newTree->Branch("Q2",      &Q2,      "Q2/D");
  double B_plus_M = 0;
  newTree->Branch("B_plus_M", &B_plus_M, "B_plus_M/D");

  TFile *output = new TFile(GetGenLevelFile(ReferenceChannel,PHSP).c_str(),"RECREATE");

  double mm_pe, mm_px, mm_py, mm_pz;
  double mp_pe, mp_px, mp_py, mp_pz;
  double b_pe, b_px, b_py, b_pz;

  ch->SetBranchStatus("*",1);
  ch->SetBranchAddress("muplus_TRUEP_E", &mp_pe);
  ch->SetBranchAddress("muplus_TRUEP_X", &mp_px);
  ch->SetBranchAddress("muplus_TRUEP_Y", &mp_py);
  ch->SetBranchAddress("muplus_TRUEP_Z", &mp_pz);

  ch->SetBranchAddress("muminus_TRUEP_E", &mm_pe);
  ch->SetBranchAddress("muminus_TRUEP_X", &mm_px);
  ch->SetBranchAddress("muminus_TRUEP_Y", &mm_py);
  ch->SetBranchAddress("muminus_TRUEP_Z", &mm_pz);

  ch->SetBranchAddress("Bplus_TRUEP_E", &b_pe);
  ch->SetBranchAddress("Bplus_TRUEP_X", &b_px);
  ch->SetBranchAddress("Bplus_TRUEP_Y", &b_py);
  ch->SetBranchAddress("Bplus_TRUEP_Z", &b_pz);

   TLorentzVector tmp_plus;
   TLorentzVector tmp_minus;
   TLorentzVector tmp_J_psi;
   TLorentzVector tmp_Bplus;
   vector<double> double_E;
   for(unsigned int e = 0; e < N; e++){
       ch->GetEntry(e);
       tmp_plus.SetPxPyPzE(mp_px,mp_py,mp_pz,mp_pe);
       tmp_minus.SetPxPyPzE(mm_px,mm_py,mm_pz,mm_pe);
       tmp_Bplus.SetPxPyPzE(b_px,b_py,b_pz,b_pe);
       double_E.push_back(mm_pe);
       double_E.push_back(mp_pe);
       tmp_J_psi = tmp_plus + tmp_minus;
       J_psi_M  = tmp_J_psi.M();
       J_psi_PX = tmp_J_psi.Px();
       J_psi_PY = tmp_J_psi.Py();
       J_psi_PZ = tmp_J_psi.Pz();
       J_psi_PE = tmp_J_psi.E();
       Q2       = tmp_J_psi.M2();
       B_plus_M = tmp_Bplus.M();
       newTree->Fill();
   }
   //Event number multiple candidates
   int counter = 0;
   for_indexed(auto firstNumber: double_E){
       for (vector<double>::iterator evt = double_E.begin()+i+1; evt < double_E.end(); evt++){
           if (double_E.at(i)==*evt)counter++;
       }
   }
   coutInfo("Total number of double candidates: " +to_string(counter));

   output->cd();
   newTree->Write("",TObject::kWriteDelete);
   output->Close();

   return 1;

}

void getQ2weights(){

  bool onlyCandOne = false;
  bool onlynTotOne = false;

  bool WriteToROOTfile = true;

  gStyle -> SetOptStat(0);
  LHCbStyle();
  gROOT->SetBatch(kTRUE);
  TH1::SetDefaultSumw2(kTRUE);

  double hmin = 0.1, hmax = 19.0;

  TChain * ch = new TChain("MCDecayTreeTuple/MCDecayTree");

  ch->Add("MCTruth300k_*.sim.root");

  unsigned int N = ch->GetEntries();

  if(N == 0){
    std::cout << "[ERROR]\tNo entries found in TFiles" << std::endl;
    return;
  }
  else
    std::cout << "[INFO]\t\tProcess " << N << " events!" << std::endl;

  double mm_pe, mm_px, mm_py, mm_pz;
  double mp_pe, mp_px, mp_py, mp_pz;

  ch->SetBranchAddress("muplus_TRUEP_E", &mp_pe);
  ch->SetBranchAddress("muplus_TRUEP_X", &mp_px);
  ch->SetBranchAddress("muplus_TRUEP_Y", &mp_py);
  ch->SetBranchAddress("muplus_TRUEP_Z", &mp_pz);

  ch->SetBranchAddress("muminus_TRUEP_E", &mm_pe);
  ch->SetBranchAddress("muminus_TRUEP_X", &mm_px);
  ch->SetBranchAddress("muminus_TRUEP_Y", &mm_py);
  ch->SetBranchAddress("muminus_TRUEP_Z", &mm_pz);

  unsigned int ncand, ntot;

  ch->SetBranchAddress("nCandidate", &ncand);
  ch->SetBranchAddress("totCandidates", &ntot);

  TH1D * h = new TH1D("q2_eff", "q2_eff; q^{2} [GeV^{2}]; events", 100, hmin, hmax);

  for(unsigned int e = 0; e < N; e++){

    if (e % (N/100) == 0)
      std::cout << std::setprecision(2) << std::fixed << "\r[" << round(double(e*100)/N) << "%]" << std::flush;

    ch->GetEntry(e);

    if(onlynTotOne){
      if(ntot != 1){
        continue;
      }
    }
    else if(onlyCandOne){
      if(ncand != 1){
        continue;
      }
    }

    h->Fill(((mm_pe+mp_pe)*(mm_pe+mp_pe)
           - (mm_px+mp_px)*(mm_px+mp_px)
           - (mm_py+mp_py)*(mm_py+mp_py)
           - (mm_pz+mp_pz)*(mm_pz+mp_pz))
           * 1.0e-6);
  }
  std::cout << std::endl;

  h->Scale(100.0 / h->Integral());  

  TH1D * w = new TH1D("flatq2weights", "flatq2weights; q^{2} [GeV^{2}]; weights", 100, hmin, hmax);
  for(int i = 0; i < 100; i++)
    w->SetBinContent(i+1, 1.0);
 
  bool div_success = w->Divide(h);
  if(!div_success)
    std::cout << "Devition was not successfull!" << std::endl;

  TCanvas * cQ2 = new TCanvas("cQ2", "cQ2");
  cQ2->cd();
  h->GetYaxis()->SetRangeUser(0.5, 1.2);
  h->Draw();
  cQ2->SaveAs("q2efficiency.eps", "eps");
  w->GetYaxis()->SetRangeUser(0.0, 2.0);
  w->Draw();
  TLine * l = new TLine(hmin, 1.0, hmax, 1.0);
  l->SetLineColor(kGray+1);
  l->SetLineStyle(kDashed);
  l->Draw("SAME");
  cQ2->SaveAs("q2weights.eps", "eps");

  if(WriteToROOTfile){
    TFile * output = new TFile("/auto/data/dgerick/B2Kstmumu/FCNCfitter/config/phspweightq2.root", "UPDATE");
    output->cd();
    w->Write("",TObject::kWriteDelete);  
    output->Close();
  }

  return;

}