program pion_decay_muon_polarization
use sim_utils
use precision_def

implicit none

real(rp) mpi,mmu,pmu_r, Emu_r,gamma_r, beta_r
real(rp) pmu_r_4vec(1:4), smu_r_4vec(1:4), pmumax_r_4vec(1:4)
real(rp) theta, Epi, ppi, s_long, s_perp,x, p_mag, s_mag, theta_star, phi, tanthetamax
real(rp) gamma, beta
real(rp) sdotp, cosphi
real(rp) pdotp
real(rp) boost(1:4,1:4), p_3vec(1:3), s_3vec(1:3), g(1:4,1:4)
real(rp) A_mu(1:4,1:4), pmu_rest(1:4), smu_rest(1:4)
real(rp) zero/0./, one/1./
real(rp) dt
real(rp) sigma_T,b, sigma_T0
real(rp) betax, betaz,cosphi_ana, phi_ana, theta_ana
real(rp) tan_theta, sperp
real(rp) snux,snuz,pmux,pmuz
real(rp) xx

character*10 ctheta, cepi
integer i,j,k
logical first/.true./

!    call get_command_argument(1, ctheta)
!    read(ctheta,*)theta_star
    call get_command_argument(1, cepi)
    read(cepi,*)Epi
 
mpi= 139.57
mmu=105.658

gamma = epi/mpi
beta = sqrt(1. - 1/gamma**2)
ppi = gamma*beta*mpi
b=mmu/mpi
pmu_r = (mpi**2-mmu**2)/2./mpi
Emu_r = sqrt(pmu_r**2+mmu**2)
gamma_r = Emu_r/mmu
beta_r = pmu_r/Emu_r

!pnu_r = pmu_r
!Enu_r=pnu_r


dt=0.005

do j=-100,100
 theta_star = j*dt*twopi

pmu_rest=(/mmu,zero,zero,zero/)

!pnu = pnu_r*(/gamma*(1-beta*cos(theta_star)),sin(theta_star),0,gamma*(cos(theta_star)-beta))

snux=sin(theta_star); snuz=cos(theta_star)
pmux=pmu_r*sin(theta_star); pmuz=gamma*(beta*Emu_r+pmu_r*cos(theta_star))
sperp = (snux*pmuz-snuz*pmux)/(pmux**2+pmuz**2)**.5

xx = (Emu_r+pmu_r/beta*cos(theta_star))/mpi

tan_theta=pmu_r*sin(theta_star)/gamma/(beta*Emu_r+pmu_r*cos(theta_star))

theta=atan(tan_theta)

smu_rest=(/zero,sin(theta_star),zero,cos(theta_star)/)

pmumax_r_4vec=(/Emu_r,zero,zero,pmu_r/)

betax=beta_r*sin(theta_star)
betaz=beta_r*cos(theta_star)

! A_mu boosts from muon rest frame to pion rest frame
A_mu = 0
A_mu(1,1)=gamma_r
A_mu(1,2)=gamma_r*beta_r*sin(theta_star)
A_mu(1,4)=gamma_r*beta_r*cos(theta_star)
A_mu(2,2) = 1+(gamma_r-1)*sin(theta_star)**2
A_mu(2,4)=(gamma_r-1)*sin(theta_star)*cos(theta_star)
A_mu(3,3)=1
A_mu(4,4)=1+(gamma_r-1)*cos(theta_star)**2
A_mu(2,1)=A_mu(1,2)
A_mu(4,1)=A_mu(1,4)
A_mu(4,2)=A_mu(2,4)

print '(/,a,es12.4)','theta_star= ', theta_star
do i=1,4
print '(4es12.4)',(A_mu(i,k),k=1,4)
end do


! pmu_r_4vec is muon 4 vector in pion rest frame
pmu_r_4vec = matmul(A_mu,pmu_rest)

g=0
g(1,1)=-1.
forall(i=2:4)g(i,i)=1

!print '(a,4es12.4)',' p=',pmu_r_4vec
!print '(a,4es12.4)',' gp=',matmul(g,pmu_r_4vec)

pdotp= -dot_product(pmu_r_4vec,matmul(g,pmu_r_4vec))
!print '(2(a,es12.4))',' pdotp = ', pdotp,' root(pdotp)=',sqrt(pdotp)


!boost boosts from pion rest frame to lab frame
boost = 0
boost(1,1)=gamma
boost(2,2)=1
boost(3,3)=1
boost(4,4)=gamma
boost(1,4)=beta*gamma
boost(4,1)=beta*gamma

!write(13,'(5es12.4)')theta,pmu_r_4vec(2),pmu_r_4vec(4), smu_r_4vec(2), smu_r_4vec(4)

!pmu_r_4vec is muon 4 vector in lab frame
pmu_r_4vec = matmul(boost,pmu_r_4vec)
pmumax_r_4vec=matmul(boost,pmumax_r_4vec)
!the muon is traveling in the direction pmu_r_4vec(2) xhat + pmu_r_4vec(4) zhat

p_3vec(1:3) = pmu_r_4vec(2:4)
s_3vec(1:3) = smu_rest(2:4)
!projection of spin onto muon direction, cos(phi) = p_3vec dot s_3vec
!   phi is the angle between the polarization of the lab frame muon and the rest frame muon

sdotp = dot_product(p_3vec,s_3vec)
s_mag = sqrt(dot_product(s_3vec,s_3vec))
p_mag = sqrt(dot_product(p_3vec,p_3vec))  !muon momentum in lab frame

cosphi = sdotp/s_mag/p_mag
phi = acos(cosphi) *sign(one,theta_star) ! angle between muon rest frame polarization and lab frame muon momentum

!tantheta = p_3vec(1)/p_3vec(3)
theta = atan2(p_3vec(1), p_3vec(3)) !angle of muon trajectory in lab frame with respect to pion direction

s_long = dot_product(p_3vec,s_3vec)/(s_mag*p_mag) !polarization parallel to lab frame muon momentum
s_perp = sqrt(1.-s_long**2)  !polarization perpendicular to lab frame muon momentum
x= p_3vec(3)/ppi             !lab longitudinal muon momentum/lab pion momentum

sigma_T0 = 2*mmu/(mpi**2-mmu**2)*p_mag*sin(theta)

sigma_T = 2*b/x/(1-b**2)*sqrt((1-x)*(x-b**2))*sign(one,theta_star)

cosphi_ana = -(betax**2+gamma*(betaz**2+beta*betaz))/sqrt(betax**2+gamma**2*(betaz**2+beta**2+2*beta*betaz))/beta_r
phi_ana=acos(cosphi_ana)
theta_ana = atan2(-betax/gamma,(beta+betaz))

!if(first)then
!  write(12,'(14a16)')'theta_star','theta', 'x', 's_long', 's_perp', 'p_3vec(3)/p_mag', 'p_3vec(1)/p_mag','phi','s_mag','sigma_T','sigma_T0','cosphi__ana','phi_ana', 'theta_ana'
!  write(11, '(4a12,2a20)') 'theta_star', 'phi', 'dE/E', 'angle xp',' asin(sigma_T)','asin(gamma*beta/beta_r*theta)'
!  first=.false.
!endif
!write(11, '(4es12.4,2es20.4)') theta_star, acos(cosphi)*sign(one,theta), (pmumax_r_4vec(4)-p_mag)/pmumax_r_4vec(4), atan2(p_3vec(1), p_3vec(3)), asin(sigma_T), asin(gamma*beta/beta_r * theta)
!write(12,'(14es16.4)')theta_star, theta, x, s_long, s_perp, p_3vec(3)/p_mag, p_3vec(1)/p_mag, phi, s_mag, sigma_T, sigma_T0, cosphi_ana,phi_ana, theta_ana


write(12,'(8es12.4)')theta_star, theta, sperp, sigma_T, sigma_T0,x,xx, b




end do

tanthetamax= beta_r/gamma/beta/sqrt(1-beta_r**2/beta**2)

print '(2(a,es12.4))',' tanthetamax',tanthetamax, ' theta_max' ,atan(tanthetamax)
print '(2(a,es12.4))',' E_pi=',Epi,' theta=',theta
print '(5(a,es12.4))','mpi=',mpi,' mmu=',mmu,' pmu_r =', pmu_r,' gamma_r=',gamma_r,' beta_r=',beta_r
print '(2(a,es12.4))','beta=', beta,' gamma=', gamma
print '(2(a,es12.4))',' cosphi=',cosphi, ' phi =', acos(cosphi)
print '(3(a,es12.4))',' px=', p_3vec(1),' pz=', p_3vec(3),' angle=',atan2(p_3vec(1), p_3vec(3))
print '(2(a,es12.4))',' Emu=', pmu_r_4vec(1), ' delta E/E_max= ', (pmumax_r_4vec(1)-pmu_r_4vec(1))/pmumax_r_4vec(1) 


end program pion_decay_muon_polarization