#input is betaz=
amu=0.00116592
c_light=2.99792458e2
R0=7.112
gamma0= 29.303464   #magic momentum
beta= (1.-1/gamma0**2)**.5


#setup fixed pitch for betaz=?
# this version is for fixed ttransverse momentum
gammaz = 1./(1-betaz**2)**.5
gamma = gammaz * gamma0
betaxy = gamma0/gamma
R = R0
omega = betaxy*c_light/R

omega0=omega*gamma0*(amu+1/gamma0) - amu*gamma0/(1+gamma0) * betaz**2
bomega=omega*gamma0*amu*gamma0/(gamma0+1)*betaz*betaxy
omegap = omega*gamma0*amu*(1-betaz**2)**0.5

beta_x(x)=betaxy*sin(omega*x)
beta_y(x)=-betaxy*cos(omega*x)
beta_z=0
h(x)=((cos(omegap * x/2))**2-((omega-omega0)**2-bomega**2)/(omegap**2) *(sin(omegap * x/2))**2)*cos(omega *x)
j(x)=((cos(omegap * x/2))**2-((omega-omega0)**2-bomega**2)/(omegap**2) *(sin(omegap * x/2))**2)*sin(omega *x) 
hh(x)=- (2/omegap*cos(omegap*x/2)*sin(omegap*x/2)*(omega-omega0))*sin(omega*x)
jj(x)=  (2/omegap*cos(omegap*x/2)*sin(omegap*x/2)*(omega-omega0))*cos(omega*x)

sx(x)=h(x)+hh(x)
sy(x)=j(x)+jj(x)
sz(x) = 2*(omega-omega0)*bomega/omegap**2 * (sin(omegap*x/2))**2

sdotb(x) = sx(x)*beta_x(x)+ sy(x)*beta_y(x)+ sz(x)*betaz
f(x) = 1+cos(omegap*x)
g(x) = 1-cos(omegap*x)
ff(x) = f(x) - ((omega-omega0)**2-bomega**2)/omegap**2*g(x)
gg(x)= betaz*(omega-omega0)*bomega/omegap**2 * g(x)
bdotsfg(x) = ff(x)+gg(x)
bdots(x) = (1+cos(omegap*x) - 1/omegap**2*((omega-omega0)**2 -bomega**2)*(1-cos(omegap*x))*beta \
+betaz*(omega-omega0)*bomega/omegap**2)*(1-cos(omegap*x))