/*
 * IBDCrossSec.cc
 *
 *  Created on: Aug 23, 2015
 *      Author: nbarros
 */

#include <RAT/IBDCrossSec.hh>
#include <RAT/Log.hh>

// -- CLHEP includes
#include <CLHEP/Units/PhysicalConstants.h>

//-- STL includes
#include <sstream>
#include <cstring>
#include <cstdlib>
#include <cmath>

// -- ROOT includes
#include <TFile.h>
#include <TH2F.h>

using CLHEP::MeV;
using CLHEP::proton_mass_c2;
using CLHEP::neutron_mass_c2;
using CLHEP::electron_mass_c2;
using CLHEP::pi;
using CLHEP::hbarc;

namespace RAT {

  const map<IBDCrossSec::Strategy,string> IBDCrossSec::_str_map =  IBDCrossSec::init_map();

  //
  // -- Define a series of invariant constants that are used in the calculations below
  // -- These are independent of the strategies used
  //
  const double IBDCrossSec::fPhase_f = 1.00;
  const double IBDCrossSec::fPhase_f2 = 3.706;
  const double IBDCrossSec::fPhase_g = 1.27641; // +/- 0.00078: https://doi.org/10.1103/PhysRevLett.122.242501

  const double IBDCrossSec::fNeutronLifetime = 879.4  / 6.58e-22; // +/- 0.6 / 6.58e-22

  const double IBDCrossSec::fDelta = neutron_mass_c2 - proton_mass_c2;

  const double IBDCrossSec::fGFermi = 1.16639e-11 / MeV / MeV; // +/- 0.0510e-11  MeV^-2
  const double IBDCrossSec::fCosThetaC = (0.9741+0.9756)/2; // +/- 0.00021



  IBDCrossSec::IBDCrossSec() : fInputFile(0), fXsTable(0) {

    // The default strategy is the on-the-fly calculation with sigma0 calculated based on
    // the neutron lifetime
    // NB: Before Oct 2019 the calculation of Sigma0 used the Sigma0_Original() method
    SetStrategy(VogelNeutron);
#ifdef RAT_XS_DEBUG
    detail << "IBDCrossSec::IBDCrossSec : In constructor instantiating with strategy " << fStrategy << newline;
#endif

  }

  IBDCrossSec::~IBDCrossSec() {
    if (fInputFile) fInputFile->Close();
    delete fInputFile;
    if (fXsTable) delete fXsTable;
    fXsTable = NULL;
  }


  IBDCrossSec::Strategy IBDCrossSec::convert_strategy(const short int strat)
  {
    if (strat <= Unknown || strat >= NoStrategy)
    {
      Log::Die("Unknown strategy. Be sure that a valid value is selected.");
    }
    return static_cast<Strategy>(strat);

  }

  void IBDCrossSec::SetStrategy(const short int strat)
  {

    SetStrategy(convert_strategy(strat));
  }

  void IBDCrossSec::SetStrategy(Strategy strat) {
#ifdef RAT_XS_DEBUG
    detail << "IBDCrossSec::SetStrategy : Setting strategy to " << static_cast<int>(strat) << newline;
#endif
    fStrategy = static_cast<short int>(strat);
    switch (strat) {
      case VogelTable:
      case VogelNeutronTable:
        LoadTablesDB();
        break;
      case Vogel:
      case VogelNeutron:
        // -- nothing to be done
        break;
      default:
        std::ostringstream msg("");
        msg << "Unknown strategy [" << static_cast<int>(strat) << "]. Be sure that a valid value is selected.";
        Log::Die(msg.str(),1);
    }
  }



  double IBDCrossSec::dSigmadTheta(const double Enu, const double cosThetaLab) const {
    // This gives room to implement different calculation strategies to improve performance
    // and/or precision

    switch(static_cast<Strategy>(fStrategy)) {
      // Strategy 0 is the default
      case Vogel:
      case VogelNeutron:
        return dSigmadThetaVogel(Enu,cosThetaLab);
        break;
      case VogelTable:
      case VogelNeutronTable:
        return dSigmadThetaVogelTable(Enu,cosThetaLab);
        break;
      default:
        warn << "IBDCrossSec::dSigmadTheta : Invalid Cross section strategy set." << newline;
        return 0.0;
    }
  }


  double IBDCrossSec::Sigma0() const
  {
    static int n_warns = 0;
    static bool warns_done = false;
    switch(fStrategy)
    {
      case 0: // Vogel
        return Sigma0_Original();
        break;
      case 2: // VogelNeutron
        return Sigma0_Neutron();
        break;
      case 1: // VogelTable
      case 3: // VogelNeutronTable
        if (!warns_done)
        {
          warn << "IBDCrossSec::Sigma0 : Sigma0 already included in table for strategy 2. Returning 1.0 as scaling factor." << newline;
          n_warns++;
          if (n_warns == 10) warns_done = true;
        }
        return 1.0;
        break;
      default:
        std::stringstream msg("");
        msg << "IBDCrossSec::Sigma0 : Invalid strategy (" << fStrategy << " ) set.";
        Log::Die(msg.str());
    }
  }

  double IBDCrossSec::Sigma0_Neutron() const
  {

    // Cross section expressed in terms of neutron lifetime.
    // This reduces the result by approx 0.5% but the
    // associated uncertainty is an order of magnitude lower.
    // Values updated with 2019 PDG values unless stated otherwise.
    // https://doi.org/10.1103/PhysRevD.98.030001
    // Default calculation as of October 2019

    //Phase space corrections constant
    const double PhaseFactor = 1.727641; // +/- 0.00009

    // overall scale
    const double Sigma0 = 2*pi*pi /
            (electron_mass_c2*electron_mass_c2*electron_mass_c2*electron_mass_c2*electron_mass_c2*PhaseFactor*fNeutronLifetime) /
            (fPhase_f*fPhase_f+3*fPhase_g*fPhase_g);

    return Sigma0;
  }

  double IBDCrossSec::Sigma0_Original() const
  {
    // Cross section constants.  Some for overall scale are just
    // to allow absolute comparison to published article
    // (Vogel & Beacom, 1999)
    // Physical Review D, 60(5), 053003.
    // http://doi.org/10.1103/PhysRevD.60.053003
    // This was the default calculation before October 2019

    // Radiative correction constant
    const double RadCor = 0.024; // +/- 0.00022: arXiv:1812.03352

    // overall scale
    const double Sigma0 = fGFermi*fGFermi*fCosThetaC*fCosThetaC/pi*(1+RadCor);

    return Sigma0;

  }

  double IBDCrossSec::dSigmadThetaVogel(const double Enu, const double cosThetaLab) const
  {
    double sigma = 0.0;

    // check for threshold
    static const double EminBeta = (
           (neutron_mass_c2+electron_mass_c2)*
           (neutron_mass_c2+electron_mass_c2)
           -proton_mass_c2*proton_mass_c2
           )  /2.0/proton_mass_c2;

    // If energy is below threshold just return 0
    if(Enu<EminBeta) return 0.0;


    /// -- Calculation of the cross section proper


    // -- order 0 terms
    double E0 = Enu - fDelta;
    if(E0<electron_mass_c2) E0=electron_mass_c2;
    double p0 = sqrt(E0*E0-electron_mass_c2*electron_mass_c2);
    double v0 = p0/E0;

    // -- order 1 terms
    const double Ysquared = (fDelta*fDelta-electron_mass_c2*electron_mass_c2)/2;
    double E1 = E0*(1.0-Enu/proton_mass_c2*(1-v0*cosThetaLab))-Ysquared/proton_mass_c2;
    if(E1<electron_mass_c2) E1=electron_mass_c2;
    double p1 = sqrt(E1*E1-electron_mass_c2*electron_mass_c2);
    double v1 = p1/E1;

    double Gamma =
      2*(fPhase_f+fPhase_f2)*fPhase_g*((2*E0+fDelta)*(1-v0*cosThetaLab)-
      electron_mass_c2*electron_mass_c2/E0)
      +(fPhase_f*fPhase_f+fPhase_g*fPhase_g)*(fDelta*(1+v0*cosThetaLab)+electron_mass_c2*electron_mass_c2/E0)
      +(fPhase_f*fPhase_f+3*fPhase_g*fPhase_g)*((E0+fDelta)*(1-cosThetaLab/v0)-fDelta)
      +(fPhase_f*fPhase_f-fPhase_g*fPhase_g)*((E0+fDelta)*(1-cosThetaLab/v0)-fDelta)*v0*cosThetaLab;

    sigma =
      ((fPhase_f*fPhase_f+3*fPhase_g*fPhase_g)+(fPhase_f*fPhase_f-fPhase_g*fPhase_g)*v1*cosThetaLab)*E1*p1
      -Gamma/proton_mass_c2*E0*p0;
    sigma *= Sigma0()/2;
    sigma *= hbarc * hbarc; // Convert from MeV^-2 to mm^2 (native units for GEANT4)

    return sigma;
  }

  double IBDCrossSec::dSigmadThetaVogelTable(const double Enu, const double cosThetaLab) const
  {
    return fXsTable->Interpolate(Enu,cosThetaLab);
  }

  void IBDCrossSec::LoadTablesDB() {
    bool has_error = false;
    std::ostringstream msg("");
    Strategy s = static_cast<Strategy>(fStrategy);
    info << "IBDCrossSec::LoadTablesDB : Loading pre-computed cross section table. Strategy : " << _str_map.at(s) << newline;
    // Check if the strategy is of table type
    if (s != VogelTable && s!=VogelNeutronTable) {
      warn << "IBDCrossSec::LoadTablesDB : Strategy " << _str_map.at(s) << " does not require loading a table. Aborting." << newline;
      return;
    }

    // See if the $IBDDATADIR environment variable exists.
    std::stringstream dirName;
    char* ibdDataDir = getenv("IBDDATADIR");
    if( ibdDataDir == NULL )
    {
      //No specific directory set, so use the rat default
      dirName << getenv("RATROOT") << "/data";
    }
    else
    {
      dirName << ibdDataDir;
    }
    dirName << "/ibd_cross_section_tables.root";

    fInputFile = TFile::Open(dirName.str().c_str());
    if (!fInputFile) {
      Log::Die("IBDCrossSec::LoadTablesDB : Input file not found " + dirName.str(),1);
    }

    switch (s) {
      case VogelTable:
        fXsTable = (TH2F*)fInputFile->Get("ibd_dSigmadCosTheta_strat1");
        // to decouple it from the open file directory
        fXsTable->SetDirectory(0);
        break;
      case VogelNeutronTable:
        fXsTable = (TH2F*)fInputFile->Get("ibd_dSigmadCosTheta_strat3");
        // to decouple it from the open file directory
        fXsTable->SetDirectory(0);
        break;

      default:
        has_error = true;
        msg << "IBDCrossSec::LoadTablesDB : Unknown strategy.";
    }
    // close the file if it exists
    if (fInputFile)
    {
      fInputFile->Close();
      delete fInputFile;
    }
    fInputFile = NULL;

    if (has_error)
    {
      Log::Die(msg.str(),1);
    }

  }

  double IBDCrossSec::Sigma(const double Enu) const
  {
    // This method calculates an integral over all angles
    // First do the most basic check
    if (Enu == 0) return 0.0;
    if (Enu < 0) {
      std::stringstream ss;
      ss << "IBDCrossSec::Sigma : Invalid neutrino Energy ( Enu = "
          << Enu << " ).";
      RAT::Log::Die(ss.str(),1);
    }

    double   sigma = 0.0;
    double cos_theta;
    const int num_steps = 200;
    const double dstep = 2./(double)num_steps;// changing between 200 and 2000 bins changes result by 0.05%
    for (int i = 0; i < num_steps; ++i ) {
      cos_theta = -1.0 + i*dstep;
      sigma += dSigmadTheta(Enu,cos_theta)*dstep;
    }
    return sigma;

  }
} /* namespace RAT */