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

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

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

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

#include "TClonesArray.h"

#include "ExRootAnalysis/ExRootClasses.h"

#include "ExRootAnalysis/ExRootTreeReader.h"

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

using namespace std;

/*
LHC Olympics format discription from http://www.jthaler.net/olympicswiki/doku.php?id=lhc_olympics:data_file_format

    * The first column of each row is just a counter that labels the object.
    * The event begins with a row labelled "0"; this row contains the event number and the triggering information. The last row of the event is always the missing transverse momentum (MET).
    * The second column of each row gives the type of object being listed [0, 1, 2, 3, 4, 6 = photon, electron, muon, hadronically-decaying tau, jet, missing transverse energy].
    * The next three columns give the pseudorapidity, the azimuthal angle, and the transverse momentum of the object.
    * The sixth column gives the invariant mass of the object.
    * The seventh column gives the number of tracks associated with the object; in the case of a lepton, this number is multiplied by the charge of the lepton.
    * The eighth column is 1 or 2 for a jet that has been "tagged" as containing a b-quark (actually a heavy flavor tag that sometimes indicates c-quarks), otherwise it is 0. For muons, the integer part of this number is the identity of the jet (see column 1) that is closest ot this muon in Delta R.
    * The ninth column is the ratio of the hadronic versus electromagnetic energy deposited in the calorimeter cells associated with the object. For muons to the left of the decimal point is the summed pT in a R=0.4 cone (excluding the muon). To the right of the decimal point is etrat, which is a percentage between .00 and .99. It is the ratio of the transverse energy in a 3x3 grid surrounding the muon to the pT of the muon.
*/

struct LHCOlympicsObject
{
  enum {maxIntParam = 2, maxDblParam = 9};

  Int_t intParam[maxIntParam];
  Double_t dblParam[maxDblParam];
};

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

class LHCOlympicsWriter
{
public:
  LHCOlympicsWriter(ExRootTreeReader *treeReader,
                    ofstream *outputStream);
  ~LHCOlympicsWriter();

  void ProcessEvent();

private:

  void ResetLHCOlympicsObject();
  void WriteLHCOlympicsObject();

  void AnalyseEvent();

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

  void AnalyseMissingET();

  LHCOlympicsObject fCurrentObject;

  Long64_t fTriggerWord, fEventNumber;

  ExRootTreeReader *fTreeReader;
  ofstream *fOutputStream;

  TClonesArray *fBranchEvent;
  TClonesArray *fBranchPhoton;
  TClonesArray *fBranchElectron;
  TClonesArray *fBranchMuon;
  TClonesArray *fBranchTau;
  TClonesArray *fBranchJet;
  TClonesArray *fBranchMissingET;

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

};

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

LHCOlympicsWriter::LHCOlympicsWriter(ExRootTreeReader *treeReader,
                                     ofstream *outputStream) :
  fTriggerWord(0), fEventNumber(1), fTreeReader(0), fOutputStream(0)
{
  fTreeReader = treeReader;
  fOutputStream = outputStream;

  // information about reconstructed event
  fBranchEvent = fTreeReader->UseBranch("Event");
  // reconstructed photons
  fBranchPhoton = fTreeReader->UseBranch("Photon");
  fItPhoton = fBranchPhoton->MakeIterator();
  // reconstructed electrons
  fBranchElectron = fTreeReader->UseBranch("Electron");
  fItElectron = fBranchElectron->MakeIterator();
  // reconstructed muons
  fBranchMuon = fTreeReader->UseBranch("Muon");
  fItMuon = fBranchMuon->MakeIterator();
  // reconstructed hadronically-decaying tau leptons
  fBranchTau = fTreeReader->UseBranch("Tau");
  fItTau = fBranchTau->MakeIterator();
  // reconstructed jets
  fBranchJet = fTreeReader->UseBranch("Jet");
  fItJet = fBranchJet->MakeIterator();
  // missing transverse energy
  fBranchMissingET = fTreeReader->UseBranch("MissingET");
}

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

LHCOlympicsWriter::~LHCOlympicsWriter()
{
}

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

void LHCOlympicsWriter::ProcessEvent()
{
  fCurrentObject.intParam[0] = 0;

  AnalyseEvent();

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

  AnalyseMissingET();
}

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

void LHCOlympicsWriter::ResetLHCOlympicsObject()
{
  int i;
  for(i = 1; i < LHCOlympicsObject::maxIntParam; ++i)
  {
    fCurrentObject.intParam[i] = 0;
  }

  for(i = 0; i < LHCOlympicsObject::maxDblParam; ++i)
  {
    fCurrentObject.dblParam[i] = 0.0;
  }
}

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

void LHCOlympicsWriter::WriteLHCOlympicsObject()
{
  int i;
  for(i = 0; i < LHCOlympicsObject::maxIntParam; ++i)
  {
    (*fOutputStream) << fCurrentObject.intParam[i] << "\t";
  }

  for(i = 0; i < LHCOlympicsObject::maxDblParam; ++i)
  {
    (*fOutputStream) << fCurrentObject.dblParam[i] << "\t";
  }

  (*fOutputStream) << endl;

  ++fCurrentObject.intParam[0];
}

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

void LHCOlympicsWriter::AnalyseEvent()
{
  ExRootEvent *element;

  element = static_cast<ExRootEvent*>(fBranchEvent->At(0));

  (*fOutputStream) << fCurrentObject.intParam[0] << "\t";
  (*fOutputStream) << element->Number << "\t";
  (*fOutputStream) << element->Trigger << endl;

  ++fCurrentObject.intParam[0];
}

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

void LHCOlympicsWriter::AnalysePhotons()
{
  ExRootPhoton *element;

  fItPhoton->Reset();
  while((element = static_cast<ExRootPhoton*>(fItPhoton->Next())))
  {
    ResetLHCOlympicsObject();

    fCurrentObject.intParam[1] = 0;

    fCurrentObject.dblParam[0] = element->Eta;
    fCurrentObject.dblParam[1] = element->Phi;
    fCurrentObject.dblParam[2] = element->PT;

    fCurrentObject.dblParam[6] = element->EhadOverEem;

    WriteLHCOlympicsObject();
  }
}

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

void LHCOlympicsWriter::AnalyseElectrons()
{
  ExRootElectron *element;

  fItElectron->Reset();
  while((element = static_cast<ExRootElectron*>(fItElectron->Next())))
  {
    ResetLHCOlympicsObject();

    fCurrentObject.intParam[1] = 1;

    fCurrentObject.dblParam[0] = element->Eta;
    fCurrentObject.dblParam[1] = element->Phi;
    fCurrentObject.dblParam[2] = element->PT;

    fCurrentObject.dblParam[4] = element->Ntrk * element->Charge;

    fCurrentObject.dblParam[6] = element->EhadOverEem;

    WriteLHCOlympicsObject();
  }
}

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

void LHCOlympicsWriter::AnalyseMuons()
{
  ExRootMuon *element;

  fItMuon->Reset();
  while((element = static_cast<ExRootMuon*>(fItMuon->Next())))
  {
    ResetLHCOlympicsObject();

    fCurrentObject.intParam[1] = 2;

    fCurrentObject.dblParam[0] = element->Eta;
    fCurrentObject.dblParam[1] = element->Phi;
    fCurrentObject.dblParam[2] = element->PT;
    
    fCurrentObject.dblParam[4] = element->Ntrk * element->Charge;

    fCurrentObject.dblParam[5] = element->JetIndex;

    fCurrentObject.dblParam[6] = element->PTiso + element->ETiso;

    WriteLHCOlympicsObject();
  }
}

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

void LHCOlympicsWriter::AnalyseTaus()
{
  ExRootTau *element;

  fItTau->Reset();
  while((element = static_cast<ExRootTau*>(fItTau->Next())))
  {
    ResetLHCOlympicsObject();

    fCurrentObject.intParam[1] = 3;

    fCurrentObject.dblParam[0] = element->Eta;
    fCurrentObject.dblParam[1] = element->Phi;
    fCurrentObject.dblParam[2] = element->PT;
    
    fCurrentObject.dblParam[4] = element->Ntrk * element->Charge;

    fCurrentObject.dblParam[6] = element->EhadOverEem;
    
    WriteLHCOlympicsObject();
  }
}

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

void LHCOlympicsWriter::AnalyseJets()
{
  ExRootJet *element;

  fItJet->Reset();
  while((element = static_cast<ExRootJet*>(fItJet->Next())))
  {
    ResetLHCOlympicsObject();

    fCurrentObject.intParam[1] = 4;

    fCurrentObject.dblParam[0] = element->Eta;
    fCurrentObject.dblParam[1] = element->Phi;
    fCurrentObject.dblParam[2] = element->PT;
    
    fCurrentObject.dblParam[3] = element->Mass;

    fCurrentObject.dblParam[4] = element->Ntrk;
    
    fCurrentObject.dblParam[5] = element->BTag;

    fCurrentObject.dblParam[6] = element->EhadOverEem;
    
    WriteLHCOlympicsObject();
  }
}

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

void LHCOlympicsWriter::AnalyseMissingET()
{
  ExRootMissingET *element;

  element = static_cast<ExRootMissingET*>(fBranchMissingET->At(0));

  ResetLHCOlympicsObject();

  fCurrentObject.intParam[1] = 6;

  fCurrentObject.dblParam[1] = element->Phi;
  fCurrentObject.dblParam[2] = element->MET;

  WriteLHCOlympicsObject();
}

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

int main(int argc, char *argv[])
{
  char *appName = "ExRootLHCOlympicsWriter";

  if(argc != 3)
  {
    cout << " Usage: " << appName << " input_file" << " output_file" << endl;
    cout << " input_file - input file in ROOT format." << endl;
    cout << " output_file - output file in LHEF format," << endl;
    return 1;
  }

  int appargc = 1;
  char *appargv[] = {appName};
  TApplication app(appName, &appargc, appargv);

  // Open a stream connected to an output file:
  ofstream outputFileStream(argv[2]);

  if(!outputFileStream.is_open())
  {
    cerr << "** ERROR: Can't open '" << argv[2] << "' for output" << endl;
    return 1;
  }

  TChain *chain = new TChain("LHCO");
  chain->Add(argv[1]);

  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 LHC Olympics converter:
    LHCOlympicsWriter *writer = new LHCOlympicsWriter(treeReader, &outputFileStream);

    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;
      }
  
      writer->ProcessEvent();

      progressBar.Update(entry);
    }
    progressBar.Finish();
    
    delete writer;
  }

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

  delete treeReader;
  delete chain;
}