! Supplies the azimuthal derivatives of the position, time and momentum of a charged particle
!   traveling through an EM field in a storage ring

! the input "coords" is an array containing the radial and vertical position, the momentum, and the
! time when the particle is at a particular angle, theta (measured from the center of the storage ring):

!       theta increases as the muon travels clockwise around the storage ring

!       y = (x,y,t,px,py,pzdev) where
!            x = radial position measured from the center of the beam tube
!            y = vertical position
!            px = radial momentum
!            py = vertical momentum
!            pzdev = the difference between azimuthal momentum (tangential to storage ring) and 
!                                            the desired value of pz = magicmomentum
       
! the output "derivs" contains the theta-derivative of each, in the same order

! the differential equations are written in dimensionless form; this requires:
!       position measured in units of (c*timeunit)
!       angle measured in radians
!       momentum measured in units of (mc)
!       electric field measured in units of (mc/q*timeunit)
!       magnetic field measured in units of (m/q*timeunit)
!    where
!           c = speed of light
!           m = mass of muon
!           q = charge of muon
!           timeunit = whatever time unit you are using (probably nanoseconds)

! To use the routine, need to specify Efield and Bfield below

SUBROUTINE derivs(theta,coords,dcoords_dtheta)

USE nrtype
USE precision_def
USE parameters ! module with coordinates and constants
IMPLICIT NONE
REAL(RP), INTENT(IN) :: theta
REAL(RP), DIMENSION(:), INTENT(IN) :: coords
REAL(RP), DIMENSION(:), INTENT(OUT) :: dcoords_dtheta

REAL(RP), DIMENSION(3) :: Efield,Bfield
REAL(RP) :: x,y,t,px,py,pzdev                            ! coordinates
REAL(RP) :: xderiv,yderiv,tderiv,pxderiv,pyderiv,pzderiv ! theta derivates of coordinates
REAL(RP) :: r,pz,pfactor


! Separate the input vector into individual components, for clarity:
x     = coords(1)
y     = coords(2)
t     = coords(3)
px    = coords(4)
py    = coords(5)
pzdev = coords(6)

r = storageradius + x       ! radial position measured from center of circular storage ring
pz = pzdev + magicmomentum ! actual azimuthal momentum


! This quantity appears several times in the differential equations:
pfactor = sqrt(1+(px**2)+(py**2)+(pz**2)) 


! Input the electric and magnetic fields
! Efield = [Ex,Ey,Ez], Bfield=[Bx,By,Bz] (radial component, vertical component, azimuthal component)

if (quad) then
   Efield = quadEfield(x,y) ! function defined below
else
   Efield = [0.0_rp, 0.0_rp, 0.0_rp]
end if

if (kick) then
   Bfield = kickerBfield(x,y,t) + dipoleBfield ! kickerBfield defined in function below
else                                           ! dipoleBfield defined in parameters module
   Bfield = dipoleBfield
end if  



! Here are the differential equations for a charged particle in an EM field:
xderiv  = r*px/pz
yderiv  = r*py/pz
tderiv  = r*pfactor/pz
pxderiv = (r/pz)*(Efield(1)*pfactor + py*Bfield(3) - pz*Bfield(2)) + pz
pyderiv = (r/pz)*(Efield(2)*pfactor + pz*Bfield(1) - px*Bfield(3))
pzderiv = (r/pz)*(Efield(3)*pfactor + px*Bfield(2) - py*Bfield(1)) - px


! Bundle the derivatives back into one vector:
dcoords_dtheta = [xderiv,yderiv,tderiv,pxderiv,pyderiv,pzderiv]


CONTAINS

FUNCTION kickerBfield(x,y,t)
! RLC circuit time dependence; uniform spatial dependence within kicker region
USE parameters ! go here to change physical size, RLC for the circuit, etc.
USE nrtype
USE precision_def
IMPLICIT NONE
REAL(RP), DIMENSION(3) :: kickerBfield
REAL(RP), INTENT(IN) :: x,y,t
REAL(RP) :: damp,natfreq,freq,tzerotomax,maxval,kickertimedep

! RLC circuit time dependence (underdamped harmonic oscillator):
!      kickerR, kickerL and kickerC are set in the parameters module
!      The time origin is when the center of the muon bunch enters the ring
damp = kickerR/(2*kickerL) ! comes out in units of 1/ns
natfreq = 1/sqrt(kickerL*kickerC)
freq = sqrt(natfreq**2 - damp**2)
tzerotomax = (1/freq)*atan(freq/damp) ! time from kicker switched on to max strength
maxval = (freq/natfreq)*exp(-(damp/freq)*atan(freq/damp)) ! will divide out so maximum value is one
kickertimedep = (1/maxval)*exp(-damp*(t+tzerotomax-kickermaxtime))*sin(freq*(t+tzerotomax-kickermaxtime))

! kickerBfield is vertical, downward
kickerBfield = [0.0_rp, -kickermagnitude*kickertimedep, 0.0_rp]
END FUNCTION kickerBfield


FUNCTION quadEfield(x,y)
! only contains a couple of the higher multipoles
USE parameters
USE nrtype
USE precision_def
IMPLICIT NONE
REAL(RP), DIMENSION(3) :: quadEfield
REAL(RP), INTENT(IN) :: x,y
quadEfield = [-2*Qb2*x - 4*Qb4*(x**3-3*x*y**2) - 6*Qb6*(x**5-10*x**3*y**2+5*x*y**4), &
               2*Qb2*y + 4*Qb4*(3*x**2*y-y**3) + 6*Qb6*(5*x**4*y-10*x**2*y**3+y**5), &
               0.0_rp]
END FUNCTION quadEfield



END SUBROUTINE derivs