//____________________________________________________________________________
/*!

\class   genie::GGenerateEventGenie.cxx

\brief   A GENIE generation driver for a neutrino telescope.


\author  Carla Distefano <distefano_c@lns.infn.it>
         LNS-INFN, Catania
         Included modifications by Alfonso Garcia <alfonsog@nikhef.nl>
         NIKHEF, Amsterdam (January 2019)

\created July 19, 2017

\cpright Copyright (c) 2017-2019, The KM3NeT Collaboration
         For the full text see $GSEAGEN/LICENSE

*/
//____________________________________________________________________________


#include <cassert>
#include <cmath>
#include <vector>

#include "SeaNuDrivers/GGenerateEventGenie.h"

//___________________________________________________________________________
GGenerateEventGenie::GGenerateEventGenie(GenParam * GenPar):GGenerateEvent(GenPar)
{

  fNTot = 0;

  this->Initialize();

  this->CalcXSecWater();

  fGenPar->XSecTotSW=this->CalcXsecTot(fGenPar->SeaWaterComp);
  fGenPar->XSecTotSR=this->CalcXsecTot(fGenPar->RockComp);
  fGenPar->XSecTotMantle=this->CalcXsecTot(fGenPar->MantleComp);
  fGenPar->XSecTotCore=this->CalcXsecTot(fGenPar->CoreComp);

  // create the GENIE monte carlo job driver
  fMcjDriver = new GMCJDriver;
  if(gSystem->Getenv("SYSTLIST"))fSeaSyst = new GSeaSyst;
  fMcjMonitor = new GMCJMonitor(fGenPar->RunNu);

  fMcjDriver->UseFluxDriver(dynamic_cast<GFluxI *>(fFluxDriver));

  if(fGenPar->GenMode.compare("EVT")==0 || fGenPar->GenMode.compare("BIN")==0) fMcjDriver->UseFluxDriver(dynamic_cast<GFluxI *>(fFileDriver));
  else fMcjDriver->UseFluxDriver(dynamic_cast<GFluxI *>(fFluxDriver));

  if      (fGenPar->WeightOpt==0) fMcjDriver->ForceSingleProbScale();
  //Forcing interaction only work with GENIE-HEDIS version
  #ifdef _KM3NET_ENABLED__
  else if (fGenPar->WeightOpt==2) fMcjDriver->ForceInteraction();
  fMcjDriver->SetPmaxSafetyFactor(1.);
  #endif
  fMcjDriver->UseSplines();
  fMcjDriver->SetEventGeneratorList(fGenPar->EventGeneratorList);

  if(fGenPar->GenMode.compare("EVT")==0 || fGenPar->GenMode.compare("BIN")==0) fMcjDriver->KeepOnThrowingFluxNeutrinos(false);


  if(fGenPar->NativeOutput){
    fNtpw = new NtpWriter(kNFGHEP, fGenPar->RunNu);
    fNtpw->CustomizeFilenamePrefix(fGenPar->EvFilePrefix);
    fNtpw->Initialize();
  }


}
//___________________________________________________________________________
GGenerateEventGenie::~GGenerateEventGenie(){

  if(fGenPar->NativeOutput){
    fNtpw->Save();
  }


  if(fMcjDriver){fMcjDriver=0; delete fMcjDriver;}
  if(fMcjMonitor){fMcjMonitor=0; delete fMcjMonitor;}
  if(fSeaSyst){fSeaSyst=0; delete fSeaSyst;}
  if(event != nullptr){event=0; delete event;}
  if(fNtpw){fNtpw=0; delete fNtpw;}

}
//________________________________________________________________________________________
void GGenerateEventGenie::Initialize(){

#ifdef _GENIEVERSIONGTEQ3__
  RunOpt::Instance()->SetEventGeneratorList(fGenPar->EventGeneratorList);
  RunOpt::Instance()->SetTuneName(fGenPar->GenTune);
  if ( ! RunOpt::Instance()->Tune() ) {
    LOG("GGenerateEventGenie", pFATAL) << " No TuneId in RunOption";
    exit(-1);
  }
  RunOpt::Instance()->BuildTune();
#endif

  // Initialization of random number generators, cross-section table,
  // messenger thresholds, cache file
  utils::app_init::MesgThresholds(RunOpt::Instance()->MesgThresholdFiles());
  utils::app_init::CacheFile(RunOpt::Instance()->CacheFile());
  utils::app_init::RandGen((long int)fGenPar->RanSeed);
  utils::app_init::XSecTable(fGenPar->InpXSecFile, false);
  // Set GHEP print level
  GHepRecord::SetPrintLevel(RunOpt::Instance()->EventRecordPrintLevel());
  return;
}
//________________________________________________________________________________________
void GGenerateEventGenie::Configure(int NTot){
  fNTot=NTot;

  fMcjDriver->Configure();

  return;
}
//________________________________________________________________________________________
void GGenerateEventGenie::CleanEvent(void){
  fSeaEvent->Reset();
  if(event)delete event;
}
//________________________________________________________________________________________
void GGenerateEventGenie::SetGeometry(string RootGeomFile){

  double GeomLUnits = genie::utils::units::UnitFromString("cm");
  double GeomDUnits = genie::utils::units::UnitFromString("g_cm3");

  if(fGenPar->PropMode)GeomLUnits = genie::utils::units::UnitFromString("m");


#ifdef __GENIE_GEOM_DRIVERS_ENABLED__

  // creating & configuring a root geometry driver
  geometry::ROOTGeomAnalyzer * rgeom =
          new geometry::ROOTGeomAnalyzer(RootGeomFile);
  rgeom -> SetLengthUnits  (GeomLUnits);
  rgeom -> SetDensityUnits (GeomDUnits);
  rgeom -> SetTopVolName   ("");
  #ifdef _KM3NET_ENABLED__
  rgeom -> SetMaxPlSafetyFactor (1.05);
  #endif
  // getting the bounding box dimensions along z so as to set the
  // appropriate upstream generation surface for the JPARC flux driver
  TGeoVolume * topvol = rgeom->GetGeometry()->GetTopVolume();
  if(!topvol) {
    LOG("GGenerateEventGenie", pFATAL) << " ** Null top ROOT geometry volume!";
    gAbortingInErr = true;
    exit(1);
  }

  // casting to the GENIE geometry driver interface
  GeomAnalyzerI * geom_driver = dynamic_cast<GeomAnalyzerI *> (rgeom);
  fMcjDriver->UseGeomAnalyzer(geom_driver);

#else
  LOG("GGenerateEventGenie", pFATAL) << "You need to enable the GENIE geometry drivers first!";
  LOG("GGenerateEventGenie", pFATAL) << "Use --enable-geom-drivers at the configuration step.";
  gAbortingInErr = true;
  exit(1);
#endif

  return;
}

//________________________________________________________________________________________
void GGenerateEventGenie::InitGenBin(double Emin, double Emax, double RL, double RT, double X0, double Y0, double Z0)
{
  fNTot+=this->GetNFluxNeutrinos();

  if(fGenPar->GenMode.compare("EVT")==0 || fGenPar->GenMode.compare("BIN")==0)
    fFileDriver->InitGenBin(Emin,Emax,RL,RT,X0,Y0,Z0);
  else
    fFluxDriver->InitGenBin(Emin,Emax,RL,RT,X0,Y0,Z0);

  return;
}

//________________________________________________________________________________________
double GGenerateEventGenie::GetNFluxNeutrinos(void)
{
  if(fGenPar->GenMode.compare("EVT")==0 || fGenPar->GenMode.compare("BIN")==0)
    return fFileDriver->NFluxNeutrinos();
  else
    return fFluxDriver->NFluxNeutrinos();
}
//________________________________________________________________________________________
void GGenerateEventGenie::FillSeaEvent(GSeaEvent * SeaEvent){

  Interaction * interaction = event->Summary();

  SeaEvent->Bx=event->Summary()->Kine().x(true);
  SeaEvent->By=event->Summary()->Kine().y(true);

  SeaEvent->ScattId = interaction->ProcInfoPtr()->ScatteringTypeId();
  SeaEvent->InterId = interaction->ProcInfoPtr()->InteractionTypeId();

  SeaEvent->XSec  = (1/units::cm2) * event->XSec();
  SeaEvent->DXSec = (1/units::cm2) * event->DiffXSec();

  SeaEvent->TargetZ   = interaction->InitStatePtr()->TgtPtr()->Z();
  SeaEvent->TargetA   = interaction->InitStatePtr()->TgtPtr()->A();

  GHepParticle * p = event->Probe();
  SeaEvent->Neutrino.Set(1, p->Status(), p->Pdg(), p->FirstMother(), 0., p->GetP4(), p->GetX4(), &SeaEvent->Vertex, 0. , 0. );

  SeaEvent->SWXSec    = this->GetXSecSW(p->Pdg(), p->Energy());
  SeaEvent->WaterXSec = this->GetXSecWater(p->Pdg(), p->Energy());

  if(SeaEvent->WaterXSec>0.)SeaEvent->WaterIntLen=1./(kNA*SeaEvent->WaterXSec)*1.E-6; // interation lenght in pure water in mwe

  p = event->FinalStatePrimaryLepton();
  SeaEvent->PrimaryLepton.Set(1, p->Status(), p->Pdg(), p->FirstMother(), p->Mass(), p->GetP4(), p->GetX4(), &SeaEvent->Vertex, 0. , 0. );

  SeaEvent->PScale=fMcjDriver->GlobProbScale();

// compute the weigths
  this->ComputeWeights(SeaEvent);

// compute systematics errors if activated
  if(gSystem->Getenv("SYSTLIST"))SeaEvent->SysWgt=fSeaSyst->ComputeWeights(event);

  this->FillEvent();
  SeaEvent->EarthNeutrino = gEarthNeutrino;

  SeaEvent->Translate(fGenPar->OriginX,fGenPar->OriginY,fGenPar->OriginZ);  // move to the geometry driver

  return;
}
//________________________________________________________________________________________
void GGenerateEventGenie::ComputeWeights(GSeaEvent * SeaEvt)
{
  fFluxDriver->ComputeWeights(SeaEvt);

  if (fGenPar->WeightOpt!=0) {
    SeaEvt->GenWeight *= event->Weight();
    SeaEvt->EvtWeight *= event->Weight();
  }

  LOG("GGenerateEventGenie", pDEBUG) << "GlobalGenWeight = " << fGenPar->GlobalGenWeight;
  LOG("GGenerateEventGenie", pDEBUG) << "PEarth          = " << SeaEvt->PEarth;
  LOG("GGenerateEventGenie", pDEBUG) << "PScale          = " << SeaEvt->PScale;
  LOG("GGenerateEventGenie", pDEBUG) << "NInter          = " << SeaEvt->NInter;
  LOG("GGenerateEventGenie", pDEBUG) << "flux            = " << SeaEvt->EvtWeight/SeaEvt->GenWeight;
  LOG("GGenerateEventGenie", pDEBUG) << "GenWeight       = " << SeaEvt->GenWeight;
  LOG("GGenerateEventGenie", pDEBUG) << "EvtWeight       = " << SeaEvt->EvtWeight;
  LOG("GGenerateEventGenie", pDEBUG) << "GenieW          = " << event->Weight();

  event->SetWeight(SeaEvt->EvtWeight);
}
//________________________________________________________________________________________
void GGenerateEventGenie::GenerateEvent(void)
{

  if(fGenPar->GenMode.compare("EVT")==0 || fGenPar->GenMode.compare("BIN")==0){  // CORSIKA tracks
    int Chances=0;
    vector<GSeaTrack> Tracks;
    while(1){

      if ( Chances == 0 ) {
        Chances = fGenPar->ChancesPerShower; // resetting the number of chances for the next shower
        event = new EventRecord();  // we create the event to delete it....
        Tracks = fFileDriver->ReadFileEvent();
        if (Tracks.size()==0) {EoF = true; break;}
      }

      fSeaEvent->VerInCan = 0;
      fSeaEvent->LepInCan = this->PropagateTracks( Tracks );   // propagate tracks up the can

      bool aretheretrks = this->WriteTracks( fSeaEvent->LepInCan );
      if( aretheretrks ) {
        fFileDriver->ComputeWeights(fSeaEvent);
        this->FillEvent();
        fNEvt++;
        fSeaEvent->iEvt=fNEvt;
        fSeaEvent->NRetries = fGenPar->ChancesPerShower - Chances; // save the number of retries
        Chances=0;
      }
      else if ( Chances > 0 ) { // showers that missed the can get more chances
        LOG("GGenerateEventGenie", pNOTICE) << "Event missed the can [ " << Chances-1 << " chances left ]";
        Chances-=1;
        continue;
      }

      if (aretheretrks) return; //leave the function to save the event

      this->CleanEvent();
    }

  }
  else{

    while(1){

      event = fMcjDriver->GenerateEvent();

      fMcjMonitor->Update(this->GetNFluxNeutrinos(),event);

      if(this->GetNFluxNeutrinos()>fNTot)break;

      fSeaEvent->Vertex=this->GetVertex();

      fSeaEvent->VerInCan = this->PosInCan(fSeaEvent->Vertex.X(),fSeaEvent->Vertex.Y(),fSeaEvent->Vertex.Z());   // check the vertex is inside the can
      assert( (fGenPar->HBedRock==0. && fSeaEvent->VerInCan!=2) || fGenPar->HBedRock>0. );
      fSeaEvent->LepInCan = this->PropagateTracks( this->GetGeneratorTracks() );   // propagate tracks up the can

      bool aretheretrks = this->WriteTracks( fSeaEvent->LepInCan );
      if( aretheretrks ) {
        this->FillSeaEvent(fSeaEvent);
        fNEvt++;
        fSeaEvent->iEvt=fNEvt;
      }

      if( gEarthNeutrino.PdgCode != 0 ){
        double ene = gEarthNeutrino.E;
        double px  = gEarthNeutrino.D1*ene;
        double py  = gEarthNeutrino.D2*ene;
        double pz  = gEarthNeutrino.D3*ene;
        TLorentzVector *p4 = new TLorentzVector(px,py,pz,ene);
        TLorentzVector *x4 = new TLorentzVector(gEarthNeutrino.V1,gEarthNeutrino.V2,gEarthNeutrino.V3,gEarthNeutrino.T);
        if      (fGenPar->GenMode.compare("DIFFUSE")==0) flux_driver_diffuse -> SetNeutrino(p4,x4,gEarthNeutrino.PdgCode,false);
        else if (fGenPar->GenMode.compare("POINT")==0)   flux_driver_point   -> SetNeutrino(p4,x4,gEarthNeutrino.PdgCode,false);
      }

      if (aretheretrks) return; //leave the function to save the event

      this->CleanEvent();
    }
  }


}
//________________________________________________________________________________________
vector<GSeaTrack> GGenerateEventGenie::GetGeneratorTracks(void){
// return a vector with genie particles at the interaction final state

  gEarthNeutrino.Reset();

  vector<GSeaTrack> GeneratorTracks;
  GSeaTrack track;

  double Momentum;

  GHepParticle * p = 0;

  TObjArrayIter piter(event);

  while((p =(GHepParticle *)piter.Next())) {

    track.Id       = event->ParticlePosition(p);
    track.Status   = p->Status();
    track.PdgCode  = p->Pdg();
    track.MotherId = p->FirstMother();
    if      ( track.Status==0 && track.Id!=0 )      track.Status = 5; //change status of initial state nuclei
    else if ( track.Status==1 && track.PdgCode>1e6) track.Status = kIStIntermediateState; //change status of remnant nuclei
    track.E        = p->Energy();
    Momentum = sqrt(pow(p->Px(),2)+pow(p->Py(),2)+pow(p->Pz(),2));  // to avoid problems with negative sqrt (this can happen for very energetic particles).
    if ( Momentum > 0 ) {
      track.D1 = p->Px()/Momentum;
      track.D2 = p->Py()/Momentum;
      track.D3 = p->Pz()/Momentum;
    }
    else track.D1 = track.D2 = track.D3 = 0;
    track.V1 = event->Vertex()->X()+p->Vx()*1E-15; // interaction vertex in m, particle position in fm with respect the target
    track.V2 = event->Vertex()->Y()+p->Vy()*1E-15;
    track.V3 = event->Vertex()->Z()+p->Vz()*1E-15;
    track.Length = 0.;
    track.T  = p->Vt() * 1e9; //convert from seconds to nanoseconds
    if (this->PosInCan(track.V1,track.V2,track.V3)==0) track.Status *= -1; //particle outside the can

    GeneratorTracks.push_back(track);

    if (fGenPar->PropMode && fSeaEvent->VerInCan!=0) {
      if ( track.Status==-1 && (fabs(track.PdgCode)==12 || fabs(track.PdgCode)==14 || fabs(track.PdgCode)==16) ) {
        if (track.E > gEarthNeutrino.E ) gEarthNeutrino = track;
      }
    }

  }

  return GeneratorTracks;
}
//________________________________________________________________________________________
void GGenerateEventGenie::CalcXSecWater(void){

// total cross section in pure water

  int NuType;
  int TgtPdgCode;
  double TgtPercentage;

  map<int, double> MediaComp;
  MediaComp.insert(map<int, double>::value_type(1000080160,0.8888));            // O16
  MediaComp.insert(map<int, double>::value_type(1000010010,0.1112));            // H1

  GEVGDriver evg_driver;
  InitialState init_state;
  const Spline * splsum;

  evg_driver.SetEventGeneratorList(fGenPar->EventGeneratorList);

  const int kNP = 1000;

  size_t nComp=MediaComp.size();
  TGraph * gr[nComp];
  double PA[nComp];

  double x,y,ytot;

  for(unsigned i=0; i<fGenPar->NuList.size(); i++){

    unsigned iComp=0.;

    NuType=fGenPar->NuList[i];

    map<int,double>::iterator iter = MediaComp.begin();
    for ( ; iter != MediaComp.end(); ++iter) {
      TgtPdgCode=iter->first;
      TgtPercentage=iter->second;
      PA[iComp]=TgtPercentage/(double)pdg::IonPdgCodeToA(TgtPdgCode);

      init_state.SetPdgs(TgtPdgCode, NuType);
      evg_driver.Configure(init_state);

      Range1D_t rangeE = evg_driver.ValidEnergyRange();
      evg_driver.CreateSplines(kNP, rangeE.max, true); //splines has been computed up to from 1e2 to 1e10 with 200knots (avoid zero at edges)
      evg_driver.CreateXSecSumSpline (kNP, rangeE.min, rangeE.max);

      splsum = evg_driver.XSecSumSpline();
      gr[iComp] = splsum->GetAsTGraph(kNP,true,true,1.,1./units::m2);
      iComp++;
    }

    TGraph XSecTot;
    for (int iNP=0; iNP<kNP; iNP++){
      ytot=0.;
      for(iComp=0; iComp<nComp; iComp++){
        gr[iComp]->GetPoint(iNP,x,y);
        ytot+=PA[iComp]*y*x;
      }
      if(ytot>0)XSecTot.SetPoint(iNP,x,ytot);     // sometimes extremes have ziro xsec!
    }

    fGenPar->XSecWater.insert(map<int,TGraph>::value_type(NuType,XSecTot));

  }

  return;
}
//__________________________________________________________________________
map<int, TGraph> GGenerateEventGenie::CalcXsecTot(map<int, double> MediaComp){

  map<int, TGraph> XSecMedia;

  int NuType;
  int TgtPdgCode;
  double TgtPercentage;


  GEVGDriver evg_driver;
  InitialState init_state;
  const Spline * splsum;

  if      (fGenPar->EventGeneratorList.compare(2,5,"HEDIS")==0) evg_driver.SetEventGeneratorList("CCHEDISGLRES");
  else if (fGenPar->EventGeneratorList.compare(0,5,"GLRES")==0) evg_driver.SetEventGeneratorList("CCHEDISGLRES");
  else evg_driver.SetEventGeneratorList("CC");

  const int kNP = 1000;

  size_t nComp=MediaComp.size();
  TGraph * gr[nComp];
  double PA[nComp];

  double x,y,ytot;

  unsigned iComp=0;

  for(unsigned i=0; i<fGenPar->NuList.size(); i++){

    iComp=0.;

    NuType=fGenPar->NuList[i];

    map<int,double>::iterator iter = MediaComp.begin();

    for ( ; iter != MediaComp.end(); ++iter) {

      TgtPdgCode=iter->first;
      TgtPercentage=iter->second;

      PA[iComp]=TgtPercentage/(double)pdg::IonPdgCodeToA(TgtPdgCode);

      init_state.SetPdgs(TgtPdgCode, NuType);
      evg_driver.Configure(init_state);

      Range1D_t rangeE = evg_driver.ValidEnergyRange();
      evg_driver.CreateSplines(kNP, rangeE.max, true); //splines has been computed up to from 1e2 to 1e10 with 200knots (avoid zero at edges)
      evg_driver.CreateXSecSumSpline (kNP, rangeE.min, rangeE.max);

      splsum = evg_driver.XSecSumSpline();
      gr[iComp] = splsum->GetAsTGraph(kNP,true,true,1.,1./units::m2);

      iComp++;
    }

    TGraph XSecTot;
    for (int iKP=0; iKP<kNP; iKP++){
      ytot=0.;
      for(unsigned iiComp=0; iiComp<nComp; iiComp++){
        gr[iiComp]->GetPoint(iKP,x,y);
        ytot+=PA[iiComp]*y*x;
      }
      if(ytot>0)XSecTot.SetPoint(iKP,x,ytot);     // sometimes extremes have ziro xsec!
    }
    XSecMedia.insert(map<int,TGraph>::value_type(NuType,XSecTot));

  }

  return XSecMedia;

}

//___________________________________________________________________________
void GGenerateEventGenie::WriteNative(GSeaEvent * SeaEvt){

  if(fGenPar->NativeOutput)
    fNtpw->AddEventRecord(SeaEvt->iEvt, event);

  return;
}