# Converts epn (energy in MeV) to rigin (rigidity in MV)
from math import sqrt
import numpy as np
def energytorigidity(epn):
  ptype = str(raw_input("Please enter particle type (eg: proton, helium): "))
  if ptype == 'proton':
    pamu = 1.0081451
    epamu = 931.141
    anuc = 1.
    zcharg = 1.
  if ptype == 'helium':
    pamu = 4.001506
    zcharg = 2.
    epamu = 3727.379/4 #/2 da per amu?!
    anuc = 4.
  if ptype == 'electron':
    pamu = 0.0005485799
    zcharg = 1.
    epamu = 931.141
    anuc = 1.
  rsmspn = (pamu/anuc)*epamu
  gmarg =(epn/rsmspn)+1.0
  rigin=sqrt(gmarg**2-1.0)*(rsmspn*anuc)/zcharg
  return rigin

# for manual testing:
# print energytorigidity(140)




# Converts rigin (rigidity in MV) to epn (Energy in MEV)
from math import sqrt
def rigiditytoenergy(rigin):
  ptype = str(raw_input("Please enter particle type (eg: proton, helium): "))
  if ptype == 'proton':
    pamu = 1.0081451
    epamu = 931.141
    anuc = 1.
    zcharg = 1.
  if ptype == 'helium':
    pamu = 4.001506
    zcharg = 2.
    epamu = 3727.379/2 #/2 da per amu?!
    anuc = 2.
  if ptype == 'electron':
    pamu = 0.0005485799
    zcharg = 1.
    epamu = 931.141
    anuc = 1.
  rsmspn = (pamu/anuc)*epamu
  gmarg = sqrt(((rigin*zcharg)/(rsmspn*anuc))**2+1.0)
  epn = (gmarg-1.0)*rsmspn  
  return epn
# for manual testing:
# print rigiditytoenergy(500)
"""
energy in MeV
pamu is rest mass in physical atomic mass units
epamu is mass-energy conversion = 931.141 MeV per amu
epn is kinetic energy per nucleon
anuc is atomic number
zcharg is charge
rigidity in mv
rsmspn is rest mass per nucleon in MeV"""


print "Select Calculation:"
print "(1) energy to rigidity"
print "(2) rigidity to energy"


mode=raw_input("Mode: ")
if mode=="1":
  e=np.float64(raw_input("kinetic energy in MeV: "))
  rig=energytorigidity(e)
  print rig

else:
  r=float(raw_input("rigidity in MV: "))
  e=rigiditytoenergy(r)
  print e