/*

EnergyLoss.cpp
Code for calculating the energy loss of a charged, massive particle through an arbitrary medium.
Based on code written by D.W. Visser at Yale for the original SPANC. Uses energy loss calulations
described by Ziegler in various SRIM textbooks.

Written by G.W. McCann Aug. 2020

*/

#include "Eloss_Tables.h"
#include "EnergyLoss.h"
#include "KinematicsExceptions.h"
#include <iostream>

EnergyLoss::EnergyLoss() {
  comp_denom = 0;
  ZP = -1;
}

EnergyLoss::~EnergyLoss() {}

/*Targets are defined by their atomic number, total number of nucleons, and their stoichiometry within the target compound*/
void EnergyLoss::SetTargetComponents(std::vector<int>& Zt, std::vector<int>& At, std::vector<int>& Stoich) {
  comp_denom = 0;
  ZT = Zt;
  AT = At;
  for(unsigned int i=0; i<Stoich.size(); i++) {
    comp_denom += Stoich[i];
    if(ZT[i] > MAX_Z) {
      throw ELossException();
    }
  }
  targ_composition.resize(Stoich.size());
  for(unsigned int i=0; i<Stoich.size(); i++) {
    targ_composition[i] = Stoich[i]/comp_denom;
  }
}

/*
  Returns units of MeV; thickness in ug/cm^2; e_initial in units of MeV
  Energy loss going through the target
 */
double EnergyLoss::GetEnergyLoss(int zp, int ap, double e_initial, double thickness) {
  if( ZP != zp) { 
    ZP = zp;
    AP = ap;
    MP = MASS.FindMass(ZP, AP)*MEV2U;
  }

  double e_final = e_initial;
  double x_traversed = 0;
  double x_step = 0.25*thickness; //initial step in x
  double e_step = GetTotalStoppingPower(e_final)*x_step/1000.0; //initial step in e
  double e_threshold = 0.05*e_initial;

  int depth=0;

  
  if(thickness == 0.0) return 0;
  else if(e_initial == 0.0) return 0;

  bool go = true;
  while(go) {
    //If intial guess of step size is too large, shrink until in range
    if(e_step/e_final > MAX_FRACTIONAL_STEP && depth < MAX_DEPTH) {
      depth++;
      x_step *= 0.5;
      e_step = GetTotalStoppingPower(e_final)*x_step/1000.0;
    } else if((x_step + x_traversed) >= thickness) { //last chunk
      go = false;
      x_step = thickness - x_traversed; //get valid portion of last chunk
      e_final -= GetTotalStoppingPower(e_final)*x_step/1000.0;
      if(depth > 20)std::cout<<"depth: "<<depth<<std::endl;
      if(e_final <= e_threshold) {
        return e_initial;
      }
    } else if(depth == MAX_DEPTH) {
      return e_initial;
    } else {
      x_traversed += x_step;
      e_step = GetTotalStoppingPower(e_final)*x_step/1000.0;
      e_final -= e_step;
      if(e_final <= e_threshold) {
        return e_initial;
      }
    }
  }
  return e_initial - e_final;
}

/*
  Returns units of MeV; thickness in ug/cm^2; e_final in units of MeV
  Energy loss going through the target using energy of particle after traversal
*/
double EnergyLoss::GetReverseEnergyLoss(int zp, int ap, double e_final, double thickness) {
  if( ZP != zp) { 
    ZP = zp;
    AP = ap;
    MP = MASS.FindMass(ZP, AP)*MEV2U;
  }

  double e_initial = e_final;
  double x_traversed = 0.0;
  double x_step = 0.25*thickness; //initial step in x
  double e_step = GetTotalStoppingPower(e_initial)*x_step/1000.0; //initial step in E

  bool go = true;
  while(go) {
    if(e_step/e_initial > MAX_FRACTIONAL_STEP) {
      x_step *= 0.5;
      e_step = GetTotalStoppingPower(e_initial)*x_step/1000.0;
    } else if (x_traversed+x_step > thickness) {
      go = false;
      x_step = thickness - x_traversed;
      e_initial += GetTotalStoppingPower(e_initial)*x_step/1000.0;
    } else {
      x_traversed += x_step;
      e_step = GetTotalStoppingPower(e_initial)*x_step/1000.0;
      e_initial += e_step;
    }
  }

  return e_initial-e_final;
}

/*
  Returns units of keV/(ug/cm^2)
  Calculates Electronic stopping power by first calculating SE for hydrogen through the target and then using
  corrections to calculate SE for projectile of interest
*/
double EnergyLoss::GetElectronicStoppingPower(double energy) {
  //Wants in units of keV
  energy *= 1000.0;
  double e_per_u = energy/MP;
  std::vector<double> values;
  if(e_per_u > MAX_H_E_PER_U) {
    throw ELossException();
  } else if (e_per_u > 1000.0) {
    for(auto& z: ZT) {
      values.push_back(Hydrogen_dEdx_High(e_per_u, energy, z));
    }
  } else if (e_per_u > 10.0) {
    for(auto& z: ZT) {
      values.push_back(Hydrogen_dEdx_Med(e_per_u, z));
    }
  } else if (e_per_u > 0.0) {
    for(auto& z: ZT) {
      values.push_back(Hydrogen_dEdx_Low(e_per_u, z));
    }
  } else {
    throw ELossException();
  }

  if(values.size() == 0) {
    throw ELossException();
  }

  if(ZP > 1) { //not hydrogen, need to account for effective charge
    for(unsigned int i=0; i<values.size(); i++) {
      values[i] *= CalculateEffectiveChargeRatio(e_per_u, ZT[i]);
    }
  }
  
  double stopping_total = 0;
  double conversion_factor = 0;
  for(unsigned int i=0; i<ZT.size(); i++) {
    conversion_factor += targ_composition[i]*NATURAL_MASS[ZT[i]];
    stopping_total += values[i]*targ_composition[i];
  }
  stopping_total *= AVOGADRO/conversion_factor;

  return stopping_total;
}

//Returns units of keV/(ug/cm^2)
double EnergyLoss::GetNuclearStoppingPower(double energy) {
  energy *= 1000.0;
  double stopping_total = 0.0;
  double sn, x, epsilon, conversion_factor;
  for(unsigned int i=0; i<ZT.size(); i++) {
    x = (MP + NATURAL_MASS[ZT[i]]) * std::sqrt(std::pow(ZP, 2.0/3.0) + std::pow(ZT[i], 2.0/3.0));
    epsilon = 32.53*NATURAL_MASS[ZT[i]]*energy/(ZP*ZT[i]*x);
    sn = 8.462*(0.5*std::log(1.0+epsilon)/(epsilon+0.10718*std::pow(epsilon, 0.37544)))*ZP*ZT[i]*MP/x;
    conversion_factor = AVOGADRO/NATURAL_MASS[ZT[i]];
    stopping_total += sn*conversion_factor*targ_composition[i];
  }

  return stopping_total;
}

/*Wrapper function for aquiring total stopping (elec + nuc)*/
double EnergyLoss::GetTotalStoppingPower(double energy) {
  return GetElectronicStoppingPower(energy)+GetNuclearStoppingPower(energy);
}

/*Charge rel to H*/
double EnergyLoss::CalculateEffectiveChargeRatio(double e_per_u, int z) {
  double z_ratio;
  if(ZP == 2) {
    double ln_epu = std::log(e_per_u);
    double gamma = 1.0+(0.007+0.00005*z)*std::exp(-std::pow(7.6-ln_epu,2.0));
    double alpha = 0.7446 + 0.1429*ln_epu + 0.01562*std::pow(ln_epu, 2.0) - 0.00267*std::pow(ln_epu,3.0)
                   + 1.338E-6*std::pow(ln_epu,8.0);
     z_ratio = gamma*(1.0-std::exp(-alpha))*2.0; //test this; SPANC has factor of 2. mult
  } else if (ZP == 3) {
    double ln_epu = std::log(e_per_u);
    double gamma = 1.0+(0.007+0.00005*z)*std::exp(-std::pow(7.6-ln_epu,2.0));
    double alpha = 0.7138+0.002797*e_per_u+1.348E-6*std::pow(e_per_u, 2.0);
    z_ratio = gamma*(1-std::exp(-alpha))*3.0; //test this; SPANC has factor of 3. mult
  } else {
    double B = 0.886*std::pow(e_per_u/25.0, 0.5)/std::pow(ZP, 2.0/3.0);
    double A = B + 0.0378*std::sin(PI/2.0*B);
    z_ratio = 1.0 - std::exp(-A)*(1.034-0.1777*std::exp(-0.08114*ZP))*z; //test this; SPANC has factor of ZT[i] mult
  }
  return z_ratio*z_ratio; //for stopping power uses ratio sq. 
}

double EnergyLoss::Hydrogen_dEdx_Low(double e_per_u, int z) {
  return std::sqrt(e_per_u)*HYDROGEN_COEFF[z][0];
}

double EnergyLoss::Hydrogen_dEdx_Med(double e_per_u, int z) {
  double x = HYDROGEN_COEFF[z][1]*std::pow(e_per_u, 0.45);
  double y = HYDROGEN_COEFF[z][2]/e_per_u * std::log(1.0+HYDROGEN_COEFF[z][3]/e_per_u+HYDROGEN_COEFF[z][4]*e_per_u);
  return x*y/(x+y);
}

double EnergyLoss::Hydrogen_dEdx_High(double e_per_u, double energy, int z) {
  energy /= 1000.0; //back to MeV for ease of beta calc
  double beta_sq = energy * (energy+2.0*MP/MEV2U)/std::pow(energy+MP/MEV2U, 2.0);
  double alpha = HYDROGEN_COEFF[z][5]/beta_sq;
  double epsilon = HYDROGEN_COEFF[z][6]*beta_sq/(1.0-beta_sq) - beta_sq - HYDROGEN_COEFF[z][7];
  for(int i=1; i<5; i++) {
    epsilon += HYDROGEN_COEFF[z][7+i]*std::pow(std::log(e_per_u), i);
  }
  return alpha * std::log(epsilon);
}

double EnergyLoss::GetRange(double energy) {return 0.0;} //unimplemented