// BetaFunction.cc
// Contact person:  Logan Sibley <lsibley@ualberta.ca>
// See BetaFunction.hh for more details
//———————————————————————//

#include <RAT/BetaFunction.hh>
#include <RAT/FermiFunction.hh>

#include <RAT/Log.hh>
#include <TMath.h>
#include <CLHEP/Units/SystemOfUnits.h>
#include <Randomize.hh>
using namespace CLHEP;

namespace RAT {

//Static variables
  TDirectory *BetaFunction::fDirectory;

  BetaFunction::BetaFunction()
  {
    Reset();
  }

  BetaFunction::BetaFunction(const std::string Name)
  {
    Reset();
    SetTargetName(Name);
  }

  BetaFunction::BetaFunction(const std::string Name, double A, double Z, int reac, double tau)
  {
    Reset();
    SetTarget(Name, A, Z, reac, tau);
  }

  BetaFunction::~BetaFunction()
  {
    Reset();
  }

  void BetaFunction::Reset()
  {
    SetTargetName("\0");
    SetTargetCharge(0.);
    SetTargetMass(0.);
    SetTargetDecayType(DecayBeta);       // beta decay by default
    SetNeutrinoMass(0.);
    SetTargetLifeTime(0.);

    SetIsVerbose(false);
    SetIsCulmulative(true);
    fNumberOfBranches = 0;

    if(fDirectory == NULL) fDirectory = new TDirectory("fDir","BetaHistos");
    fDirectory->cd();

    // See if the $RATDecayDataDir environment variable exists.
    char* path = getenv("RATDecayDataDir");
    std::string dirName;
    if ( path == 0 ) {
	// Environment variable not found... so just use a name
	// relative to the current directory.
	dirName = "./data";
    } else {
	dirName = std::string(path);
    }
    std::string fileName = dirName + "/beta_decays.dat";
    SetInputFileName(fileName);
  }

  void BetaFunction::Show()
  {
    printf("Target Name       : %s \n", GetTargetName().c_str());
    printf("Target Mass       : %f \n", GetTargetMass());
    printf("Target Charge     : %f \n", GetTargetCharge());
    printf("Decay  Type       : %d \n", GetTargetDecayType());
    printf("Target Lifetime   : %f \n", GetTargetLifeTime());
    printf("Neutrino Mass     : %f \n", GetNeutrinoMass());
    printf("Number of Branches: %d \n", GetNuberOfBranches());
    printf("Branch Prob. \t Spin \t Endpoint \t NGammas \t Energies \n");

    int nBr = GetNuberOfBranches();
    if (nBr > 0) {

      for (int i = 0; i < nBr; i++) {
	printf("%f \t", GetBranch(i));
	printf("%d \t", GetSpin(i));
	printf("%f \t", GetEndPoint(i));
	int nP = GetNParticles(i) - 1;
	printf("%d \t", nP);
	if (nP > 0) {
	  for (int j = 1; j <= nP; j++) {
	    int id = GetParticleID(i, j);
	    if (id == DecayGamma) {
	      printf("%f \t", GetEnergy(i, j));
	    }
	  }
	  printf(" \n");
	} else {
	  printf(" \n");
	}
      }
      printf(" \n");
    }
  }

  void BetaFunction::SetTarget(const std::string Name, double Z, double A,
                               int reaction, double tau)
  {
    SetTargetName(Name);
    SetTargetMass(A);
    SetTargetCharge(Z);
    SetTargetDecayType(reaction);
    SetNeutrinoMass(0.);
    SetTargetLifeTime(tau);
  }

  void BetaFunction::SetBranches(double Branch, int Spin, double EndPoint, float Corr[])
  {
    int iZ = static_cast<int>(GetTargetCharge());
    if (iZ == 0) {
      printf("No charge has been set. \n");
      return;
    }

    int nBr = GetNuberOfBranches();

    fParticleBranch[nBr] = Branch;
    fParticleSpin[nBr] = Spin;
    fParticleEndPoint[nBr] = EndPoint;
    fParticleEnergy[nBr][0] = EndPoint;

    fNumberOfParticles[nBr] = 1;
    if (GetTargetDecayType() == DecayAlpha) {
      fParticleID[nBr][0] = DecayAlpha;
    } else if (GetTargetDecayType() == DecayGamma) {
      fParticleID[nBr][0] = DecayGamma;//For gamma decay (including electron
      fNumberOfParticles[nBr] = 0;     //capture), a beta spectrum is not
      fNumberOfBranches++;             //needed, so skip it! (Trying with a
      return;                          //zero endpoint breaks the code)
    } else {
      if (iZ > 0){                       // Beta decay
        fParticleID[nBr][0] = DecayBeta;
        SetTargetDecayType(DecayBeta);
      } else if (iZ < 0) {               // Positron emission
        fParticleID[nBr][0] = DecayEC;
        SetTargetDecayType(DecayEC);
      }
    }
    SetNorm(nBr, Corr);
    fNumberOfBranches++;
  }

  void BetaFunction::RemoveBranch(int iBranch)
  {
    fParticleBranch.erase(iBranch);
    fParticleSpin.erase(iBranch);
    fParticleEndPoint.erase(iBranch);
    fParticleID.erase(iBranch);
    fParticleEnergy.erase(iBranch);
    fNumberOfParticles.erase(iBranch);
//    fDecayNorm.erase(iBranch);
    fNumberOfBranches--;
  }

  void BetaFunction::SetGammas(int iBranch, int pid, double energy)
  {
    int iParticle = GetNParticles(iBranch);
    fParticleID[iBranch][iParticle] = pid;
    fParticleEnergy[iBranch][iParticle] = energy;
    fNumberOfParticles[iBranch]++;
  }

  void BetaFunction::SetGammas(double energy)
  {
    int iBranch = GetNuberOfBranches() - 1;
    int iParticle = GetNParticles(iBranch);

    fParticleID[iBranch][iParticle] = DecayGamma;
    fParticleEnergy[iBranch][iParticle] = energy;
    fNumberOfParticles[iBranch]++;
  }

  void BetaFunction::SetParticleID(int iBranch, int n, int id)
  {
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(2);
      return;
    }
    fParticleID[iBranch][n] = id;
  }

  void BetaFunction::SetTargetMass(double A)
  {

    fTargetMass = 0.;
    double Z = GetTargetCharge();
    if ((A < 1) && (Z > 0)) {
      fTargetMass =
	1.82 + 1.90 * Z + 0.01271 * pow(Z, 2) - 0.00006 * pow(Z, 3);
    } else {
      fTargetMass = A;
    }
  }

  void BetaFunction::SetNorm(int iBranch, float Corr[])
  {
    std::string s;  // If an isotope both beta decays and emits a positron
    if (GetTargetDecayType() == DecayBeta) s = " beta ";
    else if (GetTargetDecayType() == DecayEC) s = " positron ";

    std::ostringstream os;
    os << iBranch;
    std::string hname = GetTargetName()+s+" Decay Branch "+os.str();
    TH1D *hDecaySpect = NULL;
    double nSample = 1000.;
    if (GetTargetDecayType() == DecayAlpha) {
      hname = hname+" alpha";
      hDecaySpect = (TH1D*)fDirectory->Get(hname.c_str());
      if(!hDecaySpect) {
	hDecaySpect = new TH1D(hname.c_str(),hname.c_str(),1,0.,1.);
        hDecaySpect->SetDirectory(fDirectory);
        hDecaySpect->Fill(0.5,1.);
      }
    }
    else {
      hDecaySpect = (TH1D*)fDirectory->Get(hname.c_str());
      if(!hDecaySpect) {
        double eMin = ElectronMass;
        double eMax = eMin + GetEndPoint(iBranch);
        double dE = (eMax - eMin) / nSample;
        hDecaySpect = new TH1D(hname.c_str(),hname.c_str(),
                               static_cast<int>(nSample)+1,eMin,eMax);
        hDecaySpect->SetDirectory(fDirectory);
        for (double energy = eMax-dE; energy >= eMin; energy = energy - dE){
          if (energy < eMin + dE) energy = eMin;
	  double F = GetValue(energy, iBranch, Corr);
// When the energy becomes low, F becomes zero, so set those bins equal to the
// last non-zero bin (thus why the histogram is cycled from emax to emin).
          if (F <= 0.) F = hDecaySpect->GetBinContent(static_cast<int>(nSample)+1);
          hDecaySpect->Fill(energy,F);
          nSample = nSample-1.;
        }
      }
    }

  }

//The old way of setting the norm
/*  void BetaFunction::SetNorm(int iBranch)
  {
    double Norm = 0.;
    if (GetTargetDecayType() == DecayAlpha) {
      Norm = 1.;
    } else {
      double eMin = ElectronMass;
      double eMax = eMin + GetEndpoint(iBranch);
      double dE = (eMax - eMin) / 1000.;
      for (double energy = eMin; energy < eMax; energy = energy + dE) {
        double F = GetValue(energy, iBranch);
        if (F > Norm) Norm = F;
      }
    }
    if (Norm <= 0.) Norm = 1.;
    fDecayNorm = Norm * 1.11;
  }*/

  int BetaFunction::GetNParticles(int iBranch)
  {
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(1);
      return iBAD;
    }
    return fNumberOfParticles[iBranch];
  }

  int BetaFunction::GetParticleID(int iBranch, int n)
  {
    int id;
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(2);
      return iBAD;
    }
    if (n > fNumberOfParticles[iBranch]) {
      id = NullParticle;
    } else {
      id = fParticleID[iBranch][n];
    }
    return id;
  }

  int BetaFunction::GetSpin(int iBranch)
  {
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(0);
      return iBAD;
    }
    return fParticleSpin[iBranch];
  }

  double BetaFunction::GetBranch(int iBranch)
  {
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(0);
      return iBAD;
    }
    return fParticleBranch[iBranch];
  }

  double BetaFunction::GetEndPoint(int iBranch)
  {
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(0);
      return iBAD;
    }
    return fParticleEndPoint[iBranch];
  }

  double BetaFunction::GetEnergy(int iBranch, int n)
  {
    double energy = 0.;
    if ((iBranch > fNumberOfBranches) || (iBranch < 0)) {
      ErrorLog(2);
      return iBAD;
    }

    if (n > fNumberOfParticles[iBranch]) {
      return energy;
    } else {
      energy = fParticleEnergy[iBranch][n];
    }

    return energy;
  }

  double BetaFunction::GetValue(double energy, int iBranch, float Corr[])
  {
    double Value = 0.;
    int iBeta = GetParticleID(iBranch, 0);

    if ((iBeta != DecayEC) && (iBeta != DecayBeta))
      return Value;

    int iSpin = GetSpin(iBranch);
    double Z = GetTargetCharge();
    double A = GetTargetMass();
    double nu = GetNeutrinoMass();
    double ePoint = GetEndPoint(iBranch);
    double W0 = 1. + ePoint / ElectronMass;
    double W = energy / ElectronMass;
    bool scr = GetScreening();
    Value = Nucl_Beta(iBeta, Z, A, W, W0, iSpin, nu, Corr, scr);
    return Value;
  }

  void BetaFunction::GenerateEvent()
  {
    for (int i = 0; i < PartMax; i++) {
      fEn[i] = 0.;
      fIDn[i] = 0;
    }

    int iBr = GetNuberOfBranches();
    int itag = 0;
    double prob = 0.;
    double r = GetRandomNumber();

    for (int i = 0; i < iBr; i++) {
      if (fIsCulmulative) {
	prob = GetBranch(i);
      } else {
	prob = prob + GetBranch(i);
      }

      if (r < prob) {
	itag = i;
	break;
      }
    }

    int nP = GetNParticles(itag);
    for (int n = 0; n < nP; n++) {
      int id = GetParticleID(itag, n);
      if ((id == DecayBeta) || (id == DecayEC)) { // Get a beta energy!
        std::string s;
        if (id == DecayBeta) s = " beta ";
        else if (id == DecayEC) s = " positron ";
        std::ostringstream os;
        os << itag;
        std::string hname = GetTargetName()+s+" Decay Branch "+os.str();
        TH1D *hDecaySpect = (TH1D*)fDirectory->Get(hname.c_str());

// If the branch is not set for some reason, set it and re-generate the event.
// This should never happen... but just in case.
        if(!hDecaySpect){
          warn << "Branch was not previously set! Something may be wrong.\n";
          float corr[3] = {0., 0., 0.};
          SetNorm(itag, corr);
          GenerateEvent();
        } else {
          // double energy = hDecaySpect->GetRandom();
          /* ROOT's GetRandom() function is independent of RAT's global random
             generator. It is replaced by a function GetRandomEnergy(), which is
             based on ROOT's function with needed modifications and uses the
             G4UniformRand random generator instead of TRandom.
         */
          double energy = GetRandomEnergy(hDecaySpect);

//        The old method of sampling the Fermi Function, if you wish to use it
//        bool ithrow = true;
//        while(ithrow) {
//          double R = GetNorm(itag) * GetRandomNumber();
//          double energy = ElectronMass + GetEndPoint(itag) *GetRandomNumber();
//          double F = GetValue(energy, itag);
//          W = energy / ElectronMass;
//          if (R < F)
//            ithrow = false;
//        }

          fIDn[n] = id;
          fEn[n] = energy - ElectronMass; // Because the hist is total energy
        }
      } else {
        if ((id == DecayGamma) && (GetEnergy(itag,n) < 0.)) {
          fIDn[n] = DecayBeta;  // Energies of any conversion electrons
          fEn[n] = fabs(GetEnergy(itag,n));
        } else {
          fIDn[n] = id;         // Get energies for alpha decay or gamma decay
          fEn[n] = GetEnergy(itag,n);
        }
      }
    }
    SetEventTime();
    fNGenerated = nP;
  }

  Double_t BetaFunction::GetRandomEnergy(TH1D* h)
  {
  // This function is created to replace the h1->GetRandom() function of ROOT
  // The reason is to use the same random generator as the rest of the RAT simulation.
    Int_t nbinsx = h->GetNbinsX();
    Double_t integral;
    Double_t  fEntries  = h->GetEntries();
    Double_t* fIntegral = h->GetIntegral();

    if (fIntegral){
      if (fIntegral[nbinsx+1] != fEntries)
        integral = ((TH1*)this)->ComputeIntegral();
    } else {
      integral = ((TH1*)this)->ComputeIntegral();
      if (integral == 0 || fIntegral == 0) return 0;
    }

    Double_t r1 = G4UniformRand();
    Int_t ibin = TMath::BinarySearch(nbinsx,fIntegral,r1);
    Double_t en = h->GetBinLowEdge(ibin+1);
    if (r1 > fIntegral[ibin])
      en += h->GetBinWidth(ibin+1)*(r1-fIntegral[ibin])/(fIntegral[ibin+1]-fIntegral[ibin]);

    return en;
 }

  void BetaFunction::SetEventTime()
  {
    double tau = GetTargetLifeTime();
    double r = GetRandomNumber();
    fEventTime = -tau * log(r) / log(2.);
  }

  int BetaFunction::GetEventID(int n)
  {
    if (n > fNGenerated) {
      ErrorLog(3);
      return iBAD;
    }
    return fIDn[n];
  }

  double BetaFunction::GetEventEnergy(int n)
  {
    if (n > fNGenerated) {
      ErrorLog(3);
      return iBAD;
    }
    return fEn[n];
  }

  double BetaFunction::GetEventTotE()
  {

    int nG = GetNGenerated();
    double eTot = 0.;

    for (int i = 0; i < nG; i++) {
      eTot = eTot + fabs(GetEventEnergy(i)); // fabs ensures we don't subtract
    }                                        //  conversion electron energies
    return eTot;
  }

  double BetaFunction::GetRandomNumber(double rmin, double rmax)
  {
    double rnd = G4UniformRand(); // random number from 0 to 1.
    double value = rmin + (rmax - rmin) * rnd;
    return value;
  }

  bool BetaFunction::ReadInputFile(const std::string dName)
  {
    bool iFound = ReadInputFile(dName, -1, false);
    return iFound;
  }

  bool BetaFunction::ReadInputFile(const std::string dName, int iType, bool screen)
  {
    bool iFound = ReadInputFile(dName, -1, -1, iType, screen);
    return iFound;
  }

  bool BetaFunction::ReadInputFile(const std::string dName, int iZ, int iA, int iType, bool screen)
  {

    const int nReadGamma = 6;

    std::string eProbe = "END";
    std::string aProbe = "ALPHA:";
    std::string gProbe = "GAMMA:";
    std::string bProbe = "DECAY:";
    char dummy[80];
    std::string dProbe;
    char tName[80];

    if (iType == DecayAlpha) {        // alpha decay
      dProbe = aProbe;
    } else if(iType == DecayGamma) {  // gamma decay
      dProbe = gProbe;
    } else {
      dProbe = bProbe;                // beta decay/positron emission
    }

    bool iFound = false;

    int iSpin, nP;
    int Z, A;
    float tau;
    float iBr, W0, eP[nReadGamma], aC[3];

    SetIsCulmulative(true);

    FILE* inputFile = fopen(GetInputFileName().c_str(), "r");

    if (!inputFile) {
      printf("Error opening file : %s \n", GetInputFileName().c_str());
      return iFound;
    }

    bool iRead = true;

    while ((iRead) && (!iFound)) {
      iRead = (fscanf(inputFile, "%s", dummy) > 0);
      std::string iString(dummy);
      if (iString == dProbe) {
	fscanf(inputFile, "%s", tName);
	fscanf(inputFile, "%d %d %f", &Z, &A, &tau);
	std::string iString(tName);
	iFound = (dName == iString);
	iFound = ((iFound) || ((iA == A) && (iZ == Z)));
	if (iFound) {
	  if (fIsVerbose)
	    printf("Reading %s \n \n", tName);
	  SetTargetMass(static_cast<double>(A));
	  SetTargetCharge(static_cast<double>(Z));
	  SetTargetLifeTime(static_cast<double>(tau));
	  SetTargetDecayType(iType);
          SetScreening(screen);

	  bool iScan = true;
	  while (iScan) {
	    fscanf(inputFile, "%f %d %f %d", &iBr, &iSpin, &W0, &nP);
	    for (int j = 0; j < nReadGamma; j++) {
	      fscanf(inputFile, "%f", &eP[j]);
	    }
// Use the shape correction factors when setting branches - Branches must be
// set before gammas!
	    fscanf(inputFile, "%f %f %f", &aC[0], &aC[1], &aC[2]);
	    SetBranches(static_cast<double>(iBr),iSpin,static_cast<double>(W0),aC);
	    for (int j = 0; j < nReadGamma; j++) {
	      if (eP[j] != 0.0) SetGammas(static_cast<double>(eP[j]));
	    }
	    if (iBr >= 1.)
	      iScan = false;
	  }
// Set the parent A and Z, depending on decay type and daughter A and Z
          double parentMass = GetTargetMass();
          double parentCharge = fabs(GetTargetCharge()) -
                                  static_cast<double>(GetTargetDecayType());
          if(GetTargetDecayType() == DecayAlpha){
            parentMass += 2.*DecayAlpha;
            parentCharge += 2.*DecayAlpha;
          } else if((GetTargetDecayType() == DecayGamma) &&
                    (GetTargetCharge() < 0)){
            parentCharge += 1.;    // An electron capture
          }
          SetParentMass(parentMass);
          SetParentCharge(parentCharge);

	} else if (iString == eProbe) {
	  iRead = false;
	} else {
	}
      }
    }

    fclose(inputFile);

    if (iFound) {
      if (fIsVerbose)
	Show();
    } else {
      printf("No such element found: %s .\n", dName.c_str());
    }
    return iFound;
  }

  void BetaFunction::ErrorLog(int iFlag)
  {
    printf("Error detected: %d \n", iFlag);
  }

} // namespace RAT