#include "TChain.h"
#include "TFile.h"
#include "TROOT.h"
#include "Math/Vector4D.h"
#include <iostream>
#include <thread>

using LorentzVector = ROOT::Math::PxPyPzEVector;

void minimal() {

  // Enable ROOT multithreading
  // const UInt_t nThreads = std::thread::hardware_concurrency();
  // ROOT::EnableImplicitMT(nThreads);

  Double_t muplus_PT, muplus_PE, muplus_PX, muplus_PY, muplus_PZ,
           muminus_PT, muminus_PE, muminus_PX, muminus_PY, muminus_PZ,
           Kplus_PT, Kplus_PE, Kplus_PX, Kplus_PY, Kplus_PZ, Kplus_PIDK,
           Kminus_PT, Kminus_PE, Kminus_PX, Kminus_PY, Kminus_PZ, Kminus_PIDK;

  TString title;

  LorentzVector muon1, muon2, kaon1, kaon2;

  // Input files
  TChain* ch = new TChain( "ch_All" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2016_Down.root/myTree/DecayTree" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2016_Up.root/myTree/DecayTree" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2017_Down.root/myTree/DecayTree" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2017_Up.root/myTree/DecayTree" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2018_Down.root/myTree/DecayTree" );
  ch->Add( "/eos/lhcb/wg/QEE/CEP_JPsiPhi/v3/JPsiPhi_selection/2018_Up.root/myTree/DecayTree" );

  // Deactivate all branches
  ch->SetBranchStatus( "*", 0 );
  // Reactivate needed ones
  ch->SetBranchStatus( "chi_c1_1P_P", 1 );
  ch->SetBranchStatus( "chi_c1_1P_PT", 1 );
  ch->SetBranchStatus( "chi_c1_1P_PE", 1 );
  ch->SetBranchStatus( "chi_c1_1P_PX", 1 );
  ch->SetBranchStatus( "chi_c1_1P_PY", 1 );
  ch->SetBranchStatus( "chi_c1_1P_PZ", 1 );
  ch->SetBranchStatus( "chi_c1_1P_M", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_X", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_Y", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_Z", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_XERR", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_YERR", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_ZERR", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_CHI2", 1 );
  ch->SetBranchStatus( "chi_c1_1P_ENDVERTEX_NDOF", 1 );
  ch->SetBranchStatus( "jpsi_P", 1 );
  ch->SetBranchStatus( "jpsi_PT", 1 );
  ch->SetBranchStatus( "jpsi_PE", 1 );
  ch->SetBranchStatus( "jpsi_PX", 1 );
  ch->SetBranchStatus( "jpsi_PY", 1 );
  ch->SetBranchStatus( "jpsi_PZ", 1 );
  ch->SetBranchStatus( "jpsi_M", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_X", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_Y", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_Z", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_XERR", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_YERR", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_ZERR", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_CHI2", 1 );
  ch->SetBranchStatus( "jpsi_ENDVERTEX_NDOF", 1 );
  ch->SetBranchStatus( "jpsi_FD_ORIVX", 1 );
  ch->SetBranchStatus( "jpsi_FDCHI2_ORIVX", 1 );
  ch->SetBranchStatus( "jpsi_DIRA_ORIVX", 1 );
  ch->SetBranchStatus( "muplus_P", 1 );
  ch->SetBranchStatus( "muplus_PT", 1 );
  ch->SetBranchStatus( "muplus_PE", 1 );
  ch->SetBranchStatus( "muplus_PX", 1 );
  ch->SetBranchStatus( "muplus_PY", 1 );
  ch->SetBranchStatus( "muplus_PZ", 1 );
  ch->SetBranchStatus( "muplus_PIDe", 1 );
  ch->SetBranchStatus( "muplus_PIDmu", 1 );
  ch->SetBranchStatus( "muplus_PIDK", 1 );
  ch->SetBranchStatus( "muplus_PIDp", 1 );
  ch->SetBranchStatus( "muplus_PIDd", 1 );
  ch->SetBranchStatus( "muplus_ProbNNe", 1 );
  ch->SetBranchStatus( "muplus_ProbNNk", 1 );
  ch->SetBranchStatus( "muplus_ProbNNp", 1 );
  ch->SetBranchStatus( "muplus_ProbNNpi", 1 );
  ch->SetBranchStatus( "muplus_ProbNNmu", 1 );
  ch->SetBranchStatus( "muplus_ProbNNd", 1 );
  ch->SetBranchStatus( "muplus_ProbNNghost", 1 );
  ch->SetBranchStatus( "muplus_Hlt2LowMultMuonDecision_TIS"           , 1 );
  ch->SetBranchStatus( "muplus_Hlt2LowMultMuonDecision_TOS"           , 1 );
  ch->SetBranchStatus( "muplus_Hlt1LowMultMuonDecision_TIS"           , 1 );
  ch->SetBranchStatus( "muplus_Hlt1LowMultMuonDecision_TOS"           , 1 );
  ch->SetBranchStatus( "muplus_L0Muon,lowMultDecision_TIS"            , 1 );
  ch->SetBranchStatus( "muplus_L0Muon,lowMultDecision_TOS"            , 1 );
  ch->SetBranchStatus( "muminus_P", 1 );
  ch->SetBranchStatus( "muminus_PT", 1 );
  ch->SetBranchStatus( "muminus_PE", 1 );
  ch->SetBranchStatus( "muminus_PX", 1 );
  ch->SetBranchStatus( "muminus_PY", 1 );
  ch->SetBranchStatus( "muminus_PZ", 1 );
  ch->SetBranchStatus( "muminus_PIDe", 1 );
  ch->SetBranchStatus( "muminus_PIDmu", 1 );
  ch->SetBranchStatus( "muminus_PIDK", 1 );
  ch->SetBranchStatus( "muminus_PIDp", 1 );
  ch->SetBranchStatus( "muminus_PIDd", 1 );
  ch->SetBranchStatus( "muminus_ProbNNe", 1 );
  ch->SetBranchStatus( "muminus_ProbNNk", 1 );
  ch->SetBranchStatus( "muminus_ProbNNp", 1 );
  ch->SetBranchStatus( "muminus_ProbNNpi", 1 );
  ch->SetBranchStatus( "muminus_ProbNNmu", 1 );
  ch->SetBranchStatus( "muminus_ProbNNd", 1 );
  ch->SetBranchStatus( "muminus_ProbNNghost", 1 );
  ch->SetBranchStatus( "muminus_Hlt2LowMultMuonDecision_TIS"           , 1 );
  ch->SetBranchStatus( "muminus_Hlt2LowMultMuonDecision_TOS"           , 1 );
  ch->SetBranchStatus( "muminus_Hlt1LowMultMuonDecision_TIS"           , 1 );
  ch->SetBranchStatus( "muminus_Hlt1LowMultMuonDecision_TOS"           , 1 );
  ch->SetBranchStatus( "muminus_L0Muon,lowMultDecision_TIS"            , 1 );
  ch->SetBranchStatus( "muminus_L0Muon,lowMultDecision_TOS"            , 1 );
  ch->SetBranchStatus( "phi_P", 1 );
  ch->SetBranchStatus( "phi_PT", 1 );
  ch->SetBranchStatus( "phi_PE", 1 );
  ch->SetBranchStatus( "phi_PX", 1 );
  ch->SetBranchStatus( "phi_PY", 1 );
  ch->SetBranchStatus( "phi_PZ", 1 );
  ch->SetBranchStatus( "phi_M", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_X", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_Y", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_Z", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_XERR", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_YERR", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_ZERR", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_CHI2", 1 );
  ch->SetBranchStatus( "phi_ENDVERTEX_NDOF", 1 );
  ch->SetBranchStatus( "phi_FD_ORIVX", 1 );
  ch->SetBranchStatus( "phi_FDCHI2_ORIVX", 1 );
  ch->SetBranchStatus( "phi_DIRA_ORIVX", 1 );
  ch->SetBranchStatus( "Kplus_P", 1 );
  ch->SetBranchStatus( "Kplus_PT", 1 );
  ch->SetBranchStatus( "Kplus_PE", 1 );
  ch->SetBranchStatus( "Kplus_PX", 1 );
  ch->SetBranchStatus( "Kplus_PY", 1 );
  ch->SetBranchStatus( "Kplus_PZ", 1 );
  ch->SetBranchStatus( "Kplus_PIDK", 1 );
  ch->SetBranchStatus( "Kplus_PIDe", 1 );
  ch->SetBranchStatus( "Kplus_PIDmu", 1 );
  ch->SetBranchStatus( "Kplus_PIDp", 1 );
  ch->SetBranchStatus( "Kplus_PIDd", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNe", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNk", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNp", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNpi", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNmu", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNd", 1 );
  ch->SetBranchStatus( "Kplus_ProbNNghost", 1 );
  ch->SetBranchStatus( "Kplus_isMuon", 1 );
  ch->SetBranchStatus( "Kminus_P", 1 );
  ch->SetBranchStatus( "Kminus_PT", 1 );
  ch->SetBranchStatus( "Kminus_PE", 1 );
  ch->SetBranchStatus( "Kminus_PX", 1 );
  ch->SetBranchStatus( "Kminus_PY", 1 );
  ch->SetBranchStatus( "Kminus_PZ", 1 );
  ch->SetBranchStatus( "Kminus_PIDK", 1 );
  ch->SetBranchStatus( "Kminus_PIDe", 1 );
  ch->SetBranchStatus( "Kminus_PIDmu", 1 );
  ch->SetBranchStatus( "Kminus_PIDp", 1 );
  ch->SetBranchStatus( "Kminus_PIDd", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNe", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNk", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNp", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNpi", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNmu", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNd", 1 );
  ch->SetBranchStatus( "Kminus_ProbNNghost", 1 );
  ch->SetBranchStatus( "Kminus_isMuon", 1 );
  ch->SetBranchStatus( "nPVs", 1 );
  ch->SetBranchStatus( "nLongTracks", 1 );
  ch->SetBranchStatus( "nVeloTracks", 1 );
  ch->SetBranchStatus( "nBackTracks", 1 );
  ch->SetBranchStatus( "nMuonTracks", 1 );
  ch->SetBranchStatus( "nTracks", 1 );
  ch->SetBranchStatus( "nDownstreamTracks", 1 );
  ch->SetBranchStatus( "nUpstreamTracks", 1 );
  ch->SetBranchStatus( "nTTracks", 1 );
  ch->SetBranchStatus( "nVeloClusters", 1 );
  ch->SetBranchStatus( "nSPDHits", 1 );
  ch->SetBranchStatus( "L0Data_Spd_Mult", 1 );
  ch->SetBranchStatus( "nParts", 1 );
  ch->SetBranchStatus( "AllPhotons_PX", 1 );
  ch->SetBranchStatus( "AllPhotons_PY", 1 );
  ch->SetBranchStatus( "AllPhotons_PZ", 1 );
  ch->SetBranchStatus( "AllPhotons_PE", 1 );
  ch->SetBranchStatus( "log_hrc_fom", 1 );
  ch->SetBranchStatus( "log_hrc_fom_B", 1 );
  ch->SetBranchStatus( "log_hrc_fom_F", 1 );
  ch->SetBranchStatus( "L0Muon,lowMultDecision", 1 );
  ch->SetBranchStatus( "Hlt1LowMultMuonDecision", 1 );
  ch->SetBranchStatus( "Hlt1LowMultVeloCut_LeptonsDecision", 1 );
  ch->SetBranchStatus( "Hlt2LowMultMuonDecision", 1 );
  ch->SetBranchStatus( "runNumber", 1 );
  ch->SetBranchStatus( "Polarity", 1 );

  // Set address for variabes used in the selection
  ch->SetBranchAddress( "muplus_PT", &muplus_PT );
  ch->SetBranchAddress( "muplus_PE", &muplus_PE );
  ch->SetBranchAddress( "muplus_PX", &muplus_PX );
  ch->SetBranchAddress( "muplus_PY", &muplus_PY );
  ch->SetBranchAddress( "muplus_PZ", &muplus_PZ );
  ch->SetBranchAddress( "muminus_PT", &muminus_PT );
  ch->SetBranchAddress( "muminus_PE", &muminus_PE );
  ch->SetBranchAddress( "muminus_PX", &muminus_PX );
  ch->SetBranchAddress( "muminus_PY", &muminus_PY );
  ch->SetBranchAddress( "muminus_PZ", &muminus_PZ );
  ch->SetBranchAddress( "Kplus_PT", &Kplus_PT );
  ch->SetBranchAddress( "Kplus_PE", &Kplus_PE );
  ch->SetBranchAddress( "Kplus_PX", &Kplus_PX );
  ch->SetBranchAddress( "Kplus_PY", &Kplus_PY );
  ch->SetBranchAddress( "Kplus_PZ", &Kplus_PZ );
  ch->SetBranchAddress( "Kplus_PIDK", &Kplus_PIDK );
  ch->SetBranchAddress( "Kminus_PT", &Kminus_PT );
  ch->SetBranchAddress( "Kminus_PE", &Kminus_PE );
  ch->SetBranchAddress( "Kminus_PX", &Kminus_PX );
  ch->SetBranchAddress( "Kminus_PY", &Kminus_PY );
  ch->SetBranchAddress( "Kminus_PZ", &Kminus_PZ );
  ch->SetBranchAddress( "Kminus_PIDK", &Kminus_PIDK );

  // Output file for elected events
  TFile* ofile = new TFile( "./minimal.root", "RECREATE" );
  TTree* chnew = ch->CloneTree(0);

  // Selection cuts
  const Double_t mu_pt_min = 200.;
  const Double_t k_pt_min = 200.;
  const Double_t eta_min = 2.0;
  const Double_t eta_max = 5.0;
  const Double_t PIDK_min = 0.;

  const Int_t nentries = ch->GetEntries() / 300;
  for ( Int_t i = 0; i < nentries; i++ ) {
    ch->GetEntry(i);

    muon1.SetPxPyPzE( muminus_PX, muminus_PY, muminus_PZ, muminus_PE );
    muon2.SetPxPyPzE( muplus_PX, muplus_PY, muplus_PZ, muplus_PE );
    kaon1.SetPxPyPzE( Kminus_PX, Kminus_PY, Kminus_PZ, Kminus_PE );
    kaon2.SetPxPyPzE( Kplus_PX, Kplus_PY, Kplus_PZ, Kplus_PE );

    if ( muplus_PT > mu_pt_min && muminus_PT > mu_pt_min && Kplus_PT > k_pt_min && Kminus_PT > k_pt_min ) {
      if ( muon1.Eta() >= eta_min && muon1.Eta() <= eta_max && muon2.Eta() >= eta_min && muon2.Eta() <= eta_max ) {
        if ( kaon1.Eta() >= eta_min && kaon1.Eta() <= eta_max && kaon2.Eta() >= eta_min && kaon2.Eta() <= eta_max ) {
          if ( Kminus_PIDK > PIDK_min && Kplus_PIDK > PIDK_min ) {
            chnew->Fill();
          }
        }
      }
    }
  }

  ofile->Write();
  ofile->Close();
}