//____________________________________________________________________________
/*!

\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){
  this->Initialize();
}
//___________________________________________________________________________
GGenerateEventGenie::GGenerateEventGenie(GenParam * GenPar, bool EvtGeneration):GGenerateEvent(GenPar){
  if(EvtGeneration)this->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(){
// GENIE initialization done in the main with function InitializeGenie,
// this is rquired since GENIE 3.2.0
// this method is mantained because it is a virtual method of the base class





  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);

  ostringstream MonFileName;
  MonFileName << fGenPar->EvFilePrefix  << "." << fGenPar->RunNu << ".status";
  fMcjMonitor->CustomizeFilename(MonFileName.str());
  fMcjMonitor->SetRefreshRate(1000);

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

  if(fGenPar->GenCode.compare("CORSIKA")==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__
#if __GENIE_RELEASE_CODE__ >= GRELCODE(3,2,0)
  else if (fGenPar->WeightOpt==2) fMcjDriver->ForceInteraction();
  fMcjDriver->SetPmaxSafetyFactor(1.);
#endif
  fMcjDriver->UseSplines();
  fMcjDriver->SetEventGeneratorList(fGenPar->EventGeneratorList);

  if(fGenPar->GenCode.compare("CORSIKA")==0) fMcjDriver->KeepOnThrowingFluxNeutrinos(false);


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







}
//________________________________________________________________________________________
void GGenerateEventGenie::Configure(int NTot){
  fNTot=NTot;

  fMcjDriver->Configure();

  return;
}
//________________________________________________________________________________________
void GGenerateEventGenie::CleanEvent(void){
  fSeaEvent->Reset();
  if(!(fGenPar->GenCode.compare("CORSIKA")==0)) { // we do not create event for CORSIKA processing
    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->GenCode.compare("CORSIKA")==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->GenCode.compare("CORSIKA")==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();
  TVector3  pol;
  p->GetPolarization(pol);
  SeaEvent->Neutrino.Set(1, p->Status(), p->Pdg(), p->FirstMother(), 0., p->GetP4(), p->GetX4(), &SeaEvent->Vertex, 0., 0., 0., pol);

  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();
  p->GetPolarization(pol);
  SeaEvent->PrimaryLepton.Set(1, p->Status(), p->Pdg(), p->FirstMother(), p->Mass(), p->GetP4(), p->GetX4(), &SeaEvent->Vertex, 0., 0, 0., pol);

  SeaEvent->PScale=fMcjDriver->GlobProbScale()*event->Weight();

// 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);

  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)
{
  bool aretheretrks=false;
  if(fGenPar->GenCode.compare("CORSIKA")==0){  // CORSIKA tracks
    while(1){
      fFileDriver->ReadFileEvent(); // reads in fFileDriver->OtherTracks; and fFileDriver->SecondaryTracks;
      fSeaEvent->DistaMax = fFileDriver->fDistaMax;
      if (fFileDriver->FullRun) {EoF = true; break;}
      fSeaEvent->iEvtMC=fFileDriver->EventNumber; // save the actual MC id from CORSIKA
      if (fGenPar->NoPropagationCORSIKA) fSeaEvent->LepInCan=true; // disable propagation and save everything at sea
      else fSeaEvent->LepInCan = this->PropagateTracks();   // propagate tracks up the can
      aretheretrks = this->WriteTracks( fSeaEvent->LepInCan );
      if( aretheretrks ) {
        fFileDriver->ComputeWeights(fSeaEvent); // does not really compute, rather copies the weights from file driver to the event
        this->FillEvent(); // adds the tracks to be written to the event
        fNEvt++;
        fSeaEvent->iEvt=fNEvt;
        return; //leave the function to save the event
      }
      else this->CleanEvent(); // free the memory
    }
  }
  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->Tracks = this->GetGeneratorTracks();
      fSeaEvent->LepInCan = this->PropagateTracks();   // 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(){
// 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.SetStatus(5); //change status of initial state nuclei
    else if ( track.Status==1 && track.PdgCode>1e6) track.SetStatus(kIStIntermediateState); //change status of remnant nuclei
    track.E        = p->Energy();
    Momentum = sqrt((p->Px())*(p->Px()) + (p->Py())*(p->Py()) + (p->Pz())*(p->Pz()));  // 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() * 1e-24 * 1e9; //convert from yoctosecond to nanoseconds
    if (this->PosInCan(track.V1,track.V2,track.V3)==0) track.Status *= -1; //particle outside the can
    p->GetPolarization(track.Polarization);

/*
    if (p->PolzIsSet()) {
      cout<<"check pol "<< p->PolzIsSet()<<" "<<p->Pdg()<<endl;
      cout<<"direction "<<track.D1<<" "<<track.D2<<" "<<track.D3<<endl;
      cout<<"polarization "<<track.Polarization.X()<<" "<<track.Polarization.Y()<<" "<<track.Polarization.Z()<<endl;
      cout<<"polarization angles "<<p->PolzPolarAngle()*180./3.1415<<" "<< p->PolzAzimuthAngle()*180./3.1415<< endl;
      cout<<"direction angles "<<acos(track.D3)*180./3.1415<<" "<<atan2( track.D2, track.D1 )*180./3.1415<<endl;
    }
*/
    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

  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

  fGenPar->XSecWater = this->CalcXsec (MediaComp, fGenPar->EventGeneratorList);

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

  map<int, TGraph> XSecMedia;

  string EventGeneratorList;

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

  XSecMedia = this->CalcXsec (MediaComp, EventGeneratorList);

  return XSecMedia;

}

//___________________________________________________________________________
map<int, TGraph> GGenerateEventGenie::CalcXsec(map<int, double> MediaComp, string EventGeneratorList){

  map<int, TGraph> XSecMedia;

  int NuType;
  int TgtPdgCode;
  double TgtPercentage;


  GEVGDriver evg_driver;
  InitialState init_state;
  const Spline * splsum;
  
  evg_driver.SetEventGeneratorList(EventGeneratorList);

  int kNP = 1000; // = 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();

////      kNP = (int) (15 * TMath::Log10(rangeE.max-rangeE.min));
//      kNP = (int) (15 * (TMath::Log10(rangeE.max)-TMath::Log10(rangeE.min)));
//      kNP = TMath::Max(kNP,30);

      evg_driver.CreateSplines(kNP, rangeE.max, true);
      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;
      }
      XSecTot.SetPoint(iKP,x,ytot);
    }


    if(EventGeneratorList=="CC" && XSecTot.GetHistogram()->GetMinimum()==0){ // remove points below the CC threshold
      int iBin=-1;
      double x0=0, y0=0,y1=0;
      for (int iKP=0; iKP<kNP; iKP++){
        XSecTot.GetPoint(iKP,x,y);
	if(x!=0 && y !=0) {iBin = iKP; x0=x; y0=y; break;}
      }
      
      if(iBin>0){
        XSecTot.GetPoint(iBin+1,x,y1);
        if(fabs(log10(y0)-log10(y1))>1) XSecTot.SetPoint(iBin,x0,0);
      }
    }

    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;
}