//-------------------------------------------------
//
//   Class: MuonAnalysis
//
//   Description: fill simple ROOT tree and perform analysis
//                  
//                
//   Authors :
//   N. Neumeister            CERN PH
//
//--------------------------------------------------
#include "Utilities/Configuration/interface/Architecture.h"

//-----------------------
// This Class's Header --
//-----------------------

#include "Workspace/MuonAnalysis.h"

//---------------
// C++ Headers --
//---------------

#include <iostream>
#include <iomanip>
#include <vector>
#include <string>

//-------------------------------
// Collaborating Class Headers --
//-------------------------------

#include "TFile.h"
#include "TH1.h"
#include "TTree.h"

#include "CARF/G3Event/interface/G3EventProxy.h"
#include "CARF/SimEvent/interface/SimEvent.h"
#include "CARF/BaseSimEvent/interface/SimVertex.h"
#include "CARF/Reco/interface/RecQuery.h"
#include "CARF/Reco/interface/RecCollection.h"
#include "CARF/Reco/interface/FilteredRecCollection.h"

#include "GeneratorInterface/RecParticle/interface/RawRecParticleTypeFilter.h"
#include "GeneratorInterface/RecParticle/interface/RawRecParticleStableFilter.h"
#include "GeneratorInterface/RecParticle/interface/RawRecParticle.h"

#include "Trigger/L1Trigger/interface/L1Trigger.h"
#include "Muon/PersistentMuon/interface/RecMuon.h"

const int MuonAnalysis::MAXGEN;
const int MuonAnalysis::MAXREC;

//----------------
// Constructors --
//----------------
MuonAnalysis::MuonAnalysis(const string& filename) {

  // first define a ROOT file
  theFile = new TFile(filename.c_str(),"RECREATE");
  
  // book histograms
  theHisto1 = new TH1F("H1","muon pt spectrum (gen)",100,0.0,200.0); 
  theHisto2 = new TH1F("H2","muon pt spectrum (rec)",100,0.0,200.0);
  theHisto3 = new TH1F("H3","invariant mass (gen)",100,70.0,120.0);
  theHisto4 = new TH1F("H4","invariant mass (rec)",100,70.0,120.0);
  
  // book a simple ROOT Tree
  theTree = new TTree("T1","Muon analysis tree");

  // GENERAL block
  theTree->Branch("Run", &m_Run, "Run/I");
  theTree->Branch("Event", &m_Event, "Event/I");
  theTree->Branch("Accept", &m_Accept, "Accept/B");
  theTree->Branch("Single", &m_SingleMu, "Single/B");
  theTree->Branch("Di", &m_DiMu, "Di/B");
  
  // GEANT block
  theTree->Branch("Ngen", &m_Ngen, "Ngen/I");
  theTree->Branch("Pxgen", m_Pxgen, "Pxgen[Ngen]/F");
  theTree->Branch("Pygen", m_Pygen, "Pygen[Ngen]/F");
  theTree->Branch("Pzgen", m_Pzgen, "Pzgen[Ngen]/F");
  theTree->Branch("Pgen", m_Pgen, "Pgen[Ngen]/F");
  theTree->Branch("Ptgen", m_Ptgen, "Ptgen[Ngen]/F");
  theTree->Branch("Etagen", m_Etagen, "Etagen[Ngen]/F");
  theTree->Branch("Phigen", m_Phigen, "Phigen[Ngen]/F");
  theTree->Branch("Chagen", m_Chagen, "Chagen[Ngen]/I");
  theTree->Branch("Vxgen", m_Vxgen, "Vxgen[Ngen]/F");
  theTree->Branch("Vygen", m_Vygen, "Vygen[Ngen]/F");
  theTree->Branch("Vzgen", m_Vzgen, "Vzgen[Ngen]/F");

  // reconstructed muon block
  theTree->Branch("Nrec", &m_Nrec, "Nrec/I");
  theTree->Branch("Nrhitr", m_Nrhitr, "Nrhitr[Nrec]/I");       
  theTree->Branch("Dofr", m_Dofr, "Dofr[Nrec]/I");
  theTree->Branch("Chi2r", m_Chi2r, "Chi2r[Nrec]/F");
  theTree->Branch("X", m_X, "X[Nrec]/F");
  theTree->Branch("Y", m_Y, "Y[Nrec]/F");
  theTree->Branch("Z", m_Z, "Z[Nrec]/F");
  theTree->Branch("Cha", m_Cha, "Cha[Nrec]/I");
  theTree->Branch("Px", m_Px, "Px[Nrec]/F");
  theTree->Branch("Py", m_Py, "Py[Nrec]/F");
  theTree->Branch("Pz", m_Pz, "Pz[Nrec]/F");
  theTree->Branch("P", m_P, "P[Nrec]/F");
  theTree->Branch("Pt", m_Pt, "Pt[Nrec]/F");
  theTree->Branch("Eta", m_Eta, "Eta[Nrec]/F");
  theTree->Branch("Phi", m_Phi, "Phi[Nrec]/F");

}


//--------------
// Destructor --
//--------------
MuonAnalysis::~MuonAnalysis() {

  theFile->cd();
  theTree->Write();
  theHisto1->Write();
  theHisto2->Write();
  theHisto3->Write();
  theHisto4->Write();
  theFile->Close();

  delete theFile;

}


//--------------
// Operations --
//--------------


//
// analyze event
//
void MuonAnalysis::analyzeEvent(G3EventProxy* ev) {

  //
  // run and event number
  //
  m_Run = ev->simSignal()->id().runNumber();
  m_Event = ev->simSignal()->id().eventInRun();

  //
  // first check the Level-1 trigger decision
  //
  L1Trigger l1trig;
  m_Accept = l1trig.decision();
  m_SingleMu = l1trig.decisionWord().element(0);
  m_DiMu = l1trig.decisionWord().element(1);
  cout.setf(ios::boolalpha);
  cout << "Level-1 Trigger decision: " << m_Accept << endl;
  cout << "Single/Di-muon trigger: " 
       << m_SingleMu << "/" 
       << m_DiMu << endl;
  cout << endl;

  //
  // analyze Monte Carlo information
  //
  analyzeMCEvent();

  //
  // analyze simulated event (GEANT information)
  //
  analyzeSimEvent(ev);

  //
  // analyze reconstructed muons
  //
  analyzeMuons();

  //
  // analyze di-muon events
  //
  analyzeDiMuons();

  //
  // fill ROOT tree
  //
  theTree->Fill();

}


//
// analyze Monte Carlo event information
//
void MuonAnalysis::analyzeMCEvent() {

  //
  // get MC muons
  // 
  RecQuery q("RawParticleRecon");
  q.setParameter("FakeDet","Particles");
  RawRecParticleStableFilter stableFilter;
  RawRecParticleTypeFilter muonFilter("mu-","mu+");
  Capri::MultipleFilterAnd<RawRecParticle> f;
  f.add(&stableFilter);
  f.add(&muonFilter);
  FilteredRecCollection<RawRecParticle> mcMuons(q,f);
  FilteredRecCollection<RawRecParticle>::const_iterator it = mcMuons.begin();
  for ( it = mcMuons.begin(); it != mcMuons.end(); it++ ) { 
    (**it).print();
  }

}


//
// analyze GEANT information
//
void MuonAnalysis::analyzeSimEvent(G3EventProxy* ev) {

  //
  // get GEANT muons from signal event (no pile-up)
  //
  const float ptmin  = 0.5;
  const float etamax = 2.4;

  cout << "GEANT muons: " << endl;
  int igen = 0;
  int ntracks = ev->simSignal()->tracks()->size();
  for ( int it = 0; it < ntracks; it++ ) {
    SimTrack tk(ev->simSignal()->track(it));
    int ipart = tk.type();
    if ( abs(ipart) != 13 ) continue;   // select muons
    const HepLorentzVector& mom = tk.momentum();
    float pt  = mom.perp();
    float eta = mom.pseudoRapidity();
    if ( pt > ptmin && fabs(eta) < etamax ) {
      int charge = static_cast<int>(tk.charge());
      float phi = mom.phi();
      if ( phi < 0 ) phi = 2*M_PI + phi;
      const HepLorentzVector& vtx = tk.vertex().position();

      cout.setf(ios::showpoint);
      cout << setiosflags(ios::showpoint | ios::fixed)
           << setw(2) << igen+1 << " : "
           << "pt = "     << setw(5) << setprecision(1) << pt << " GeV  "
           << "charge = " << setw(2) << charge << " "
           << "eta = "    << setw(6) << setprecision(3) << eta << "  "
           << "phi = "    << setw(5) << setprecision(3) << phi << " rad \t"
           << "vertex = (" << vtx.x() << "," << vtx.y() << "," << vtx.z() << ")  " << endl;

      theHisto1->Fill(pt);

      if ( igen < MAXGEN ) {
        m_Pxgen[igen] = mom.px();
        m_Pygen[igen] = mom.py();
        m_Pzgen[igen] = mom.pz();
        m_Pgen[igen] = mom.rho();
        m_Ptgen[igen] = pt;
        m_Etagen[igen] = eta;
        m_Phigen[igen] = phi;
        m_Chagen[igen] = charge;
        m_Vxgen[igen] = vtx.x();
        m_Vygen[igen] = vtx.y();
        m_Vzgen[igen] = vtx.z();
        igen++;
      }
    }  
  }
  
  m_Ngen = igen;

}


//
// analyze reconstructed muons
//
void MuonAnalysis::analyzeMuons() {

  RecQuery q("GlobalMuonReconstructor");
  RecCollection<RecMuon> recmuons(q);
  RecCollection<RecMuon>::iterator muon = recmuons.begin();

  int nrec = recmuons.size();
  cout << "Number of reconstructed muons: " << nrec << endl;
       
  if ( nrec > 0 ) cout << "Reconstructed muons: " << endl; 
  int idx = 0;
 
  while ( muon != recmuons.end() ) {
    if ( idx >= MAXREC ) continue;
    int nmeas       = (**muon).foundHits();
    int dof         = (**muon).degreesOfFreedom();
    double chi2     = (**muon).chiSquared();
    double normChi2 = (**muon).normalisedChiSquared();

    m_Nrhitr[idx] = nmeas;
    m_Dofr[idx] = dof;
    m_Chi2r[idx] = chi2;
     
    cout << setiosflags(ios::showpoint | ios::fixed) 
         << setw(2) << idx+1 << '\t' << "number of measurements = " 
         << setw(2) << nmeas << "  " 
         << "Chi2/DoF = " << setw(6) << setprecision(2) << chi2
         << "/" << setw(2) << dof
         << " = " << setw(6) << setprecision(3) << normChi2 << endl;

    //
    // state at closeat approach to vertex in the transverse plane
    //
    FreeTrajectoryState* traj = (**muon).ipState();

    if ( traj ) {

      GlobalVector mom = traj->momentum();
      GlobalPoint pos  = traj->position();
      int charge       = traj->charge();
  
      float pt    = mom.perp();
      float eta   = mom.eta();
      float phi   = mom.phi();
      if ( phi < 0 ) phi = 2*M_PI + phi;
        
      m_X[idx] = pos.x();
      m_Y[idx] = pos.y();
      m_Z[idx] = pos.z();
      m_Px[idx] = mom.x();
      m_Py[idx] = mom.y();
      m_Pz[idx] = mom.z();
      m_P[idx] = mom.mag();
      m_Pt[idx] = pt;
      m_Eta[idx] = eta;
      m_Phi[idx] = phi;
      
      cout << setiosflags(ios::showpoint | ios::fixed)
           << setw(2) << idx+1 << '\t'
           << "pt = " << setw(5) << setprecision(1) << pt << " GeV  "
           << "charge = " << setw(2) << charge << " "
           << "eta = " << setw(6) << setprecision(3) << eta << "  "
           << "phi = " << setw(5) << setprecision(3) << phi << " rad  "
           << "position = " << pos << endl;

      theHisto2->Fill(pt);
    }
    else {
      m_X[idx] = 0.;
      m_Y[idx] = 0.;
      m_Z[idx] = 0.;
      m_Px[idx] = 0.;
      m_Py[idx] = 0.;
      m_Pz[idx] = 0.;
      m_P[idx] = 0.;
      m_Pt[idx] = 0.;
      m_Eta[idx] = 0.;
      m_Phi[idx] = 0.;
    }    
   
    cout << *(**muon).muonTrack() << endl;
    cout << *(**muon).trackerTrack() << endl;

    idx++; 
    muon++;
  }

  m_Nrec = idx;
  
}


//
// analyze di-muon events
//
void MuonAnalysis::analyzeDiMuons() {

  const double m = 0.105658357;

  //
  // analyze sim event
  // 
  if ( m_Ngen < 2 ) return;

  for ( int i = 0; i != m_Ngen; i++ ) {
    if ( m_Ptgen[i] < 5.0 ) continue;
    HepLorentzVector v1(m_Pxgen[i],m_Pygen[i],m_Pzgen[i],sqrt(m_Pgen[i]*m_Pgen[i] + m*m));
    for ( int j = i+1; j != m_Ngen; j++ ) {
      if ( m_Ptgen[j] < 5.0 ) continue;
      if (m_Chagen[i] == m_Chagen[j] ) continue;
      HepLorentzVector v2(m_Pxgen[j],m_Pygen[j],m_Pzgen[j],sqrt(m_Pgen[j]*m_Pgen[j] + m*m));
      try {
        double invm = v1.invariantMass(v2);      
        theHisto3->Fill(invm);
      } catch (...) {
	cout << " exception caught" << endl;
      }
    }
  }  
  
  //
  // analyze rec event
  //
  if ( m_Nrec < 2 ) return;
  
  for ( int i = 0; i != m_Nrec; i++ ) {
    if ( m_Pt[i] < 5.0 ) continue;
    HepLorentzVector v1(m_Px[i],m_Py[i],m_Pz[i],sqrt(m_P[i]*m_P[i] + m*m));
    for ( int j = i+1; j != m_Nrec; j++ ) {
      if ( m_Pt[j] < 5.0 ) continue;
      if (m_Cha[i] == m_Cha[j] ) continue;
      HepLorentzVector v2(m_Px[j],m_Py[j],m_Pz[j],sqrt(m_P[j]*m_P[j] + m*m));
      try {
        double invm = v1.invariantMass(v2);      
        theHisto4->Fill(invm);
      } catch (...) {
	cout << " exception caught" << endl;
      }
    }
  }

}