#include<TFile.h>
#include<TLorentzVector.h>
#include<TTree.h>
#include <chrono>
#include <iostream>

using namespace std::chrono;

Bool_t CheckDataProcessingCuts(TLorentzVector metTlv){
	Double_t metPt = metTlv.Pt();
	Double_t metPtMin = 0.5; //GeV
	Bool_t passedCuts = metPt > metPtMin;
	return passedCuts;
}

TTree* ConvertTree(TTree* inputTree) {

  Double_t phPt = 0; 
  Double_t phEta = 0; 
  Double_t phPhi = 0;

  Double_t leadJetPt = 0; 
  Double_t leadJetEta = 0; 
  Double_t leadJetPhi = 0; 
  Double_t leadJetE = 0;

  Double_t subleadJetPt = 0; 
  Double_t subleadJetEta = 0; 
  Double_t subleadJetPhi = 0; 
  Double_t subleadJetE = 0;

  Double_t metPt = 0;
  Double_t metPhi = 0;

  Int_t nJets = 0;

  inputTree->SetBranchAddress( "ph_pt", &phPt );
  inputTree->SetBranchAddress( "ph_eta", &phEta );
  inputTree->SetBranchAddress( "ph_phi", &phPhi );

  inputTree->SetBranchAddress( "jet_lead_pt", &leadJetPt );
  inputTree->SetBranchAddress( "jet_lead_eta", &leadJetEta );
  inputTree->SetBranchAddress( "jet_lead_phi", &leadJetPhi );
  inputTree->SetBranchAddress( "jet_lead_E", &leadJetE );

  inputTree->SetBranchAddress( "jet_sublead_pt", &subleadJetPt );
  inputTree->SetBranchAddress( "jet_sublead_eta", &subleadJetEta );
  inputTree->SetBranchAddress( "jet_sublead_phi", &subleadJetPhi );
  inputTree->SetBranchAddress( "jet_sublead_E", &subleadJetE );

  inputTree->SetBranchAddress( "metTST_pt", &metPt );
  inputTree->SetBranchAddress( "metTST_phi", &metPhi );
  
  inputTree->SetBranchAddress( "n_jet", &nJets );

  TTree *outputTree = new TTree( "TMVA_input", "" );

  Double_t mJJ = 0;
  outputTree->Branch( "mJJ", &mJJ, "mJJ/D" );

  Int_t nEntries = (Int_t)inputTree->GetEntries();

  Int_t passedEntries = 0;
  for ( Int_t inputTreeIdx = 0; inputTreeIdx < nEntries; inputTreeIdx++ ){

    inputTree->GetEntry( inputTreeIdx );
    
    if( nJets >= 2 ) { 
      TLorentzVector metTlv;
      metTlv.SetPtEtaPhiE( metPt, 0, metPhi, metPt );
      TLorentzVector phTlv;
      phTlv.SetPtEtaPhiM( phPt, phEta, phPhi, 0 );
      TLorentzVector leadJetTlv;
      leadJetTlv.SetPtEtaPhiE( leadJetPt, leadJetEta, leadJetPhi, leadJetE );
      TLorentzVector subleadJetTlv;
      subleadJetTlv.SetPtEtaPhiE( subleadJetPt, subleadJetEta, subleadJetPhi, subleadJetE );
      
      Bool_t dataProcessingCutsPassed = CheckDataProcessingCuts(metTlv); 
     
      if(dataProcessingCutsPassed) {
      	mJJ = (leadJetTlv + subleadJetTlv).M();
    		outputTree->Fill(); 
        passedEntries++;

      }
    }
  }

  cout << passedEntries << endl;

  return outputTree;
}

int main() {
	TFile *inputFile = new TFile("testSpeed.root");
	TTree *inputTree = (TTree*)inputFile->Get("output_tree");

	TFile *outputFile = new TFile("OutputFileCpp.root", "recreate");

  high_resolution_clock::time_point t1 = high_resolution_clock::now();
	TTree *outputTree = ConvertTree(inputTree);
  high_resolution_clock::time_point t2 = high_resolution_clock::now();
  auto duration = duration_cast<milliseconds>( t2 - t1 ).count();
  cout << duration / 1000. << " seconds to convert a tree" << endl;

	outputTree->Write();
	outputFile->Close();

    return 0;
}