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

// RAT::DecayChain_Gen
// 02Nov2009 Aleksandra Bialek - modifications
//    See comments in RATDecayChain_Gen.hh

#include <CLHEP/Units/PhysicalConstants.h>
using CLHEP::pi; using CLHEP::twopi; using CLHEP::halfpi;

#include <CLHEP/Vector/LorentzVector.h>

#include <RAT/DecayChain.hh>
#include <RAT/FermiFunction.hh>
#include <RAT/Log.hh>

#include <Randomize.hh>

#include <cstring>
#include <cmath>
#include <sstream>
using namespace std;

#undef DEBUG

namespace RAT{

DecayChain::DecayChain()
{
  Reset();
}

DecayChain::DecayChain(const std::string Name)
{
#ifdef DEBUG
  debug << "In RAT:: DecayChain::DecayChain(), with Name="
	    << Name << "\n";
#endif

  Reset();
  SetChainName(Name);
}

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

void DecayChain::Reset()
{
#ifdef DEBUG
  debug << "In RAT:: DecayChain::Reset()\n";
#endif
  SetInMiddleChain();
  SetScreenCorr();
  SetIsVerbose();
  SetIsEquilibrium();
  SelectParent();
  IncludeDaughter(true);
  fInMiddleChain=false;
  fIsChainElement = false;
  fAlphaDecayStart=false;
  fGammaDecayStart=false;
  fSetScreenCorr=false;
  SetChainName("");
  fNumberOfElements = 0;
  fTstart = 0.;
  SetNGenerated(0);

  // See if the $RATDecayDataDir environment variable exists.
  stringstream dirName;
  char* ratDecayDataDir = getenv("RATDecayDataDir");
  if( ratDecayDataDir == NULL )
    {
      //No specific directory set, so use the rat default
      dirName << getenv("RATROOT") << "/data";
    }
  else
    {
      dirName << ratDecayDataDir;
    }
  dirName << "/beta_decays.dat";
  SetInputFileName( dirName.str() );
}

void DecayChain::Show()
{
  printf("Decay Chain Name: %s .\n", GetChainName().c_str());
  int nE = GetNuberOfElements();
  if (nE > 0) {
     printf("Elements : %d \n", nE);
     printf ("Name \t Branch \t Decay Type \t Half-Life \t Weight \t Prob \t Level \n");
     for (int i = 0; i < nE; i++) {
	      printf("%s \t %d \t %d \t %g \t %g \t %g \t %d \n", GetElementName(i).c_str(),
	      GetElementBranch(i), GetElementDecay(i), GetLifetime(i),
	      GetElementWt(i), GetElementProb(i), CheckChainLevel(i));
     }
  }
  else {
     printf("No elements defined. \n");
  }
  printf(" \n");
}

void DecayChain::AddElement(const std::string Name, int iDecay, int iLocation,
			      double tau, double wt)
{
#ifdef DEBUG
  info << "RAT:: DecayChain::AddElement()\n";
#endif
  int iChain = GetNuberOfElements();

  fElement[iChain] = new BetaFunction(Name);

  fElement[iChain] -> SetInputFileName(GetInputFileName());
#ifdef DEBUG
  fElement[iChain] -> SetIsVerbose(true);
#endif
  bool scr = GetScreenCorr();
  bool iFound = fElement[iChain]->ReadInputFile(Name,iDecay,scr);
  SetElementName(iChain, Name);
  SetElementBranch(iChain, iLocation);
  SetElementDecay(iChain, iDecay);
  SetLifetime(iChain, tau);
  SetElementWt(iChain, wt);

  fNumberOfElements++;
  SetDecaySequence();

  if (fIsVerbose) {
     if (!iFound)
	     printf("No such element found in decay scheme. \n");
     fElement[iChain]->Show();
  }
}

void DecayChain::SetElementName(int iBranch, const std::string Name)
{
  if (iBranch <= fNumberOfElements) {
     fElementName[iBranch] = Name;
  }
}

void DecayChain::SetElementNumber(int iBranch)
{
  if (iBranch <= fNumberOfElements)
     fElementNumber = iBranch;
}

void DecayChain::SetElementBranch(int iBranch, int iLocation)
{
  if (iBranch <= fNumberOfElements)
     fElementBranch[iBranch] = iLocation;
}

void DecayChain::SetElementWt(int iBranch, double wt)
{
  if (iBranch <= fNumberOfElements)
     fElementWeight[iBranch] = wt;
}

void DecayChain::SetElementDecay(int iBranch, int iDecay)
{
  if (iBranch <= fNumberOfElements)
     fElementDecay[iBranch] = iDecay;
}

void DecayChain::SetLifetime(int iBranch, double tau)
{
  if (iBranch <= fNumberOfElements)
     fElementLife[iBranch] = tau;
}

bool DecayChain::ReadInputFile(const std::string dName)
{

  string dProbe = "CHAIN:";
  string eProbe = "END";
  char dummy[80];
  char tName[80];
  char sName[80];

  bool iFound = false;

  int eP = 0;
  int iChain, iDecay;
  float weight, tau;

  FILE *inputFile;

  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);
        string iString(dummy);

        if (iString == dProbe) {
          Log::Assert( fscanf(inputFile, "%s", tName) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
           string iString(tName);
           if (dName == iString && !fInMiddleChain && !fAlphaDecayStart && !fGammaDecayStart) {
	           iFound = true;
               Log::Assert( fscanf(inputFile, "%d", &eP) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
              if (fIsVerbose)
		           printf("Reading %s \n \n", tName);
	           for (int j = 0; j < eP; j++) {
                 Log::Assert( fscanf(inputFile, "%s %d %f %d %f", sName, &iChain,
                                     &weight, &iDecay, &tau) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
		            string iString(sName);
                  if (iDecay != NullParticle) {
		               AddElement(iString, iDecay, iChain, (double) tau,(double) weight);
		            }
		        }
	        }
           else if (dName !=tName && (fInMiddleChain)){
             Log::Assert( fscanf(inputFile, "%d", &eP) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
              bool partchain =false;
              int el=0;
              while (el<eP ){
                Log::Assert( fscanf(inputFile, "%s %d %f %d %f", sName, &iChain,
                                    &weight, &iDecay, &tau) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
                           string iString(sName);
                    if ((iString==dName && weight==1) || (partchain) ){
                       if (iDecay != NullParticle)
                          AddElement(iString, iDecay, iChain, (double) tau,(double) weight);
                       partchain=true;
                       iFound = true;
                    }
                    el++;
              }
           }
           else if (fAlphaDecayStart){
             Log::Assert( fscanf(inputFile, "%d", &eP) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
              int el=0;
              while (el<eP ){
                Log::Assert( fscanf(inputFile, "%s %d %f %d %f", sName, &iChain,
                                    &weight, &iDecay, &tau) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
                    std::string iString(sName);
                    if (iString==dName && iDecay==2){
                       if (weight<1) weight=1;
                          AddElement(iString, iDecay, iChain, (double) tau,(double) weight);
                       iFound = true;
                       iRead = false;
                       el=999;
                    }
                    el++;
              }
           }
           else if (fGammaDecayStart){
             Log::Assert( fscanf(inputFile,"%d",&eP) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
              int el=0;
              while (el<eP){
                Log::Assert( fscanf(inputFile, "%s %d %f %d %f", sName, &iChain,
                                    &weight, &iDecay, &tau) > 0, "DecayChain::ReadInputFile: File format is incorrect." );
                    std::string iString(sName);
                    if (iString==dName && iDecay==0){
                       if (weight<1) weight=1;
                       AddElement(iString, iDecay, iChain, (double)tau,
                                  (double)weight);
                       iFound = true;
                       iRead = false;
                       el = 999;
                    }
                    el++;
              }
           }
           else if (iString == eProbe) {
	           iRead = false;
	        }
        }

  } // end while

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

     if (fAlphaDecayStart){
        printf("The element %s undergoes alpha decay...\n", dName.c_str());
        AddElement(dName, 2);
        iFound = true;
     }
     else if(fGammaDecayStart){
        printf("The element %s undergoes gamma decay...\n", dName.c_str());
        AddElement(dName, 0);
        iFound = true;
     }
     else{
        printf("The element %s undergoes beta decay...\n", dName.c_str());
        AddElement(dName);
        iFound = true;
     }
  }
  return iFound;
}

const std::string DecayChain::GetElementName(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return "";
     return fElementName[iBranch];
}

int DecayChain::GetElementBranch(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return 0;

  return fElementBranch[iBranch];
}

int DecayChain::GetElementDecay(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return 0;

  return fElementDecay[iBranch];
}

double DecayChain::GetElementWt(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return 0.;

  return fElementWeight[iBranch];
}

double DecayChain::GetLifetime(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return 0.;

  return fElementLife[iBranch];
}

void DecayChain::GenerateDecayElement(int i)
{
  int nE = GetNuberOfElements();
  if (nE <= 0) return;

  fElement[i]->GenerateEvent();

  int nP = fElement[i]->GetNGenerated();

  fTstart += fElement[i]->GetEventTime();
  if (nP > 0) {
     int iStart = GetNGenerated();

     for (int j = 0; j < nP; j++) {
	      int k = iStart + j;
	      SetParentName(k, fElement[i]->GetTargetName());
              SetParentMass(k, fElement[i]->GetParentMass());
              SetParentCharge(k,fElement[i]->GetParentCharge());
	      int iNext = i;
	      if (i > GetNuberOfElements()-1) {
	         for (int q = GetNuberOfElements()-1; q > i; q--){
	             if (CheckInChain(i,q)) iNext = q;
	         }
	      }
	      SetDaughterName(k, fElement[iNext]->GetTargetName());
	      fParticleID[k]     = fElement[i]->GetEventID(j);
	      fParticleEnergy[k] = fElement[i]->GetEventEnergy(j);
	      SetEventTime(k,fTstart);
	      SetParticleInfo(i, j, k);
     }

// The daughter information is not currently used = fIsEquilibrium is set to
// false when generating the full chain
     bool iDaughter = ((fIncludeDaughter) && (fIsEquilibrium));
     if (iDaughter){
	     int k = iStart + nP;
	     int isoID = 100000000;
	     isoID += (int) (fElement[i]->GetTargetMass()) * 1000;
	     isoID += (int) (fElement[i]->GetTargetCharge());
	     ParticleInfo_t theDaughter = AddDaughterInfo(i);
	     SetParentName(k, fElement[i]->GetTargetName());
             SetParentMass(k, fElement[i]->GetParentMass());
             SetParentCharge(k, fElement[i]->GetParentCharge());
	     int iNext = i;
	     if (i > GetNuberOfElements()-1) {
	        for (int q = GetNuberOfElements()-1; q > i; q--){
	            if (CheckInChain(i,q)) iNext = q;
	        }
	     }
	     SetDaughterName(k, fElement[iNext]->GetTargetName());
	     fParticleID[k]     = isoID;
	     fParticleEnergy[k] = theDaughter.vector.e();
	     SetEventTime(k,fTstart);
	     fParticleInfo[k] = theDaughter;
	     nP++;
    }
  }
  SetNGenerated(GetNGenerated()+nP);
  SetElementNumber(i);
}

void DecayChain::GenerateDecayElement(const std::string iElement)
{
  for (int i=0; i<fNumberOfElements; i++) {
      const std::string tName = GetElementName(i);
      if (tName == iElement) {
	      GenerateDecayElement(i);
	      return;
      }
  }
}

void DecayChain::GenerateDecayChain(const std::string iNameStart)
{

  SetIsEquilibrium(true);
  int nE = GetNuberOfElements();
  int iStart = 0;
  int iEnd = nE;

  if (nE <= 0)
     return;

  for (int i = iStart; i < iEnd; i++) {
      const string tName = GetElementName(i);
      if (tName == iNameStart)
	       iStart = i;
  }

  double rMax = 0.;
  for (int i = iStart; i < iEnd; i++)
      rMax = rMax + GetElementProb(i);

  SetNGenerated(0);
  fIsChainElement = false;
  double r = rMax * GetRandomNumber();
  double Prob = 0.;
  int iTag = -1;
  for (int i = iStart; i < iEnd; i++) {
      Prob = Prob + GetElementProb(i);
      if ((iTag < 0) && (r < Prob))
	      iTag = i;
  }

  fTstart = 0.;
  GenerateDecayElement(iTag);
}

void DecayChain::GenerateDecayChain(int iBranch)
{
  int nE = GetNuberOfElements();
  if (iBranch > nE)
     return;
  const string tName = GetElementName(iBranch);
  GenerateDecayChain(tName);
}

const string DecayChain::GetParentName(int iEvent)
{
  if (iEvent > GetNGenerated()) return "";
  return fParentName[iEvent];
}

double DecayChain::GetParentMass(int iEvent)
{
  if (iEvent > GetNGenerated()) return 0.;
  return fParentMass[iEvent];
}

double DecayChain::GetParentCharge(int iEvent)
{
  if (iEvent > GetNGenerated()) return 0.;
  return fParentCharge[iEvent];
}

const std::string DecayChain::GetDaughterName(int iEvent)
{
  if (iEvent > GetNGenerated()) return "";
  return fDaughterName[iEvent];
}

int DecayChain::GetEventID(int iEvent)
{
  if (iEvent > GetNGenerated()) return 0;
  if (!CheckParentMatch(iEvent)) return 0;

  return fParticleID[iEvent];
}

double DecayChain::GetEventEnergy(int iEvent)
{
  if (iEvent > GetNGenerated()) return 0.;
  if (!CheckParentMatch(iEvent)) return 0;

  return fParticleEnergy[iEvent];
}

double DecayChain::GetEventTime(int iEvent)
{
  if (iEvent > GetNGenerated()) return 0.;
  if (!CheckParentMatch(iEvent)) return 0;

  return fParticleTime[iEvent];
}

void DecayChain::SetEventTime(int iEvent, double time)
{
  fParticleTime[iEvent] = time;
}

void DecayChain::GenerateFullChain(const std::string iElement)
{
  SetIsEquilibrium(false); // This kills including the daughter information!
  int nE = GetNuberOfElements();
  if (nE <= 0) return;

  int iStart = 0;
  for (int i = 0; i < nE; i++) {
      const string tName = GetElementName(i);
      if (tName == iElement)
	      iStart = i;
  }

  int iBefore = iStart;
  int iReject = -1;
  double Br = 1.;
  double RelWt = 1.;
  fTstart = 0.;
  SetNGenerated(0);
  fIsChainElement = false;

  for (int j = iStart; j < nE; j++) {
      bool iCheck = (CheckInChain(j, iStart) && CheckInChain(j, iBefore));

      if (iReject > 0)
 	      iCheck = ((iCheck) && (!CheckInChain(j, iReject)));

      if (iCheck) {
	      Br = GetElementWt(j);
         double Prob = RelWt * GetRandomNumber();
         if (Prob <= Br) {
	         GenerateDecayElement(j);
	         iBefore = j;
	         iReject = -1;
	         RelWt = 1.;
	      }
         else {
	         RelWt = RelWt - Br;
	         iReject = j;
         }
      }
      fIsChainElement = true;
  }
  fIsChainElement = false;

  double tMax = 1.0e+99;
  for (int i=0;i<GetNGenerated();i++){
      if (GetEventTime(i)<tMax) tMax = GetEventTime(i);
  }

  for(int i=0;i<GetNGenerated();i++) SetEventTime(i,GetEventTime(i)-tMax);
}

void DecayChain::GenerateFullChain(int iBranch)
{
  int nE = GetNuberOfElements();
  if (iBranch > nE)
     return;
  const string tName = GetElementName(iBranch);
  GenerateFullChain(tName);
}

bool DecayChain::CheckParentMatch(int iEvent)
{
  bool iPass = false;

  if (fSelectParent == "all") iPass = true;
  if (fSelectParent == "ALL") iPass = true;
  if (fSelectParent == "0/")  iPass = true;

  if (fSelectParent == GetParentName(iEvent)) iPass = true;

  return iPass;
}

void DecayChain::SetDecaySequence()
{
  int iStart = 0;
  int iEnd = GetNuberOfElements();
  double BrL[5];
  for (int j = iStart; j < iEnd; j++) {
      int jLevel = CheckChainLevel(j);
      BrL[jLevel] = GetElementWt(j);
      for (int k = iStart; k < j; k++) {
	       int kLevel = CheckChainLevel(k);
          if (CheckInChain(k, j)){
	          BrL[kLevel] = BrL[kLevel] * GetElementWt(k);
          }
      }

      double Br = 1;
      for (int l=0; l <= jLevel; l++) Br *= BrL[l];

      if (CheckInChain(iStart, j)) {

	      double weight = 1.;
	      double A_Source = fElement[j]->GetTargetMass();
        	if (A_Source > 0) {

	         double Lambda = log(2.) / GetLifetime(iStart);
	         double Lambda_d = log(2.) / GetLifetime(j);
	         double LTau = 1.;

	         if (Lambda_d > 1.e-12)
	            LTau = (1. - exp(-1. * Lambda_d)) / Lambda_d;

	         weight = Br * (N_A / A_Source) * Lambda * LTau;

	      }

	      if (GetIsEquilibrium()) weight = Br;

	      SetElementProb(j, weight);
      }
  }
}

void DecayChain::SetElementProb(int iBranch, double prob)
{
  if (iBranch <= GetNuberOfElements())
     fElementProb[iBranch] = prob;
}

double DecayChain::GetElementProb(int iBranch)
{
  if (iBranch > GetNuberOfElements())
     return 0;
  return fElementProb[iBranch];
}

bool DecayChain::CheckInChain(int iElement, int jElement)
{
  bool inChain = true;
  int index0 = 0;
  int index1 = 0;
  for (int l = 0; l < 3; l++) {
      int divisor = (int) pow(10., 1. * l);
      index0 = (GetElementBranch(iElement) / divisor) % 10;
      index1 = (GetElementBranch(jElement) / divisor) % 10;
      if (((index0 * index1) != 0) && (index1 != index0))
	      inChain = false;
  }
  return inChain;
}

int DecayChain::CheckChainLevel(int iElement)
{
  int level = 0;
  for (int l = 0; l < 3; l++) {
      int divisor = (int) pow(10., 1. * l);
      int index = (GetElementBranch(iElement) / divisor) % 10;
      if (index != 0) level = l;
  }
  return level;
}

void DecayChain::SetParticleInfo(int iTag, int jTag, int kTag)
{
  int pid = fElement[iTag] -> GetEventID(jTag);

  double mass = 0.;
  if ((pid == DecayBeta) || (pid ==DecayEC)) mass = ElectronMass;
  if (pid == DecayAlpha) mass = AlphaMass;

  double energy = fElement[iTag] -> GetEventEnergy(jTag);
  energy += mass;
  double pp  = sqrt(pow(energy,2) - pow(mass,2));
  double cz  = GetRandomNumber(-1.,1.);
  double phi = GetRandomNumber(0.,twopi);

  double px = pp * sqrt(1. - pow(cz,2)) * cos(phi);
  double py = pp * sqrt(1. - pow(cz,2)) * sin(phi);
  double pz = pp * cz;

  fParticleInfo[kTag].vector.setE(energy);
  fParticleInfo[kTag].vector.setPx(px);
  fParticleInfo[kTag].vector.setPy(py);
  fParticleInfo[kTag].vector.setPz(pz);
  fParticleInfo[kTag].ID = pid;

}

DecayChain::ParticleInfo_t DecayChain::AddDaughterInfo(int iTag)
{
  ParticleInfo_t theDaughterInfo;

  int iType = fElement[iTag] -> GetTargetDecayType();
  double A  = fElement[iTag] -> GetTargetMass();
  double Z  = fElement[iTag] -> GetTargetCharge();

  double massI  = 0.;
  if (iType == DecayBeta) {massI = ElectronMass;}
  if (iType == DecayEC)   {massI = ElectronMass;}
  if (iType == DecayAlpha){massI = AlphaMass;}
// Daughter information for DecayGamma not input yet... may be necessary if
// daughters are eventually used.

  double MassD = Nucl_Mass(A,Z);
  double Q  = 0.;
  for (int i = 0; i<fElement[iTag]->GetNParticles(0); i++){
      Q += fElement[iTag] -> GetEnergy(0,i);
  }

  double TRecoil = Q * (massI / MassD) + pow(Q,2.)/(2.*MassD);
  double ERecoil = TRecoil + MassD;
  double PRecoil = sqrt(pow(ERecoil,2.) - pow(MassD,2.));
  double cz  = GetRandomNumber(-1.,1.);
  double phi = GetRandomNumber(0.,twopi);

  double px = PRecoil * sqrt(1. - pow(cz,2)) * cos(phi);
  double py = PRecoil * sqrt(1. - pow(cz,2)) * sin(phi);
  double pz = PRecoil * cz;

  int isoID = 100000000;
  isoID += (int) A * 1000;
  isoID += (int) Z;

  theDaughterInfo.vector.setE(ERecoil);
  theDaughterInfo.vector.setPx(px);
  theDaughterInfo.vector.setPy(py);
  theDaughterInfo.vector.setPz(pz);
  theDaughterInfo.ID = isoID;

  return theDaughterInfo;
}

DecayChain::ParticleInfo_t DecayChain::GetParticleInfo(int iEvent)
{
  ParticleInfo_t blankInfo = {0, CLHEP::HepLorentzVector()};

  if (iEvent > GetNGenerated()) return blankInfo;
  if (!CheckParentMatch(iEvent)) return blankInfo;

  return fParticleInfo[iEvent];
}

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

}//end