#include <iostream>
#include <fstream>
#include <sstream>
#include <map>

#include "TROOT.h"
#include "TApplication.h"

#include "TFile.h"
#include "TChain.h"
#include "TString.h"
#include "TClonesArray.h"

#include "TH2.h"
#include "THStack.h"
#include "TLegend.h"
#include "TPaveText.h"
#include "TLorentzVector.h"

#include "ExRootAnalysis/ExRootClasses.h"

#include "ExRootAnalysis/ExRootTreeReader.h"
#include "ExRootAnalysis/ExRootTreeWriter.h"
#include "ExRootAnalysis/ExRootTreeBranch.h"

#include "ExRootAnalysis/ExRootUtilities.h"
#include "ExRootAnalysis/ExRootProgressBar.h"

using namespace std;

//------------------------------------------------------------------------------

class MadWeightAnalysis
{
public:
  MadWeightAnalysis(ExRootTreeReader *treeReader,
                    ExRootTreeWriter *treeWriter);
  ~MadWeightAnalysis();

  void ProcessEvent();

private:

  void AnalyseEvent();

  void AnalysePhotons();
  void AnalyseElectrons();
  void AnalyseMuons();
  void AnalyseTaus();
  void AnalyseJets();

  void AnalyseMissingET();

  Long64_t fTriggerWord, fEventNumber;

  ExRootTreeReader *fTreeReader;
  ExRootTreeWriter *fTreeWriter;

  TClonesArray *fInputEvent;
  TClonesArray *fInputPhoton;
  TClonesArray *fInputElectron;
  TClonesArray *fInputMuon;
  TClonesArray *fInputTau;
  TClonesArray *fInputJet;
  TClonesArray *fInputMissingET;

  TIterator *fItPhoton;
  TIterator *fItElectron;
  TIterator *fItMuon;
  TIterator *fItTau;
  TIterator *fItJet;

  ExRootTreeBranch *fOutputEvent;
  ExRootTreeBranch *fOutputPhoton;
  ExRootTreeBranch *fOutputElectron;
  ExRootTreeBranch *fOutputMuon;
  ExRootTreeBranch *fOutputTau;
  ExRootTreeBranch *fOutputJet;
  ExRootTreeBranch *fOutputMissingET;
};

//------------------------------------------------------------------------------

MadWeightAnalysis::MadWeightAnalysis(ExRootTreeReader *treeReader,
                                     ExRootTreeWriter *treeWriter) :
  fTriggerWord(0), fEventNumber(1), fTreeReader(0), fTreeWriter(0)
{
  fTreeReader = treeReader;
  fTreeWriter = treeWriter;

  // information about reconstructed event
  fInputEvent = fTreeReader->UseBranch("Event");
  // reconstructed photons
  fInputPhoton = fTreeReader->UseBranch("Photon");
  fItPhoton = fInputPhoton->MakeIterator();
  // reconstructed electrons
  fInputElectron = fTreeReader->UseBranch("Electron");
  fItElectron = fInputElectron->MakeIterator();
  // reconstructed muons
  fInputMuon = fTreeReader->UseBranch("Muon");
  fItMuon = fInputMuon->MakeIterator();
  // reconstructed hadronically-decaying tau leptons
  fInputTau = fTreeReader->UseBranch("Tau");
  fItTau = fInputTau->MakeIterator();
  // reconstructed jets
  fInputJet = fTreeReader->UseBranch("Jet");
  fItJet = fInputJet->MakeIterator();
  // missing transverse energy
  fInputMissingET = fTreeReader->UseBranch("MissingET");


  // information about reconstructed event
  fOutputEvent = fTreeWriter->NewBranch("Event", ExRootEvent::Class());
  // reconstructed photons
  fOutputPhoton = fTreeWriter->NewBranch("Photon", ExRootPhoton::Class());
  // reconstructed electrons
  fOutputElectron = fTreeWriter->NewBranch("Electron", ExRootElectron::Class());
  // reconstructed muons
  fOutputMuon = fTreeWriter->NewBranch("Muon", ExRootMuon::Class());
  // reconstructed hadronically-decaying tau leptons
  fOutputTau = fTreeWriter->NewBranch("Tau", ExRootTau::Class());
  // reconstructed jets
  fOutputJet = fTreeWriter->NewBranch("Jet", ExRootJet::Class());
  // missing transverse energy
  fOutputMissingET = fTreeWriter->NewBranch("MissingET", ExRootMissingET::Class());

}

//------------------------------------------------------------------------------

MadWeightAnalysis::~MadWeightAnalysis()
{
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::ProcessEvent()
{
  AnalyseEvent();

  AnalysePhotons();
  AnalyseElectrons();
  AnalyseMuons();
  AnalyseTaus();
  AnalyseJets();

  AnalyseMissingET();
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseEvent()
{
  ExRootEvent *input, *output;

  input = static_cast<ExRootEvent*>(fInputEvent->At(0));
  
  output = static_cast<ExRootEvent*>(fOutputEvent->NewEntry());

  *output = *input;

}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalysePhotons()
{
  ExRootPhoton *input, *output;

  fItPhoton->Reset();
  while((input = static_cast<ExRootPhoton*>(fItPhoton->Next())))
  {
    // Example: select central photons
    if(TMath::Abs(input->Eta) < 1.0)
    {
      output = static_cast<ExRootPhoton*>(fOutputPhoton->NewEntry());
    
      *output = *input;
    }
  }
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseElectrons()
{
  ExRootElectron *input, *output;

  fItElectron->Reset();
  while((input = static_cast<ExRootElectron*>(fItElectron->Next())))
  {
    // Example: select central electrons
    if(TMath::Abs(input->Eta) < 1.0)
    {
      output = static_cast<ExRootElectron*>(fOutputElectron->NewEntry());
    
      *output = *input;
    }
  }
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseMuons()
{
  ExRootMuon *input, *output;

  fItMuon->Reset();
  while((input = static_cast<ExRootMuon*>(fItMuon->Next())))
  {
    // Example: select central muons
    if(TMath::Abs(input->Eta) < 1.0)
    {
      output = static_cast<ExRootMuon*>(fOutputMuon->NewEntry());
    
      *output = *input;
    }
  }
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseTaus()
{
  ExRootTau *input, *output;

  fItTau->Reset();
  while((input = static_cast<ExRootTau*>(fItTau->Next())))
  {
    // Example: select central taus
    if(TMath::Abs(input->Eta) < 1.0)
    {
      output = static_cast<ExRootTau*>(fOutputTau->NewEntry());
    
      *output = *input;
    }
  }
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseJets()
{
  ExRootJet *input, *output;

  fItJet->Reset();
  while((input = static_cast<ExRootJet*>(fItJet->Next())))
  {
    // Example: select central jets
    if(TMath::Abs(input->Eta) < 1.0)
    {
      output = static_cast<ExRootJet*>(fOutputJet->NewEntry());
    
      *output = *input;
    }
  }
}

//---------------------------------------------------------------------------

void MadWeightAnalysis::AnalyseMissingET()
{
  ExRootMissingET *input, *output;

  if(fInputMissingET->GetEntriesFast() > 0)
  {
    input = static_cast<ExRootMissingET*>(fInputMissingET->At(0));
  
  
    output = static_cast<ExRootMissingET*>(fOutputMissingET->NewEntry());
  
    *output = *input;
  }
}

//---------------------------------------------------------------------------

void Analysis(const char *inputFileName, const char *outputFileName)
{
  TChain *chain = new TChain("LHCO");
  chain->Add(inputFileName);

  ExRootTreeReader *treeReader = new ExRootTreeReader(chain);

  cout << "** Calculating number of events to process. Please wait..." << endl;
  Long64_t allEntries = treeReader->GetEntries();
  cout << "** Input file contains " << allEntries << " events" << endl;

  Long64_t entry;

  if(allEntries > 0)
  {
    // Create ROOT tree writer:
    TFile *outputFile = TFile::Open(outputFileName, "RECREATE");
    ExRootTreeWriter *treeWriter = new ExRootTreeWriter(outputFile, "LHCO");

    // Create analysis object:
    MadWeightAnalysis *analysis = new MadWeightAnalysis(treeReader, treeWriter);

    ExRootProgressBar progressBar(allEntries);
    // Loop over all events
    for(entry = 0; entry < allEntries; ++entry)
    {
      if(!treeReader->ReadEntry(entry))
      {
        cout << "** ERROR: cannot read event " << entry << endl;
        break;
      }

      treeWriter->Clear();

      analysis->ProcessEvent();

      treeWriter->Fill();

      progressBar.Update(entry);
    }
    progressBar.Finish();
    
    treeWriter->Write();

    delete analysis;
    delete treeWriter;
  }

  cout << "** Exiting..." << endl;

  delete treeReader;
  delete chain;
}