#include <iostream>
#include "TTree.h"
#include "TFile.h"
#include "TROOT.h"
#include "TLorentzVector.h"
#include "TMath.h"

TTree* m_tree;

Int_t InfoNums[21] = {0,0,0,0,0,0};


Int_t           NPar;
Int_t           Type[1000];          //[NPar]
Float_t         PtGen[1000];         //[NPar]
Float_t         PhiGen[1000];        //[NPar]
Float_t         EtaGen[1000];        //[NPar]
Int_t           KMothNt[1000];       //[NPar]
Int_t           KFDauNt[1000];       //[NPar]
Int_t           KLDauNt[1000];       //[NPar]

//************************************************************************************
// returns the 4-momentum of the tau after substracting the 4-momentum of its neutrino
void VisiblePartP(Int_t TNum, TLorentzVector &VP) 
{
  Double_t x = PtGen[TNum]*cos(PhiGen[TNum]);
  Double_t y = PtGen[TNum]*sin(PhiGen[TNum]);
  Double_t z = PtGen[TNum]*sinh(EtaGen[TNum]);
  Double_t e = PtGen[TNum]*cosh(EtaGen[TNum]);

  for(Int_t j=0; j<NPar; ++j) 
  {
    if ( ( abs(Type[j])==16 ) && ( KMothNt[j] == TNum ) ) 
    {
      x = x - PtGen[j]*cos(PhiGen[j]);
      y = y - PtGen[j]*sin(PhiGen[j]);
      z = z - PtGen[j]*sinh(EtaGen[j]);
      e = e - PtGen[j]*cosh(EtaGen[j]);
    }
  }
  VP.SetXYZT(x, y, z, e);
}
//*************************************************************************************
int Initialise(char* ntupFile)
{
  TFile *m_file  = new TFile(ntupFile);
  m_tree    = (TTree*) m_file->Get("CBNT/t3333");

  //reed branches
  m_tree->SetBranchAddress("NPar" ,&NPar);
  m_tree->SetBranchAddress("Type" , Type);
  m_tree->SetBranchAddress("PtGen" , PtGen);
  m_tree->SetBranchAddress("PhiGen" , PhiGen);
  m_tree->SetBranchAddress("EtaGen" , EtaGen);
  m_tree->SetBranchAddress("KMothNt" , KMothNt);
  m_tree->SetBranchAddress("KFDauNt" ,  KFDauNt);
  m_tree->SetBranchAddress("KLDauNt" , KLDauNt);
          
  return 0;  
}

//************************************************************************************
void fillHists()
{
  TLorentzVector VP;
  for (Int_t N=0;N<m_tree->GetEntries();N++)  //loop over all events
  {
    m_tree->GetEntry(N);
    Bool_t stvar;

    Bool_t gTau;

    for (Int_t i=1;i<NPar;i++) // loop over all MC taus
    { 
      if (abs(Type[i])==15)
      {
        VisiblePartP(i,VP);
        gTau=true;
        for (Int_t i2=KFDauNt[i];i2 <= KLDauNt[i]; i2++) if (abs(Type[i2])==15) gTau=false;
        if (gTau) 
        {  
          
          if (VP.Pt()/1000 > 15 ) InfoNums[1]++; //works
          if (VP.Pt()/1000 > 10 ) InfoNums[2]++; //works
          if ( fabs( VP.Eta() ) <= 2.5) InfoNums[3]++; //works
          stvar =( VP.Pt()/1000 > 10 ) && ( fabs( VP.Eta() ) <= 2.5);
          if ( stvar ) InfoNums[4]++;
          if ( ( VP.Pt()/1000 > 15 ) && ( fabs( VP.Eta() ) <= 2.5 ) ) InfoNums[5]++;
          cout << InfoNums[1] << " " << InfoNums[2] << " "<<  InfoNums[3] <<" "<<   InfoNums[4] <<" "<<   InfoNums[5] <<endl;
        }
      }
    }
  }
}
//**********************************************************************************************
int example(char* ntupFile="example.root") 
{
  
  Initialise(ntupFile);
  fillHists();
  return 0;
}