MODULE materials_mod

USE bmad
IMPLICIT NONE

TYPE material
  REAL(RP) :: Z       ! atomic number              [e]
  REAL(RP) :: A       ! atomic mass                [g/mol]
  REAL(RP) :: density ! density                    [g/cm3]
  REAL(RP) :: radl    ! radiation length           [cm]
  REAL(RP) :: nucintl ! nuclear interaction length [cm]
  REAL(RP) :: imean   ! mean excitation energy     [eV]
END TYPE material

! Elements
TYPE (material) :: Be
TYPE (material) :: Al
TYPE (material) :: Si
TYPE (material) :: Ti
TYPE (material) :: Cu
TYPE (material) :: Nb

! Compounds
TYPE (material) :: Kapton
TYPE (material) :: InflectorCoil

CONTAINS
SUBROUTINE initializeMaterials()
  IMPLICIT NONE

  ! NIST MATERIALS DATABASE

  ! Elements
  Be%Z       = 4.     ! e
  Be%A       = 9.01   ! g/mol
  Be%density = 1.848  ! g/cm3
  Be%radL    = 35.276 ! cm
  Be%nucIntL = 39.449 ! cm
  Be%Imean   = 63.7   ! eV

  Al%Z       = 13.    ! e
  Al%A       = 26.98  ! g/mol
  Al%density = 2.699  ! g/cm3
  Al%radL    = 8.896  ! cm
  Al%nucIntL = 38.877 ! cm
  Al%Imean   = 166.0  ! eV

  Si%Z       = 14.    ! e
  Si%A       = 28.09  ! g/mol
  Si%density = 2.330  ! g/cm3
  Si%radL    = 9.366  ! cm
  Si%nucIntL = 45.753 ! cm
  Si%Imean   = 173.0  ! eV

  Ti%Z       = 22.    ! e
  Ti%A       = 47.87  ! g/mol
  Ti%density = 4.54   ! g/cm3
  Ti%radL    = 3.560  ! cm
  Ti%nucIntL = 27.939 ! cm
  Ti%Imean   = 233.0  ! eV

  Cu%Z       = 29.    ! e
  Cu%A       = 63.55  ! g/mol
  Cu%density = 8.96   ! g/cm3
  Cu%radL    = 1.436  ! cm
  Cu%nucIntL = 15.514 ! cm
  Cu%Imean   = 322.0  ! eV

  Nb%Z       = 41.    ! e
  Nb%A       = 92.91  ! g/mol
  Nb%density = 8.57   ! g/cm3
  Nb%radL    = 1.158  ! cm
  Nb%nucIntL = 18.485 ! cm
  Nb%Imean   = 417.0  ! eV

  ! Kapton
  Kapton%Z       = 6.35993 ! calculated from mass fractions
  Kapton%A       = 12.7014 ! calculated from mass fractions
  Kapton%density = 1.42    ! g/cm3
  Kapton%radL    = 28.575  ! cm
  Kapton%nucIntL = 55.867  ! cm
  Kapton%Imean   = 79.6    ! eV

  ! Inflector coil, i.e. kapton-wrapped, aluminum-stabilized NbTi/Cu superconducting wires,
  ! with Nb:Ti:Cu:Al:Kapton:Vacuum ( 0.5 : 0.5 : 0.9 : 3.7 : 1.41126 : 0.476108 )
  InflectorCoil%Z       = 22.8882 ! calculated from mass fractions
  InflectorCoil%A       = 48.9843 ! calculated from mass fractions
  InflectorCoil%density = 3.554   ! g/cm3
  InflectorCoil%radL    = 4.539   ! cm
  InflectorCoil%nucIntL = 33.338  ! cm
  InflectorCoil%Imean   = 221.129 ! eV

END SUBROUTINE initializeMaterials

SUBROUTINE scatter(mater,dx,co)
  USE bmad
  USE nr
  USE parameters_bmad
  IMPLICIT NONE

  TYPE(material),     INTENT(IN)    :: mater ! material type
  REAL(RP),           INTENT(IN)    :: dx    ! length over which to apply scattering (PDG notation)
  TYPE(coord_struct), INTENT(INOUT) :: co    ! coordinates

  REAL(RP) :: dp_p, p, betagamma, rapidity
  REAL(RP) :: x, xp, dxp, y, yp, dyp

  REAL(RP) :: X0, theta0
  REAL(RP) :: rand1, rand2

  ! For brevity
  x    = co%vec(1)
  xp   = co%vec(2)
  y    = co%vec(3)
  yp   = co%vec(4)
  X0   = mater%radl

  ! Calculate energy, momentum, etc.
  dp_p      = co%vec(6)          ! fractional momentum spread (unitless)
  p         = ( 1+dp_p )*pMagic  ! momentum magnitude (eV)
  betagamma = p/mmu              ! Lorentz betagamma (unitless)
  rapidity  = asinh(betagamma)   ! hyperbolic rotation angle

  ! "Multiple scattering through small angles," Particle Data Group (2013), Eq. (27.15).  Our momentum 
  ! units here are eV, not eV/c, so there is no need to multiply the momentum by c_light in the formula.
  theta0 = 13.6e6/( tanh(rapidity)*p ) * sqrt(dx/X0) * ( 1+0.038*log(dx/X0) )

  ! Apply scattering angle
  call gasdev(rand1)
  call gasdev(rand2)
  dxp = rand1*theta0
  dyp = rand2*theta0
  xp = xp + dxp
  yp = yp + dyp

  ! Apply energy loss -- changes co%vec(6)
  ! call energyLoss(mater,dx,co)

  ! Calculate change in particle's transverse spatial coordinates by integrating the angles across the length 
  ! of the material.  (Since the angles are so small here, this nearly looks like propagation in a drift space!)
  x = x + tan(xp)*dx                    ! the tangent is the slope, dx is the distance
  y = y + tan(yp)*dx*sqrt(1+tan(xp)**2) ! the tangent is the slope, dx*sqrt() is the distance

  ! Reassign variables
  co%vec(1) = x
  co%vec(2) = xp
  co%vec(3) = y
  co%vec(4) = yp

  ! Momentum components
  ! py = p*sin(yp)
  ! px = p*cos(yp)*sin(xp)
  ! pz = p*cos(yp)*sin(xp)

END SUBROUTINE scatter


SUBROUTINE energyLoss(mater,dx,co)
  USE bmad
  USE nr
  USE parameters_bmad
  IMPLICIT NONE

  TYPE(material),     INTENT(IN)    :: mater ! material type
  REAL(RP),           INTENT(IN)    :: dx    ! length over which to apply scattering (PDG notation)
  TYPE(coord_struct), INTENT(INOUT) :: co    ! coordinates

  ! The following are used in the "Bethe Equation," PDG(2013) Eq. (27.4), and the Landau-Vavilov distribution PDG(2013) Figure 27.7.
  REAL(RP) :: K = 0.307075               ! PDG(2013) Table 27.1 [MeV*cm2/g]
  REAL(RP) :: Tmax                       ! PDG(2013) Eq. (27.5)
  REAL(RP) :: plasmaEnergy               ! PDG(2013) Table 27.1 [eV]
  REAL(RP) :: delta_2                    ! PDG(2013) Eq. (27.6)
  REAL(RP) :: MostProbableEnergyLoss, xi ! PDG(2013) Eq. 27.11

  ! Relativistic kinematics, energy loss, etc.
  REAL(RP) :: p, dp_p, betagamma, rapidity
  REAL(RP) :: E, dE, dE_avg, dE_dx

  ! Variables for generating random energy-loss fluctuations from the Landau-Vavilov distribution
  REAL(RP) :: rand, di, frac_low, frac_high, i_interp
  INTEGER  :: i

  ! Calculate energy, momentum, etc.
  dp_p      = co%vec(6)          ! fractional momentum spread (unitless)
  p         = ( 1+dp_p )*pMagic  ! momentum magnitude (eV)
  betagamma = p/mmu              ! Lorentz betagamma (unitless)
  rapidity  = asinh(betagamma)   ! hyperbolic rotation angle
  E         = cosh(rapidity)*mmu ! energy (eV)

  ! Before we make use of the "Bethe Equation" [PDG(2013) Eq. (27.4)], etc., we need to calculate a few things

  ! "Maximum kinetic energy transferred in a collision," PDG(2013) Eq. (27.5)
  Tmax = 2*me*betagamma**2/( 1 + 2*cosh(rapidity)*me/mmu + (me/mmu)**2 ) ! eV

  ! "Plasma energy," defined in PDG(2013) Table 27.1
  plasmaEnergy = sqrt(mater%density*mater%Z/mater%A)*28.816 ! eV

  ! "Density effect," PDG(2013) Eq. (27.6)
  delta_2 = log(plasmaEnergy/mater%Imean) + log(betagamma) - 0.5 ! dimensionless

  ! Now let's put everything together in the "Bethe Equation," PDG(2013) Eq. (27.4)
  dE_dx = -K*mater%Z/mater%A/tanh(rapidity)**2 * ( 0.5*log(2*me*betagamma**2*Tmax/mater%Imean**2) - tanh(rapidity)**2 - delta_2 ) ! MeV*cm2/g

  ! Convert to MeV/cm
  dE_dx = dE_dx * mater%density ! MeV/cm

  ! Calculate the average energy loss
  dE_avg = dE_dx * dx * (1.e6) ! eV

  ! So far, we have only calculated the /average/ energy shift from the Bethe Equation.  How do we include 
  ! fluctuations?  The answer is shown in PDG(2013) Figure 27.7, i.e. the Landa-Vavilov distribution.

  ! Fluctuations in energy loss (Landau-Vavilov distribution) -- PDG(2013) Eq. (27.11) and surrounding text
  xi = K/2.*mater%Z/mater%A*dx*mater%density/tanh(rapidity)**2 ! text immediately beneath Eq. (27.11)
  MostProbableEnergyLoss = xi*( log(2*me*betagamma**2/mater%Imean) + log(xi/mater%Imean) + 0.200 - tanh(rapidity)**2 + 2*delta_2 ) ! Eq. (27.11)

  ! Generate a random number, find it in the Landau-Vavilov integral table (parameters_bmad.f90), and map it to an energy loss.
  call ran_uniform(rand)
  DO i=-5, 100-1 ! bounds of integral table
    IF ( LandauVavilov_integral(i)<rand .and. rand<LandauVavilov_integral(i+1) ) THEN
      di = LandauVavilov_integral(i+1) - LandauVavilov_integral(i)  ! values in the LandauVavilov_integral table (y-axis)
      frac_low  = ( LandauVavilov_integral(i+1) - rand )/di         ! fraction of lower value to include (linear interpolation)
      frac_high = ( rand -  LandauVavilov_integral(i)  )/di         ! fraction of upper value to include (linear interpolation)
      i_interp  = frac_low*i + frac_high*(i+1)                      ! corresponding "fractional index" (x-axis)
      dE = -(i_interp*xi + MostProbableEnergyLoss)*1.e6             ! map to an energy loss [eV]
      IF (i==99) THEN
        ! It turns out the Landau-Vavilov distribution has an EXTREMELY long tail, and so the very last little bit of integral table is handled differently.
        ! We know that the maximum energy loss is the kinetic energy of the muon, Tmax.  So linearly interpolate between dE[i=99] and Tmax
        dE = -( frac_low*(99*xi + MostProbableEnergyLoss)*1.e6 + frac_high*Tmax ) ! eV
      ENDIF
      EXIT
    ENDIF
  ENDDO

  ! TODO: See if you can reproduce PDG(2013) Figure 27.7 for 10-GeV muons traversing 1.7mm of silicon using your Landau-Vavilov-based numerical techniques
  ! DONE: I looked at 3.094 GeV mu+ on 1.0cm Cu, 1.5mm Cu, and 1.5mm Al -- the energy fluctuations seem to make good sense

  ! Change the energy
  E = E + dE

  ! Calculate the new momentum from the new energy
  betagamma = sinh(acosh(E/mmu))
  dp_p      = mmu*betagamma/pMagic - 1.

  ! Assign the new momentum
  co%vec(6) = dp_p

END SUBROUTINE energyLoss

END MODULE materials_mod
