#ifndef UESHAPESNOVTX_H
#define UESHAPESNOVTX_H

// STL includes
#include <vector>
#include <string>
#include <iostream>
#include <fstream>

// Root includes
#include "TH1.h"
#include "TH2.h"
#include "TH3.h"
#include "TTree.h"
#include "TStopwatch.h"
#include <TSelector.h>
#include <TLorentzVector.h>
#include <TRandom3.h>
#include "d3pdReader/McEventD3PDObject.h"
#include "d3pdReader/McTruthD3PDObject.h"
#include "d3pdReader/McGenD3PDObject.h"
#include "d3pdReader/TrackD3PDObject.h"
#include "d3pdReader/VertexD3PDObject.h"
#include "d3pdReader/VectorBosonD3PDObject.h"
#include "d3pdReader/BeamSpotD3PDObject.h"

#include "McGenD3PDObjectAdapter.h"
#include "McTruthD3PDObjectAdapter.h"
#include "Logging.h"
#include "EventShapes.h"
// #include "ITruth.h"
#include "NNtree.h"

using Eventshapes::EventShapes;
using namespace std;

#define MAXSYSTYPES 8

class UEShapesNoVtx : public TSelector {
    
 public:
  UEShapesNoVtx(TTree * tree);
  virtual void    Begin(TTree *tree);
  virtual Bool_t  ProcessCut(Long64_t entry); // strangely will not run Process()
  virtual void    Terminate();
  virtual Bool_t  Notify();


   // flags and run control
  bool flag_DEBUG;
  bool flag_QUIET;
  bool flag_isMC;
  bool flag_dumpNN;
  string MCGENERATOR;
  bool flag_doTrackControlPlots;
  int  evtmax;
  int  m_seed;
  int  m_lepton_sys;
  int  m_nprint;
  bool m_QCD;
  bool extractPU_Tracks;
  int m_maxAdditionalPileUp;
  bool m_genPUlib;
  bool m_muon;
  bool m_writeTrackLib;
  bool m_writeSideBands;
  double m_hbomprob;

  double m_CUT_ISO_DR;
  double m_CUT_ISO_PTFRAC;
  double m_CUT_ISO_PT;

  double m_MLL_MIN;
  double m_MLL_MAX;

  bool m_dumpReco;
  TString m_dumpRecoFile;
  ofstream m_dumpFile;
  ofstream m_printFile;

  double m_sbcut;

  // path to corrections directory
  TString m_corrpath;
  // path to pile-up library
  TString m_pileuplibpath;
 
  TString m_HBOMFile;
  TString m_HBOMMeanFile;
  // Mean HBOM correction profiles
  TProfile* m_meanCorrNtrk      ;
  TProfile* m_meanCorrSumPt     ;
  TProfile* m_meanCorrAvgPt     ;
  TProfile* m_meanCorrBeamThrust;
  TProfile* m_meanCorrThrust    ;
  TProfile* m_meanCorrMinor     ;
  TProfile* m_meanCorrFParameter;
  TProfile* m_meanCorrSpherocity;
  void bookHBOMHistos();
  void LoadMeanHBOMCorrections();

  vector < vector<double> > dump_vec;
  vector < long > dump_eventnumber;
  vector<double>  dump_vec_event;
  void openDumpFile();
  void closeDumpFile();
  void dumpStuff();
  bool m_applyHBOM;

  TString m_tracklibpath;
  // names of systematic variations to apply
  vector<std::string> *m_sysNames;
  
  TTree* pileUpLibrary;
  vector < vector < vector<double> > > pumap_vec;
  //typedef vector< vector<double> > MyNtrkPULib;
  //typedef map<int, MyNtrkPULib> MyPUMap;

  // This reads in the event number and run number tuples from text file
  void readEventNumbers(const char* fname);
  bool matchEventNumber(int eventnr);
  vector<int> mcEventNumbers;

  void readEventAndRunNumbers(const char* fname);
  bool matchEventAndRunNumber(int eventnr, int runnr);
  vector<vector<int> > mcEventAndRunNumbers;
  //MyPUMap pumap;
  map< int, vector<vector<vector<double> > > > pumap;
  map<int, vector<int> > nch2particles;  // For debugging
  map<int, int > statuscodes;  // For debugging

  void     addImagiroBranches(TTree* imtree);
  void     addImagiroBranches2(TTree* imtree2);
  void     initTransientVariables();

  void     LoadTrkCorrection(TString fname);
  void     LoadSideBands();
  void     LoadPUlib();

  void     WriteTrackLib(int n, vector<vector<vector<double> > > lib);
  void     LoadTrackLib();

  bool     truthCutVB(double VB_m, TVector3 l1, TVector3 l2);
  bool     recoCutVB(double VB_m, TVector3 l1, TVector3 l2);
  double   Dphi(TVector3 p1, TVector3 p2);

  double   getTrkEff();
  double   getTrkEff(double eta, double pt, string syst); 
  double   getTrkEffError(double eta, double pt);
  
  int      getEtaIndex(double eta, bool highpt);
  int      getPtIndex(double pt);
  
  bool     randomRemove(double eff);
  bool     SelectZ(TVector3 lep1,TVector3 lep2);
  bool     SelectZ(double Mee);
  TVector3 mkLepton(double pt, double eta, double phi, bool ismuon);
  void     smearTrack(TVector3 &vec, int updown);
  void     smearTrack(vector<double> &trk);
  TVector3 getDressedLepton(ITruth& mc, int index);
  bool     doTrackAnalysis(int suffix, int itrk, vector<TVector3>& trkvec, double z_vtx=0., double z0=0.0);
  bool     doTrackAnalysisLegacy(int suffix, int itrk, vector<TVector3>& trkvec);
  double   getEventWeight(bool gen=false);
  double   getTrackWeight(TVector3 trkvec);

  // Muon Isolation cut
  bool isIsolatedMuon(TVector3 lep1, TVector3 lep2, vector<TVector3> tracks);

  double   getRdmDz();
  // Histogram filling
  void fillTrackHistos(int suffix, TVector3 t, double w_trk); 
  void fillTrackD3PDObjectControlPlots(int suffix, int index, double w_trk);

  void fillMcMitVBHistos(double genMit, double recMit);
  void fillMcPtVBHistos (double genPt,  double recPt);

  void fillImagiroTree(TTree* imtree, TTree* imtree2, EventShapes* ES, double pt_VB, double weight,
		       pair<double,double> pt_VB_updown=make_pair(0,0));
  void fillImagiroTree(TTree* imtree, TTree* imtree2, EventShapes* ES, double eta_VB, double pt_VB, double weight,
		       pair<double,double> pt_VB_updown=make_pair(0,0));
  void fillEventShapes(int suffix, EventShapes* ES, double pt_VB, double weight, double z0=0.0, bool doSumPtTrans=false);
  void fillEventShapes3D(int suffix, EventShapes* ES, double pt_VB, double Mll, double weight, double z0=0.0, bool doSumPtTrans=false);
  void fillEventShapeSmear(int suffix, EventShapes* ES1, EventShapes* ES2,
			   double pt_VB_1, double pt_VB_2, double weight);
  void fillEventShapeSmear(int suffix, EventShapes* ES1,  double pt_VB_1, double weight);
  void fillEventShapeDiff(int suffix, EventShapes* ES1, EventShapes* ES2,
			  double pt_VB_1, double pt_VB_2, double weight);

  void fillEventShapesDump(int suffix, EventShapes* ES1, double pt_VB, vector < vector<double> > stl_vec);
  //void fillEventShapesDump(int suffix, EventShapes* ES1, EventShapes* ES2, double pt_VB);
  void fillMigrationMatrices(int suffix, EventShapes* recES, EventShapes* genES, 
			     double pt_VB_rec, double pt_VB_gen);
  void fillCTCorrection(EventShapes* recES, double weight);
  void fillOACorrection(EventShapes* recES, EventShapes* genES, double pt_VB, double weight);

  // Make the output file created in main.cxx known to UEShapesNoVtx
  void setOutPutFile(TFile* outputFile) {m_outputFile = outputFile;};

  // printing
  void printEvent(double mee, double pTZ, double PhiZ, vector<int> trkindices);
  void printEventProcessInfo();
  void rememberEvent(const char*, bool);
  vector<TString> m_events_for_inspection;

  static const int    ERRORVAL = 99999;

  void fillNNtreeReco(EventShapes* ES, double pt_VB, NNtree* nntree, double weight);
  void fillNNtreeRecoAndGen(EventShapes* ES, double pt_VB, EventShapes* ESgen, double pt_VBgen, NNtree* nntree, double weight);

  void LoadHBOMCorrections();
  double GetHBOMCorr(TString esname, double esvalue);
  double GetHBOMCorrErr(TString esname, double esvalue);

  // Shit for smearing kacke
  int getRegion(double eta,double phi);
  void SetRegions();
  void FillValues();
  double ErrorMatrix(int detRegion, double pt, double g3, double g4);
  TRandom3* gRand;
  ifstream InValues;
      /* corrections on resolution parameters */
    std::vector<double> m_p1_ID;
    std::vector<double> m_p2_ID;
    std::vector<double> m_p2_ID_TAN;
    std::vector<double> m_p1_MS;
    std::vector<double> m_p2_MS;
    /* stat. errors on resolution parameters corrections */
    std::vector<double> m_E_p1_ID;
    std::vector<double> m_E_p2_ID;   
    std::vector<double> m_E_p2_ID_TAN;   
    std::vector<double> m_E_p1_MS;
    std::vector<double> m_E_p2_MS;
    /* syst. errors on resolution parameters corrections */
    std::vector<double> m_S_p1_ID;
    std::vector<double> m_S_p2_ID;
    std::vector<double> m_S_p2_ID_TAN;   
    std::vector<double> m_S_p1_MS;   
    std::vector<double> m_S_p2_MS;   

  unsigned short int m_nb_regions;
  /* eta boundaries of the regions */
  std::vector<float> m_eta_min;
  std::vector<float> m_eta_max;
  /* phi boundaries of the regions */
  std::vector<float> m_phi_min;
  std::vector<float> m_phi_max;
 protected:
  //virtual void Rebin(TH1* harray[][400000]);

  void getSysType(std::string);
  //bool isLeptonTrack(TVector3 trkvec, TVector3 lepvec);
  
  void fillPileUpLibrary(double z_vtx, int first_index, int count, std::vector< std::vector<double> >* trkSelVec);

  TTree * m_tree; // 
  int m_entries; // # of entries in current tree

  // Imagiro trees
  TTree* imagiro_tree_rec;
  TTree* imagiro_tree_gen;
  TTree* imagiro_tree_rec2; 
  TTree* imagiro_tree_gen2;
  TTree* imagiro_tree_recTrkEff;
  TTree* imagiro_tree_recTrkSmearUp;
  TTree* imagiro_tree_recTrkSmearDown;
  TTree* imagiro_tree_rec2TrkEff;
  TTree* imagiro_tree_rec2TrkSmearUp;
  TTree* imagiro_tree_rec2TrkSmearDown;
  TTree* imagiro_tree_recPtZUp;
  TTree* imagiro_tree_rec2PtZUp;
  TTree* imagiro_tree_recPtZDown;
  TTree* imagiro_tree_rec2PtZDown;
  NNtree* nntree; // 

  double nevt;
  double nevt_weighted;
  double fraclast;
  double cum_merging_sum;
  double merging_prob;

  TStopwatch m_timer;
  double m_totaltime;

  // Setup logging
  Logger l;

  // output ntuple (goes to histofile)
  TFile *m_outputFile;    

  Long64_t master; // local event counter for d3pdReader objects

  // pointer(s) to the D3PDReader objects:
  D3PDReader::McEventD3PDObject     * mcevt;
  D3PDReader::McTruthD3PDObject     * mctruth;
  D3PDReader::McGenD3PDObject       * mcgen;
  D3PDReader::TrackD3PDObject       * trk;
  D3PDReader::VertexD3PDObject      * vtx;
  D3PDReader::VectorBosonD3PDObject * VB;
  D3PDReader::BeamSpotD3PDObject    * BS;

  McTruthD3PDObjectAdapter *mc; // leyton
  McGenD3PDObjectAdapter *mcGen;

  // HISTOGRAMMING
 
  TH1* harray[MAXSYSTYPES][400000];

  inline void fillh1f(int hid, double val, double weight) {
    //cout << m_systype << " " << hid << endl;
    if (harray[m_systype][hid]==0) { 
      // printf("Warning: fillh1f:: Histogram if null for [systematics][hid] = [%d][%d]\n", m_systype, hid);
      return; 
    }  
    //if (fabs(val)==ERRORVAL) cout << "ERROR filling 1d histo " << hid << " with value " << ERRORVAL << endl;
    // if (weight==1.) { harray[m_systype][hid]->Fill(val); }         
    harray[m_systype][hid]->Fill(val,weight);
  }
  inline void fillh2f(int hid, double valx, double valy,double weight) {
    if (harray[m_systype][hid] == 0) {
      return;
    }
    //if (fabs(valx)==ERRORVAL || fabs(valy)==ERRORVAL) cout << "ERROR filling 2d histo " << hid << " with value " << ERRORVAL << endl;
    ((TH2F*)harray[m_systype][hid])->Fill(valx,valy,weight);
  }
  inline void fillh3f(int hid, double valx, double valy, double valz, double weight) {
    if (harray[m_systype][hid] == 0) {
      return;
    }
    //if (fabs(valx)==ERRORVAL || fabs(valy)==ERRORVAL) cout << "ERROR filling 2d histo " << hid << " with value " << ERRORVAL << endl;
    ((TH3F*)harray[m_systype][hid])->Fill(valx,valy,valz,weight);
  }
  //

  void bumpIdentityInMigrationMatrix(double percentage_merged, TH2F& inputHisto);
  // Correction histograms
  //_______________________________________________________________________
  
  TH2F* mTrkEff_pt_eta;
  TH1F* mSecondaries_pt;
  TH2F* mOutside_pt_eta;
  TH1F* mVtx_merge_prob;


  TH1F* m_sb_N;
  TH1F* m_sb_phi;
  TH1F* m_sb_eta;
  TH1F* m_sb_pT;

  TH1F* m_pu_z;
  TH1F* m_pu_vtxN;
  TH2F* m_pu_vtxDz;
  TH2F* m_pu_vtxzPP;
  TH2F* m_pu_zN;
  TH1F* m_pu_ntrack;
  TH2F* m_pu_Dz_z;
  TH2F* m_pu_Dz_N;
  TH2F* m_pu_d_z;

  TH1F* m_vtx_res;

  vector<double> mkRdmTrack();
  vector<double> mirrorTrack(vector<double>);
  bool m_sidebands_loaded;

  // To be read HBOM correction histograms for resolution/smearing calculation
  TH1D* m_HBOM_AvgPt;
  TH1D* m_HBOM_BeamThrust;
  TH1D* m_HBOM_FParameter;
  TH1D* m_HBOM_Minor;
  TH1D* m_HBOM_Ntrk;
  TH1D* m_HBOM_Spherocity;
  TH1D* m_HBOM_SumPt;
  TH1D* m_HBOM_Thrust;


  // Transient variables 
  //_______________________________________________________________________
  
  int    m_systype;
  double evt_weight_gen;
  double evt_weight_reco;

  TVector3 lep1_trk_3vec, lep2_trk_3vec;
  double   lep1_cone_sumpt, lep2_cone_sumpt;
  int   lep1_cone_n, lep2_cone_n;

  EventShapes trkSelES_nominal;
  EventShapes mcES_nominal;

  vector<double> trackweights;
  
  // PileUp Library
  //_______________________________________________________________________
  vector < vector< vector<double> > > dynPileUpLibrary, dynPileUpLibrary_MCtruth;

  // Imagiro branch variables
  UInt_t EventNumber;
  double EventWeight;
  double Ntrk;
  double SumPt;
  double SumPtTrans;
  double AvgPt;
  double EtaVB;
  double PtVB;
  double PtVB_up;
  double PtVB_down;
  double Thrust;
  double ThrustPhi;
  double Minor;
  double FParameter;
  double Spherocity;
  double BeamThrust;

  // More stuff
  vector<double> mc_vx_z;
};


#endif