import math

def get_lambda(thalf):
    """
    calculates radioactive decay constant from decay half-life
    """

    return math.log(2)/thalf


def decay(t1, t2, lamb):
    """
    calculates radioactive decay
    
    t1, t2 - time in second since the end of chain reaction
    lamb - radioactive decay constant
    """
    
    return (math.exp(-lamb*t1) - math.exp(-lamb*t2))/(lamb*(t2-t1))


import numpy
%matplotlib inline
from matplotlib import pyplot as plt

a = numpy.ones(10)

plt.step(range(10), a)
plt.xlabel("Time step")
plt.ylabel("Activity concentration")

thalf = 1 * 3600.
lamb = get_lambda(thalf)
timestep = 1 *3600.  # time step in seconds
a_dec = numpy.zeros(a.shape[0])

for t in range(a_dec.shape[0]):
    a_dec[t] =  a[t] * decay(t*timestep, (t+1)*timestep, lamb)
    print a[t], a_dec[t]

#plt.plot(range(1,11), a)
plt.plot(range(1,11), a_dec)

plt.step(range(2,12), a, label="av. conc")
plt.step(range(2,12), a_dec, label="av. conc. decay")  # we shist the support to have the value valid the following interval, not preceeding

plt.xlabel("Time step")
plt.ylabel("Activity concentration")
plt.ylim(-0.1,1.1)
plt.xlim(0,11)
plt.grid()
plt.legend(loc="best")

def decay_2(t1, t2, lamb1, lamb2, branch12):
    """
    calculates decay of species A1 into species A2 which also decays with lambda2 into a stable isotope
    
    t1, t2 - time in second since the end of chain reaction
    lamb1, lamb2 - radioactive decay constants of A1 and A2, respectively
    branch12 - branching ratio of decay A1 to A2
    """
    
    ret = (math.exp(-lamb1*t1) - math.exp(-lamb1*t2))/lamb1 + (math.exp(-lamb2*t2) - math.exp(-lamb2*t1))/lamb2
    ret = ret * branch12 * lamb2 / ((lamb2-lamb1)*(t2-t1)) 
    return ret

thalf1 = 66 #* 3600. # Mo-99
thalf2 = 6 #* 3600.  # Tc-99m
lamb1 = get_lambda(thalf1)
lamb2 = get_lambda(thalf2)

print lamb1, lamb2 

branch = 1  # 0.1 goes directly to Tc-99

timestep = 0.1#*3600.  # time step in seconds
dimx = 1000  # number of time steps

a1 = numpy.ones(dimx)     # how much parent nuclide is present?
a2 = numpy.ones(dimx)*0.  # how much A2 is already present?
a1_dec = numpy.zeros(dimx)
a2_dec = numpy.zeros(dimx)

for t in range(dimx):
    a2_dec[t] =  a1[t] * decay_2(t*timestep, (t+1)*timestep, lamb1, lamb2, branch)  # decay of A1 to A2 and to stable
    a2_dec[t] +=  a2[t] * decay(t*timestep, (t+1)*timestep, lamb2)  # decay of already present A2
    
    a1_dec[t] =  a1[t] * decay(t*timestep, (t+1)*timestep, lamb1)  # decay of already present A2


plt.step(range(2,dimx+2), a1_dec, label="parent")
plt.step(range(2,dimx+2), a2_dec, label="daughter")  # we shist the support to have the value valid the following interval, not preceeding

plt.xlabel("Time step")
plt.ylabel("Activity concentration")
#plt.ylim(-0.1,1.1)
plt.xlim(0,dimx+1)
plt.grid()
plt.legend(loc="best")
#plt.yscale('log')

thalf1 = 12.77  # days
thalf2 = 1.68   # days

lamb1 = get_lambda(thalf1)  # Ba-140
lamb2 = get_lambda(thalf2)  # La-140

print lamb1, lamb2 

branch = 1 

timestep = 0.1#*3600.  # time step in seconds
dimx = 300  # number of time steps

a1 = numpy.ones(dimx)    # how much parent nuclide is present?
a2 = numpy.ones(dimx)*0.  # how much A2 is already present?
a1_dec = numpy.zeros(dimx)
a2_dec = numpy.zeros(dimx)

for t in range(dimx):
    a2_dec[t] =  a1[t] * decay_2(t*timestep, (t+1)*timestep, lamb1, lamb2, branch)  # decay of A1 to A2 and to stable
    a2_dec[t] +=  a2[t] * decay(t*timestep, (t+1)*timestep, lamb2)  # decay of already present A2
    
    a1_dec[t] =  a1[t] * decay(t*timestep, (t+1)*timestep, lamb1)  # decay of already present A2


plt.step(range(2,dimx+2), a1_dec, label="parent")
plt.step(range(2,dimx+2), a2_dec, label="daughter")  # we shist the support to have the value valid the following interval, not preceeding

plt.xlabel("Time step")
plt.ylabel("Activity concentration")
#plt.ylim(-0.1,1.1)
plt.xlim(0,dimx+1)
plt.grid()
plt.legend(loc="best")
#plt.yscale('log')

rad_data = {
'Cs-125': {'halflife': 2700.0, 'cloud': 3.01e-14, 'ground': 6.85e-16},
'Cs-126': {'halflife': 98.4, 'cloud': 4.96e-14, 'ground': 1.18e-15},
'Cs-127': {'halflife': 22500.0, 'cloud': 1.78e-14, 'ground': 3.95e-16},
'Cs-128': {'halflife': 234.0, 'cloud': 4.06e-14, 'ground': 9.54e-16},
'Cs-129': {'halflife': 115776.0, 'cloud': 1.13e-14, 'ground': 2.62e-16},
'Cs-130': {'halflife': 1792.8, 'cloud': 2.3e-14, 'ground': 5.41e-16},
'Cs-131': {'halflife': 837216.0, 'cloud': 2.38e-16, 'ground': 1.79e-17},
'Cs-132': {'halflife': 559872.0, 'cloud': 3.11e-14, 'ground': 6.69e-16},
'Cs-134': {'halflife': 64964160.0, 'cloud': 7.06e-14, 'ground': 1.48e-15},
'Cs-134m': {'halflife': 10440.0, 'cloud': 7.95e-16, 'ground': 2.25e-17},
'Cs-135': {'halflife': 7.25328e+13, 'cloud': 9.5e-18, 'ground': 2.69e-20},
'Cs-135m': {'halflife': 3178.8, 'cloud': 7.25e-14, 'ground': 1.51e-15},
'Cs-136': {'halflife': 1131840.0, 'cloud': 9.94e-14, 'ground': 2.03e-15},
'Cs-137': {'halflife': 946080000.0, 'cloud': 9.28e-17, 'ground': 2.99e-18},
'Cs-138': {'halflife': 1929.6, 'cloud': 1.15e-13, 'ground': 2.26e-15},
'I-120': {'halflife': 4860.0, 'cloud': 1.31e-13, 'ground': 2.62e-15},
'I-120m': {'halflife': 3178.8, 'cloud': 2.49e-13, 'ground': 5.01e-15},
'I-121': {'halflife': 7632.0, 'cloud': 1.78e-14, 'ground': 3.96e-16},
'I-122': {'halflife': 217.2, 'cloud': 4.31e-14, 'ground': 1.02e-15},
'I-123': {'halflife': 47520.0, 'cloud': 6.49e-15, 'ground': 1.53e-16},
'I-124': {'halflife': 361152.0, 'cloud': 5.04e-14, 'ground': 1.04e-15},
'I-125': {'halflife': 5192640.0, 'cloud': 3.73e-16, 'ground': 3.14e-17},
'I-126': {'halflife': 1123200.0, 'cloud': 2.01e-14, 'ground': 4.42e-16},
'I-128': {'halflife': 1497.6, 'cloud': 4.33e-15, 'ground': 1.71e-16},
'I-129': {'halflife': 4.951152e+14, 'cloud': 2.81e-16, 'ground': 1.95e-17},
'I-130': {'halflife': 44640.0, 'cloud': 9.67e-14, 'ground': 2.05e-15},
'I-131': {'halflife': 694656.0, 'cloud': 1.69e-14, 'ground': 3.64e-16},
'I-132': {'halflife': 8280.0, 'cloud': 1.05e-13, 'ground': 2.2e-15},
'I-132m': {'halflife': 5004.0, 'cloud': 1.42e-14, 'ground': 3.11e-16},
'I-133': {'halflife': 74880.0, 'cloud': 2.76e-14, 'ground': 6.17e-16},
'I-134': {'halflife': 3153.6, 'cloud': 1.22e-13, 'ground': 2.53e-15},
'I-135': {'halflife': 23796.0, 'cloud': 7.54e-14, 'ground': 1.47e-15},
'Ar-37': {'halflife': 3024000.0, 'cloud': 4.75e-20, 'ground': -999},
'Ar-39': {'halflife': 8483184000.0, 'cloud': 1.27e-16, 'ground': -999},
'Ar-41': {'halflife': 6588.0, 'cloud': 6.13e-14, 'ground': -999},
'Kr-74': {'halflife': 690.0, 'cloud': 5.21e-14, 'ground': -999},
'Kr-76': {'halflife': 53280.0, 'cloud': 1.85e-14, 'ground': -999},
'Kr-77': {'halflife': 4482.0, 'cloud': 4.51e-14, 'ground': -999},
'Kr-79': {'halflife': 126144.0, 'cloud': 1.12e-14, 'ground': -999},
'Kr-81': {'halflife': 6.62256e+12, 'cloud': 2.43e-16, 'ground': -999},
'Kr-83m': {'halflife': 6588.0, 'cloud': 2.43e-18, 'ground': -999},
'Kr-85': {'halflife': 337435200.0, 'cloud': 2.55e-16, 'ground': -999},
'Kr-85m': {'halflife': 16128.0, 'cloud': 6.83e-15, 'ground': -999},
'Kr-87': {'halflife': 4572.0, 'cloud': 3.94e-14, 'ground': -999},
'Kr-88': {'halflife': 10224.0, 'cloud': 9.72e-14, 'ground': -999},
'Xe-120': {'halflife': 2400.0, 'cloud': 1.74e-14, 'ground': -999},
'Xe-121': {'halflife': 2406.0, 'cloud': 8.68e-14, 'ground': -999},
'Xe-122': {'halflife': 72360.0, 'cloud': 2.2e-15, 'ground': -999},
'Xe-123': {'halflife': 7488.0, 'cloud': 2.78e-14, 'ground': -999},
'Xe-125': {'halflife': 61200.0, 'cloud': 1.08e-14, 'ground': -999},
'Xe-127': {'halflife': 3144960.0, 'cloud': 1.12e-14, 'ground': -999},
'Xe-129m': {'halflife': 691200.0, 'cloud': 9.38e-16, 'ground': -999},
'Xe-131m': {'halflife': 1028160.0, 'cloud': 3.7e-16, 'ground': -999},
'Xe-133m': {'halflife': 189216.0, 'cloud': 1.27e-15, 'ground': -999},
'Xe-133': {'halflife': 452736.0, 'cloud': 1.39e-15, 'ground': -999},
'Xe-135m': {'halflife': 918.0, 'cloud': 1.85e-14, 'ground': -999},
'Xe-135': {'halflife': 32760.0, 'cloud': 1.11e-14, 'ground': -999},
'Xe-138': {'halflife': 852.0, 'cloud': 5.44e-14, 'ground': -999},
}

SOURCE_TERM = {
    "ts" : 3600,  # time step in seconds:
    "release_start" : "20150428 12:00",
    "release_end"   : "20150428 14:00",  # e.g. a 2 hours of release, i.e. two time steps of length 1 hour
    "sources" : [  # an array with list of sources acting at that time (list of lists)
                    # TIME STEP 1
                    [{"lat0":lat0,  # source 1 for time step 1
                      "lon0":lon0,
                      "lat1":lat1,
                      "lon1":lon1,
                      "h0":  0,  # bottom of layer in meters
                      "h1": 50,  # top of layer in meters
                      "inventory": {"Cs-137": (1e10,1), "I-131": (1e11,1), "Xe-133": (1e10,0)}  # (release, species)
                     },
                     {"lat0":lat0,  # source 2 for time step 1
                      "lon0":lon0,
                      "lat1":lat1,
                      "lon1":lon1,
                      "h0":  50,  # bottom of layer in meters
                      "h1": 200,  # top of layer in meters
                      "inventory": {"Cs-137": (1e11,1), "I-131": (1e12,1) , "Xe-133": (1e13,0)}
                     }],
                    # TIME STEP 2
                    [{"lat0":lat0,  # source 1 for time step 2
                      "lon0":lon0,
                      "lat1":lat1,
                      "lon1":lon1,
                      "h0":  0,  # bottom of layer in meters
                      "h1": 50,  # top of layer in meters
                      "inventory": {"Cs-137": (1e9,1), "I-131": (1e10,1), "Xe-133": (0.,0)}
                     },
                     {"lat0":lat0,  # source 2 for time step 2
                      "lon0":lon0,
                      "lat1":lat1,
                      "lon1":lon1,
                      "h0":  50,  # bottom of layer in meters
                      "h1": 200,  # top of layer in meters
                      "inventory": {"Cs-137": (1e10,1), "I-131": (1e11,1), "Xe-133": (0.,0)}
                     }]
                ]
    }