//pfSkimEx.C

#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include "TString.h"
#include "TTree.h"
#include "TLatex.h"
#include "TStyle.h"
#include "TBranch.h"
#include "TFile.h"
#include <TTree.h>
#include <TChain.h>
#include "TCanvas.h"
#include "THStack.h"
#include "TLine.h"
#include "TMath.h"
#include "TGraph.h"
#include "TLegend.h"
#include "TF1.h"
#include "TH1.h"
#include "TH2.h"
#include "TLorentzVector.h"
using namespace std;

void pfSkimEx(string inputfile, string outfile)
{   
  TFile newfile(outfile.c_str(),"recreate");
  
  TChain* calochain = new TChain("ggHiNtuplizer/EventTree");
  TChain* jetchain = new TChain("ak4PFJetAnalyzer/t");
  TChain* hiEvtchain = new TChain("hiEvtAnalyzer/HiTree");
  TChain* skimchain = new TChain("skimanalysis/HltTree");
  TChain* trackchain = new TChain("ppTrack/trackTree");
  TChain* pfchain = new TChain("pfcandAnalyzer/pfTree");

  calochain->Add(inputfile.c_str());
  jetchain->Add(inputfile.c_str());
  hiEvtchain->Add(inputfile.c_str());
  skimchain->Add(inputfile.c_str());
  trackchain->Add(inputfile.c_str());
  pfchain->Add(inputfile.c_str());

  int caloentries = calochain->GetEntries();
  int jetentries = jetchain->GetEntries();
  int hiEvtEntries = hiEvtchain->GetEntries();
  int skimEntries = skimchain->GetEntries();
  int trackEntries = trackchain->GetEntries();
  cout << "Entries in calo chain = " << caloentries << endl;
  cout << "Entries in jet chain = " << jetentries << endl;
  cout << "Entries in hiEvt chain = " << hiEvtEntries << endl;
  cout << "Entries in skim chain = " << skimEntries << endl;
  cout << "Entries in track chain = " << trackEntries << endl;

  // Tree Variables
  int nRef=0, nZDC=0, nTrk=0, nTower=0, n=0, nPFpart=0, vtxFilter=0, clusterFilter=0;
  float vtx_z=0;
  float trkPt[10000], trkEta[10000];
  float jetPt[200]={0}, jetEta[200]={0}, jetPhi[200]={0};

  vector<float> *jetPt_vec=0, *jetEta_vec=0, *jetPhi_vec=0;
  vector<float> *caloEta=0, *caloPhi=0, *caloEn=0, *calo_hadE=0, *calo_emE=0;
  vector<float> *trkPt_vec=0, *trkEta_vec=0;
  vector<float> *pfEnergy=0, *pfEta=0, *pfPhi=0;
  vector<int> *highPurity=0;

  hiEvtchain->SetBranchAddress("vz",&vtx_z);
  jetchain->SetBranchAddress("nref", &nRef);
  jetchain->SetBranchAddress("jtpt", &jetPt);
  jetchain->SetBranchAddress("jteta", &jetEta);
  jetchain->SetBranchAddress("jtphi", &jetPhi);
  skimchain->SetBranchAddress("pprimaryVertexFilter", &vtxFilter);
  skimchain->SetBranchAddress("pclusterCompatibilityFilter", &clusterFilter);
  trackchain->SetBranchAddress("nTrk", &nTrk);
  trackchain->SetBranchAddress("trkPt", &trkPt);
  trackchain->SetBranchAddress("trkEta", &trkEta);
  pfchain->SetBranchAddress("nPFpart", &nPFpart);
  pfchain->SetBranchAddress("pfEnergy", &pfEnergy);
  pfchain->SetBranchAddress("pfEta", &pfEta);
  pfchain->SetBranchAddress("pfPhi", &pfPhi);
  pfchain->SetBranchAddress("highPurity", &highPurity);

  // New ROOT File
  TTree *jetVariables = new TTree("jetVariables","jetVariables");
  jetVariables->Branch("nRef", &nRef);
  jetVariables->Branch("jetpt", &jetPt_vec);
  jetVariables->Branch("jeteta", &jetEta_vec);
  jetVariables->Branch("jetphi", &jetPhi_vec);
  TTree *hiEvent = new TTree("hiEvent","hiEvent");
  hiEvent->Branch("Vertex_Z", &vtx_z);
  hiEvent->Branch("primaryVertexFilter", &vtxFilter);
  hiEvent->Branch("clusterCompatibilityFilter", &clusterFilter);
  hiEvent->Branch("nTrk", &nTrk);
  hiEvent->Branch("trkPt", &trkPt_vec);
  hiEvent->Branch("trkEta", &trkEta_vec);
  hiEvent->Branch("nPFpart", &nPFpart);
  hiEvent->Branch("pfEnergy", &pfEnergy);
  hiEvent->Branch("pfEta", &pfEta);
  hiEvent->Branch("pfPhi", &pfPhi);
  hiEvent->Branch("highPurity", &highPurity);

  // Event Loop
  for (int i=0; i<jetentries; i++)
    {
      jetchain->GetEntry(i);
      hiEvtchain->GetEntry(i);
      skimchain->GetEntry(i);
      calochain->GetEntry(i);
      trackchain->GetEntry(i);
      pfchain->GetEntry(i);
      
      jetPt_vec->clear();
      jetEta_vec->clear();
      jetPhi_vec->clear();
      trkPt_vec->clear();
      trkEta_vec->clear();

      // Cuts
      if (vtxFilter==0) continue;
      if (clusterFilter==0) continue;
      if (fabs(vtx_z) > 20) continue;
      if (nTrk==0) continue;
      if (nRef==0) continue;
      if (nRef==1) continue;
      if ((fabs(jetEta[0]) > 2.4) || (fabs(jetEta[1]) > 2.4)) continue;
      if (jetPt[0]<30 || jetPt[1]<20) continue;

      for (int a=0; a<nRef; a++)
        {
          jetPt_vec->push_back(jetPt[a]);
          jetEta_vec->push_back(jetEta[a]);
          jetPhi_vec->push_back(jetPhi[a]);
        }
      
      for (int j=0; j<nTrk; j++)
        {
          trkPt_vec->push_back(trkPt[j]);
          trkEta_vec->push_back(trkEta[j]);
        }
      
      jetVariables->Fill();
      hiEvent->Fill();
   
    }
  newfile.Write();
  newfile.Close();
}