//____________________________________________________________________________
/*!

\class   genie::GPEarth

\brief   A class to evaluate the probability to transmission through the Earth
         The column depth is computed according to the Earth PREM model
         A. M. Dziewonski and D.L. Anderson, Physcis of the Earth and Planetary
         Interiors, 25 (1981) 297, Table 1

\author  Carla Distefano <distefano_c@lns.infn.it>
         LNS-INFN, Catania

\created October 17, 2012


\cpright  Copyright (c) 2012-2019, The KM3NeT Collaboration
          For the full text of the license see $GSEAGEN/LICENSE
*/
//____________________________________________________________________________


#include <cassert>

#include <TMath.h>
#include <TGraph.h>
#include <TSystem.h>
#include <TRotation.h>

#include <TGeoManager.h>
#include <TPostScript.h>


#ifdef _GENIEVERSIONGTEQ3__
#include "Framework/Messenger/Messenger.h"
#include "Framework/ParticleData/PDGUtils.h"
#else
#include "Messenger/Messenger.h"
#include "PDG/PDGUtils.h"
#endif

#include "SeaNuDrivers/GPEarth.h"

using namespace std;
using namespace genie;

//____________________________________________________________________________
GPEarth::GPEarth(GenParam * GenPar)
{

  fREarth= 6371;

  fGenPar = GenPar;

  // PREM model http://www.sciencedirect.com/science/article/pii/0031920181900467
  double RR[NLAYERS] = {1221.5,3480.,5701.,5771.,5971.,6151.,6346.6,6356.,6368.,fREarth};

  double EEarthCoeff[NLAYERS][4] = {
    {13.0885,0.,-8.8381,0.},
    {12.5815,-1.2638,-3.6426,-5.5281} ,
    {7.9565,-6.4761,5.5283,-3.0807},
    {5.3197,-1.4836,0.,0.},
    {11.2494,-8.0298,0.,0.},
    {7.1089,-3.8045,0.,0.},
    {2.691,0.6924,0.,0.},
    {2.9,0.,0.,0.},
    {2.65,0.,0.,0.},
    {1.02,0.,0.,0.},
  };

  string Comp[NLAYERS] = {"Core","Core","Mantle","Mantle","Mantle","Mantle","Mantle","Mantle","Rock","SeaWater"};

  for(int i=0; i<NLAYERS; i++)fR[i]=RR[i];
  fR[NLAYERS-2]=RR[NLAYERS-1]-fGenPar->SiteDepth/1.E3;  // change the rock radius with the site depth


  for(int i=0; i<NLAYERS-2; i++)
      for(int j=0; j<4; j++) fEarthCoeff[i][j]=EEarthCoeff[i][j];

  fEarthCoeff[NLAYERS-2][0]=fGenPar->RhoSR;     // change the last rock layer with the rock density used in the code
  fEarthCoeff[NLAYERS-1][0]=fGenPar->RhoSW;     // change the ocean density with the one used in the code


  fEarthCoeff[NLAYERS-2][1]=0.;                 // forsed to be zero because they foud with strange values in fIntLayer
  fEarthCoeff[NLAYERS-1][1]=0.;
  fEarthCoeff[NLAYERS-2][2]=0.;
  fEarthCoeff[NLAYERS-1][2]=0.;
  fEarthCoeff[NLAYERS-2][3]=0.;
  fEarthCoeff[NLAYERS-1][3]=0.;


  for(int i=0; i<NLAYERS; i++)fComp[i]=Comp[i];


  // Earth centre coordinates

  fXc=0.;
  fYc=0.;
  fZc=-RR[NLAYERS-2]+fGenPar->Can1/1.E3;  // in km
  fRSeaBottom=RR[NLAYERS-2];

  if(fGenPar->PropMode){
    this->fBuildPREMEarth();
    LOG("GPEarth", pINFO)<< "Building geometry for propagation through the Earth... done.";
  }
}
//___________________________________________________________________________
GPEarth::~GPEarth(){}
//___________________________________________________________________________
void GPEarth::fBuildPREMEarth(){



  vector<Layers> VecLayers;
  Layers layer;

  double rmean=0.;
  double step=500.;
  int iLayer=0;
  bool border=false;
  double r1=0;
  double r2=step;
  int iStep=0;

  while(1){
    if(r2>fR[iLayer]){
      border=true;
      r2=fR[iLayer];
    }

    layer.r1=r1;
    layer.r2=r2;
    rmean=(r2+r1)/2.;
    layer.rho=fEarthCoeff[iLayer][0]+fEarthCoeff[iLayer][1]*rmean/fREarth+fEarthCoeff[iLayer][2]*pow(rmean/fREarth,2)+fEarthCoeff[iLayer][3]*pow(rmean/fREarth,3);
    layer.Composition=fComp[iLayer];

    VecLayers.push_back(layer);

    if(border){
      r1=r2;
      iLayer++;
      border=false;
    }
    else{
      r1=(iStep+1)*step;
    }
    r2=(iStep+2)*step;
    iStep++;

    if(r1>=fR[NLAYERS-4])break;

  }

  for(int iiLayer=NLAYERS-3; iiLayer<NLAYERS; iiLayer++){       // constant density
    layer.r1=fR[iiLayer-1];
    layer.r2=fR[iiLayer];
    layer.rho=fEarthCoeff[iiLayer][0]+fEarthCoeff[iiLayer][1]*layer.r2/fREarth+fEarthCoeff[iiLayer][2]*pow(layer.r2/fREarth,2)+fEarthCoeff[iiLayer][3]*pow(layer.r2/fREarth,3);
    layer.Composition=fComp[iiLayer];
    VecLayers.push_back(layer);
  }


  // Define media and volumes

  TGeoManager * GeoManager =  new TGeoManager("VolGenGeo", "generation volume geometry");

  // vacuum
  TGeoMaterial * MatVacuum = new TGeoMaterial("MatVacuum");
  MatVacuum->SetA(0.);
  MatVacuum->SetZ(0.);
  MatVacuum->SetDensity(0.);

  TGeoMedium * Vacuum = new TGeoMedium("Vacuum", 0, MatVacuum);


  map<int,double>::iterator iter;

  TGeoMixture * LayerMix[VecLayers.size()];
  TGeoMedium * LayerMedium[VecLayers.size()];



  TGeoVolume *TopVolume = GeoManager->MakeTube("TopVolume",Vacuum,0.,fR[NLAYERS-1]*1E3, 2*fR[NLAYERS-1]*1.E3);
  GeoManager->SetTopVolume(TopVolume);

  TGeoVolume * Layer[VecLayers.size()];

  TGeoTranslation * Trans = new TGeoTranslation(0.,0.,fZc*1.E3);


  for(unsigned iiLayer=0; iiLayer<VecLayers.size(); iiLayer++){

    char ch[100];

    sprintf(ch,"_%d",iiLayer);

    string name=VecLayers[iiLayer].Composition;
    name.append(ch);

    if(VecLayers[iiLayer].Composition=="SeaWater"){
      LayerMix[iiLayer] = new TGeoMixture(name.c_str(),fGenPar->SeaWaterComp.size(),VecLayers[iiLayer].rho);
      iter = fGenPar->SeaWaterComp.begin();
      for( ; iter != fGenPar->SeaWaterComp.end(); ++iter)
        LayerMix[iiLayer]->AddElement(pdg::IonPdgCodeToA(iter->first), pdg::IonPdgCodeToZ(iter->first), iter->second);
    }
    else if(VecLayers[iiLayer].Composition=="Rock"){
      LayerMix[iiLayer] = new TGeoMixture(name.c_str(),fGenPar->RockComp.size(),VecLayers[iiLayer].rho);
      iter = fGenPar->RockComp.begin();
      for( ; iter != fGenPar->RockComp.end(); ++iter)
        LayerMix[iiLayer]->AddElement(pdg::IonPdgCodeToA(iter->first), pdg::IonPdgCodeToZ(iter->first), iter->second);
    }
    else if(VecLayers[iiLayer].Composition=="Mantle"){
      LayerMix[iiLayer] = new TGeoMixture(name.c_str(),fGenPar->MantleComp.size(),VecLayers[iiLayer].rho);
      iter = fGenPar->MantleComp.begin();
      for( ; iter != fGenPar->MantleComp.end(); ++iter)
        LayerMix[iiLayer]->AddElement(pdg::IonPdgCodeToA(iter->first), pdg::IonPdgCodeToZ(iter->first), iter->second);
    }
    else if(VecLayers[iiLayer].Composition=="Core"){
      LayerMix[iiLayer] = new TGeoMixture(name.c_str(),fGenPar->CoreComp.size(),VecLayers[iiLayer].rho);
      iter = fGenPar->CoreComp.begin();
      for( ; iter != fGenPar->CoreComp.end(); ++iter)
        LayerMix[iiLayer]->AddElement(pdg::IonPdgCodeToA(iter->first), pdg::IonPdgCodeToZ(iter->first), iter->second);
    }



    LayerMedium[iiLayer] = new TGeoMedium(name.c_str(), iiLayer+1, LayerMix[iiLayer]);

    Layer[iiLayer] = GeoManager->MakeSphere(name.c_str(),LayerMedium[iiLayer],VecLayers[iiLayer].r1*1.E3,VecLayers[iiLayer].r2*1.E3);
    TopVolume->AddNode(Layer[iiLayer], iiLayer+1, Trans);

  }


  string filename = "geometry-Earth.root";

  GeoManager->CloseGeometry();
  GeoManager->Export(filename.c_str());

  Vacuum=0; delete Vacuum;
  MatVacuum=0; delete MatVacuum;
  for(unsigned iiLayer=0; iiLayer<VecLayers.size(); iiLayer++){
    LayerMedium[iiLayer]=0; delete LayerMedium[iiLayer];
    LayerMix[iiLayer]=0; delete LayerMix[iiLayer];
  }
  GeoManager=0; delete GeoManager;



  return;
}

//___________________________________________________________________________
double GPEarth::EarthAbsorption(double X0, double Y0, double Z0, double l, double m, double n, double Energy, int neutrino_pdg){

   // X0,V0,Z0 neutrino vertex in km
   // l,m,n neutrino track direction cosines

  fParL = l;
  fParM = m;
  fParN = n;

  fX0 = X0/1E3;
  fY0 = Y0/1E3;
  fZ0 = Z0/1E3;

  fA = 1.;
  fB = 2.*((fX0-fXc)*fParL+(fY0-fYc)*fParM+(fZ0-fZc)*fParN);
  fC = TMath::Power((fX0-fXc),2.)+TMath::Power((fY0-fYc),2.)+TMath::Power((fZ0-fZc),2.);

  this->fCalcColumnDepth();

  double sigma_tot=0.;

  double SigmaTotSW=this->fGetXsecTot(Energy, neutrino_pdg, fGenPar->XSecTotSW);
  double SigmaTotSR=this->fGetXsecTot(Energy, neutrino_pdg, fGenPar->XSecTotSR);
  double SigmaTotMantle=this->fGetXsecTot(Energy, neutrino_pdg, fGenPar->XSecTotMantle);
  double SigmaTotCore=this->fGetXsecTot(Energy, neutrino_pdg, fGenPar->XSecTotCore);


  fSigmaMean=0.;

  for(unsigned i=0; i<fColumnDepthLayer.size();i++){
    if(fComp[fColumnDepthLayer[i].iLayer]=="SeaWater")sigma_tot=SigmaTotSW;
    if(fComp[fColumnDepthLayer[i].iLayer]=="Rock")sigma_tot=SigmaTotSR;
    if(fComp[fColumnDepthLayer[i].iLayer]=="Mantle")sigma_tot=SigmaTotMantle;
    if(fComp[fColumnDepthLayer[i].iLayer]=="Core")sigma_tot=SigmaTotCore;
    fSigmaMean+=sigma_tot*fColumnDepthLayer[i].ColumnDepth;
  }

  fSigmaMean*= 1.E9;
  fPEarth=exp(-TMath::Na()*fSigmaMean);  // in water equivalent, sigma in m^2, fColumnDepth in mwe
  fSigmaMean/=fColumnDepth;

  return fPEarth;
}
//________________________________________________________________________________________
void GPEarth::fCalcColumnDepth(void)
{

  fColumnDepthLayer.clear();
  fColumnDepth=0.;

  double ColInt;
  int iLayer=0;
  double T1=0.,T2=0.;

  ColumnDepthLayer ColLayer;

  this->fExtremes();

  if(fParN<=0.|| fB*fB-4*fA*(fC-TMath::Power(fR[NLAYERS-2],2))<=0.){    // down-going and horizontal (the track intersects only the ocean)
    iLayer=NLAYERS-1;
    T1=fVecInterPoint[0].T;
    T2=0.;
    ColInt = this->fIntLayer(iLayer,T1,T2);
    ColLayer.iLayer=iLayer;
    ColLayer.ColumnDepth=ColInt;

    fColumnDepthLayer.push_back(ColLayer);
  }
  else { // up-going

      for(unsigned i=0; i<fVecExtremPoint.size()-1; i++){  // no second layer in ocean
        iLayer=fVecExtremPoint[i].iLayer;
        T1=fVecExtremPoint[i].T1;
        T2=fVecExtremPoint[i].T2;

        ColInt = this->fIntLayer(iLayer,T1,T2);
        ColLayer.iLayer=iLayer;
        ColLayer.ColumnDepth=ColInt;
        fColumnDepthLayer.push_back(ColLayer);
      }

      T1=T2;
      T2=0.;
      if(sqrt(fC)>=fR[iLayer]){  // vertex in ocean
        ColInt = this->fIntLayer(iLayer+1,T1,T2);
        ColLayer.iLayer=iLayer+1;
        ColLayer.ColumnDepth=ColInt;
        fColumnDepthLayer.push_back(ColLayer);

      }
      else { // vertex in rock
        ColInt = this->fIntLayer(iLayer,T1,T2);
        fColumnDepthLayer[fColumnDepthLayer.size()-1].ColumnDepth -= ColInt;
      }

  }

  for(unsigned i=0; i<fColumnDepthLayer.size();i++){
    fColumnDepth+=fColumnDepthLayer[i].ColumnDepth;
  }
  fColumnDepth*=1E9;    // in mwe

  return;
}
//_____________________________________________________________________
double GPEarth::fGetXsecTot(double Energy, int neutrino_pdg, map<int,TGraph> XSecTot){
  double XSec=0.;

  map<int,TGraph>::iterator iter = XSecTot.begin();
  for ( ; iter != XSecTot.end(); ++iter) {
    if(neutrino_pdg==iter->first){
       TGraph GrSecTot = iter->second;
       XSec=GrSecTot.Eval(Energy);
       break;
    }
  }
  return XSec;
}
//______________________________________________________________________
double GPEarth::fIntLayer(int iLayer, double T1, double T2){
  double ColIntLayer=0.;

  ColIntLayer +=(fEarthCoeff[iLayer][0]*this->fINTEG0(T2)+
                fEarthCoeff[iLayer][1]/fREarth*this->fINTEG1(T2)+
                fEarthCoeff[iLayer][2]/pow(fREarth,2)*this->fINTEG2(T2)+
                fEarthCoeff[iLayer][3]/pow(fREarth,3)*this->fINTEG3(T2))*
                sqrt(fA);

  ColIntLayer -=(fEarthCoeff[iLayer][0]*this->fINTEG0(T1)+
                fEarthCoeff[iLayer][1]/fREarth*this->fINTEG1(T1)+
                fEarthCoeff[iLayer][2]/pow(fREarth,2)*this->fINTEG2(T1)+
                fEarthCoeff[iLayer][3]/pow(fREarth,3)*this->fINTEG3(T1))*
                sqrt(fA);

  return ColIntLayer;
}
//_______________________________________________________________________
void GPEarth::fExtremes(void){

  fVecExtremPoint.clear();

  ExtremPoint epoint;

  this->fInterCir();

  bool check = true;
  for(unsigned i=1; i<fVecInterPoint.size(); i++){
    if(fVecInterPoint[i-1].iLayer==fVecInterPoint[i].iLayer)check=false;

    epoint.T1=fVecInterPoint[i-1].T;
    epoint.T2=fVecInterPoint[i].T;

    if(check)epoint.iLayer=fVecInterPoint[i-1].iLayer;
    else epoint.iLayer=fVecInterPoint[i].iLayer;

    fVecExtremPoint.push_back(epoint);
  }

  return;
}
//_______________________________________________________________________
void GPEarth::fInterCir(void)
{
// intersecations of straight X=X0+LT, Y=Y0+MT, Z=Z0+NT
// with the spheres X^2+Y^2+Z^2+ALPHA X+BETA Y+GAMMA Z + DELTA=0

  double Delta;
  InterPoint ipoint;

  fVecInterPoint.clear();

  for(int iLayer=0; iLayer<NLAYERS; iLayer++){

    Delta = fB*fB-4*fA*(fC-TMath::Power(fR[iLayer],2));

    if(Delta>=0){
      Delta = sqrt(Delta);
      ipoint.iLayer=iLayer;
      ipoint.T=(-fB-Delta)/(2.*fA);
      fVecInterPoint.push_back(ipoint);
      ipoint.T=(-fB+Delta)/(2.*fA);
      fVecInterPoint.push_back(ipoint);
    }
  }

  sort(fVecInterPoint.begin(),fVecInterPoint.end());

  return;
}
//_______________________________________________________________________
double GPEarth::fINTEG0(double X)
{
//  Primitive of 1 evaluated at T=X
  return X;
}
//_______________________________________________________________________
double GPEarth::fINTEG1(double X)
{
//  Primitive of SQRT(AT^2+BT+C) evaluated at T=X

  double INT1A,INT1B,INT1C,INT1;
  double ARG;

  if(fB*fB==4.*fA*fC){
    INT1= 1./(2.*sqrt(fA))*pow((fA*X+fB/2.),2);
  }
  else {
    INT1A=(2.*fA*X+fB)*sqrt(fA*X*X+fB*X+fC)/(4.*fA);
    INT1B=(4.*fA*fC-fB*fB)/(8.*fA);
    ARG=2.*sqrt(fA)*sqrt(fA*X*X+fB*X+fC)+2.*fA*X+fB;
    INT1C=1./sqrt(fA)*log(ARG);
    INT1=INT1A+INT1B*INT1C;
  }

  return INT1;
}
//_______________________________________________________________________
double GPEarth::fINTEG2(double X)
{
//  Primitive of AT^2+BT+C evaluated at T=X
  return fA*(pow(X,3)/3.)+fB*(X*X/2.)+fC*X;
}
//_______________________________________________________________________
double GPEarth::fINTEG3(double X)
{
//  Primitive of (AT^2+BT+C)^3/2 evaluated at T=X
  double INT3A,INT3B,INT3C,INT3;

  if(fB*fB==4.*fA*fC){
    INT3=1./(4.*sqrt(fA))*pow((fA*X+fB/2.),2);
  }
  else {
    INT3A=(2.*fA*X+fB)*pow((fA*X*X+fB*X+fC),(3./2.))/(8.*fA);
    INT3B=3.*(4.*fA*fC-fB*fB)/(16.*fA);
    INT3C=this->fINTEG1(X);
    INT3=INT3A+INT3B*INT3C;
  }

  return INT3;
}