#include <iostream>
#include <fstream>

#include "TLorentzVector.h"

#include "Utilities/ExRootAnalysis/interface/ExRootTreeWriter.h"
#include "Utilities/ExRootAnalysis/interface/ExRootTreeBranch.h"
#include "Utilities/ExRootAnalysis/interface/BlockClasses.h"

#include "interface/LHEFConverter.h"
#include "Utilities/ExRootAnalysis/interface/LHEF.h"

using namespace std;


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

void LHEFConverter::AnalyseEvent(LHEF::Reader *reader, ExRootTreeBranch *branch, const Long64_t eventNumber)
{
  const LHEF::HEPEUP &hepeup = reader->hepeup;

  TRootLHEFEvent *element;

  element = (TRootLHEFEvent*) branch->NewEntry();

  element->Number = eventNumber;
  element->Nparticles = hepeup.NUP;
  element->ProcessID = hepeup.IDPRUP;
  element->Weight = hepeup.XWGTUP;
  element->ScalePDF = hepeup.SCALUP;
  element->CouplingQED = hepeup.AQEDUP;
  element->CouplingQCD = hepeup.AQCDUP;
}

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

void LHEFConverter::AnalyseParticles(LHEF::Reader *reader, ExRootTreeBranch *branch)
{
  const LHEF::HEPEUP &hepeup = reader->hepeup;
  Double_t signPz;
  TLorentzVector momentum;
  TRootLHEFParticle *element;

  for(Int_t particle = 0; particle < hepeup.NUP; ++particle)
  {
    element = (TRootLHEFParticle*) branch->NewEntry();

    element->PID = hepeup.IDUP[particle];
    element->Status = hepeup.ISTUP[particle];
    element->Mother1 = hepeup.MOTHUP[particle].first;
    element->Mother2 = hepeup.MOTHUP[particle].second;
    element->ColorLine1 = hepeup.ICOLUP[particle].first;
    element->ColorLine2 = hepeup.ICOLUP[particle].second;
    element->Px = hepeup.PUP[particle][0];
    element->Py = hepeup.PUP[particle][1];
    element->Pz = hepeup.PUP[particle][2];
    element->E = hepeup.PUP[particle][3];
    element->M = hepeup.PUP[particle][4];

    momentum.SetPxPyPzE(element->Px, element->Py, element->Pz, element->E);
    element->PT = momentum.Perp();
    signPz = (element->Pz >= 0.0) ? 1.0 : -1.0;
    element->Eta = element->PT == 0.0 ? signPz*999.9 : momentum.Eta();
    element->Phi = momentum.Phi();
    element->Rapidity = element->PT == 0.0 ? signPz*999.9 : momentum.Rapidity();
    element->LifeTime = hepeup.VTIMUP[particle];
    element->Spin = hepeup.SPINUP[particle];
  }
}

LHEFConverter::~LHEFConverter()
{
}

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

LHEFConverter::LHEFConverter(const string& inputFileList, const string& outputFileName) {

  ExRootTreeWriter *treeWriter = new ExRootTreeWriter(outputFileName, "GEN");
  
  // generated event from LHEF
  ExRootTreeBranch *branchEvent = treeWriter->NewBranch("Event", TRootLHEFEvent::Class());
  
  // generated partons from LHEF
  ExRootTreeBranch *branchParticle = treeWriter->NewBranch("Particle", TRootLHEFParticle::Class());
  
  // Open a stream connected to an event file:
  ifstream infile(inputFileList.c_str());
  string filename;
  if(!infile.is_open()) {
      cerr << "** ERROR: Can't open '" << inputFileList << "' for input" << endl;
      exit(1);
  }

  while(1) { // parsing the list of files
    
      infile >> filename;
      if(!infile.good()) break;
        ifstream checking_the_file(filename.c_str());
        if(!checking_the_file.good()) { cout << filename << ": file not found\n"; continue;}
        else checking_the_file.close();

      // Create the Reader object:
      LHEF::Reader *inputReader = new LHEF::Reader(filename);
      
      cout << "** Calculating number of events to process. Please wait..." << endl;
      Long64_t allEntries = inputReader->getNumberOfEvents();
      cout << "** Input file contains " << allEntries << " events" << endl;
      
      if(allEntries > 0) {
	  // Loop over all events
	  Long64_t entry = 0;
	  while(inputReader->readEvent()) {
	      treeWriter->Clear();
	      AnalyseEvent(inputReader, branchEvent, entry + 1);
	      AnalyseParticles(inputReader, branchParticle);
	      treeWriter->Fill();
	      ++entry;
	  }
      }
  }
  treeWriter->Write();
  
  cout << "** Exiting..." << endl;
  
  delete treeWriter;
  //delete inputReader;
}