//
//  Sample test program for running EvtGen
//

#include "EvtGen/EvtGen.hh"

#include "EvtGenBase/EvtParticle.hh"
#include "EvtGenBase/EvtParticleFactory.hh"
#include "EvtGenBase/EvtPatches.hh"
#include "EvtGenBase/EvtPDL.hh"
#include "EvtGenBase/EvtRandom.hh"
#include "EvtGenBase/EvtReport.hh"
#include "EvtGenBase/EvtHepMCEvent.hh"
#include "EvtGenBase/EvtStdlibRandomEngine.hh"
#include "EvtGenBase/EvtAbsRadCorr.hh"
#include "EvtGenBase/EvtDecayBase.hh"

#ifdef EVTGEN_EXTERNAL
#include "EvtGenExternal/EvtExternalGenList.hh"
#endif

#include <iostream>
#include <string>
#include <list>

#include "TROOT.h"
#include "TFile.h"
#include "TH1F.h"
#include "TNtuple.h"

int main(int argc, char** argv) {
  // ================
  // == USER TUPLE ==
  // ================
  //unsigned long int NEvts = 1E5;
  //unsigned long int NEvts = 1E7;
  unsigned long int NEvts = 5E6;
  TFile *file=new TFile("EvtGenTuple.root", "RECREATE");

  long int EVT_ID = 0;
  double B_TRUETAU = -999, B_TRUEM = -999, B_PE = -999, B_PX = -999, B_PY = -999, B_PZ = -999;
  double gamma_PE = -999, gamma_PX = -999, gamma_PY = -999, gamma_PZ = -999;
  double KPLUS_PE = -999, KPLUS_PX = -999, KPLUS_PY = -999, KPLUS_PZ = -999;
  int B_TRUEID = 0, B_MIXED = -999, B_NDAUG = -999;
  TTree *tuple = new TTree("tuple","tuple");
  tuple->Branch("EVT_ID",&EVT_ID);
  tuple->Branch("B_TRUETAU",&B_TRUETAU);
  tuple->Branch("B_TRUEM",&B_TRUEM);
  tuple->Branch("B_TRUEID",&B_TRUEID);
  tuple->Branch("B_MIXED",&B_MIXED);
  tuple->Branch("B_NDAUG",&B_NDAUG);
  /*
  tuple->Branch("B_PE",&B_PE);
  tuple->Branch("B_PX",&B_PX);
  tuple->Branch("B_PY",&B_PY);
  tuple->Branch("B_PZ",&B_PZ);
  tuple->Branch("gamma_PE",&gamma_PE);
  tuple->Branch("gamma_PX",&gamma_PX);
  tuple->Branch("gamma_PY",&gamma_PY);
  tuple->Branch("gamma_PZ",&gamma_PZ);
  tuple->Branch("KPLUS_PE",&KPLUS_PE);
  tuple->Branch("KPLUS_PX",&KPLUS_PX);
  tuple->Branch("KPLUS_PY",&KPLUS_PY);
  tuple->Branch("KPLUS_PZ",&KPLUS_PZ);
  */

  //TH1F* hlifetime = new TH1F("hlifetime","Proper time distribution generated with EvtGen", 100, 0., 10.);
  //TH1F* hspin = new TH1F("hspin","Spin degrees of freedom", 100, 1., 3.);
  // TH1F* hspinD1 = new TH1F("hspinD1","Spin degrees of freedom Daughter 1", 100, 0.,10.);

  EvtParticle* parent(0);

  EvtStdlibRandomEngine eng;
  EvtRandom::setRandomEngine((EvtRandomEngine*)&eng);

  EvtAbsRadCorr* radCorrEngine = 0;
  std::list<EvtDecayBase*> extraModels;

#ifdef EVTGEN_EXTERNAL
  EvtExternalGenList genList;
  radCorrEngine = genList.getPhotosModel();
  extraModels = genList.getListOfModels();
#endif

  //Initialize the generator - read in the decay table and particle properties
  EvtGen myGenerator("../DECAY_2010.DEC","../evt.pdl", (EvtRandomEngine*)&eng,
  		     radCorrEngine, &extraModels);

  //If I wanted a user decay file, I would read it in now.
  myGenerator.readUDecay("Bs2phigamma.DEC");

  static EvtId B_s0 = EvtPDL::getId(std::string("B_s0sig"));
  static EvtId antiB_s0 = EvtPDL::getId(std::string("anti-B_s0sig"));

  unsigned long int nEvents(NEvts);

  // Loop to create nEvents, starting from an B_s0
  for (unsigned long i = 0; i < nEvents; i++) {

    //std::cout<<"Event number "<<i<<std::endl;
    if (fmod(i,1000)==0) std::cout<<"Event number "<<i<<std::endl;
    
    // Set up the parent particle
    EvtVector4R pInit((EvtPDL::getMass(B_s0)), 0., 0., 0.);
    /*   
    if( EvtRandom::Flat() < 0.5 ) {
      parent = EvtParticleFactory::particleFactory(B_s0, pInit);
    }
    else {
      parent = EvtParticleFactory::particleFactory(antiB_s0, pInit);
    }
    */
    parent = EvtParticleFactory::particleFactory(B_s0, pInit);
    //parent = EvtParticleFactory::particleFactory(antiB_s0, pInit);

    parent->setDiagonalSpinDensity();
  

    // Generate the event
    myGenerator.generateDecay(parent);
    
    // Write out the results
    EvtHepMCEvent theEvent;
    theEvent.constructEvent(parent);
    HepMC::GenEvent* genEvent = theEvent.getEvent();
    if(i<5) genEvent->print(std::cout);
    
    // ================
    // == USER TUPLE ==
    // ================
    EVT_ID = i;
    B_TRUETAU = parent->getLifetime()/EvtConst::c*1E12;
    B_TRUEM = parent->mass();
    
    EvtPDL *evtpdl;
    B_TRUEID = evtpdl->getStdHep( parent->getId() );

    int id = evtpdl->getStdHep( parent->getDaug(0)->getId() );
    (id == 531 || id == -531) ? B_MIXED = 1 : B_MIXED = 0;

    B_NDAUG = parent->getNDaug();
    B_PE = parent->getP4Lab().get(0);
    B_PX = parent->getP4Lab().get(1);
    B_PY = parent->getP4Lab().get(2);
    B_PZ = parent->getP4Lab().get(3);

    EvtParticle *ph, *v1, *kplus, *kminus;
    if(B_MIXED) {
      ph = parent->getDaug(0)->getDaug(1);
      v1 = parent->getDaug(0)->getDaug(0);
      kplus = v1->getDaug(0);
      kminus = v1->getDaug(1);
    }
    else {
      ph = parent->getDaug(1);
      v1 = parent->getDaug(0);
      kplus = v1->getDaug(0);
      kminus = v1->getDaug(1);
    }
    gamma_PE = ph->getP4Lab().get(0);
    gamma_PX = ph->getP4Lab().get(1);
    gamma_PY = ph->getP4Lab().get(2);
    gamma_PZ = ph->getP4Lab().get(3);
    KPLUS_PE = kplus->getP4Lab().get(0);
    KPLUS_PX = kplus->getP4Lab().get(1);
    KPLUS_PY = kplus->getP4Lab().get(2);
    KPLUS_PZ = kplus->getP4Lab().get(3);
    
    tuple->Fill();

    parent->deleteTree();

  }

  file->Write();
  file->Close();

  return 0;

}