import numpy as np
import matplotlib.pyplot as plt

# Given an initial number of thallium-208 atoms NTl and
# their probability of decay p in time h, randomly simulates
# their decay to Pb-208 over a total time tmax
# Every time step h, checks if each remaining Tl atom will
# decay if a random uniform deviate is smaller than p

# Constants
NTl = 1000            # initial number of thallium atoms
NPb = 0               # initial number of lead atoms
t_half = 3.053*60        # Half life of thallium in seconds
h = 1.0               # Size of time-step in seconds
p = 1 - 2**(-h/t_half)   # Probability of decay in time h
tmax = 1000           # Total simulation time in seconds

# Lists of plot points
tpoints = np.arange(0.0,tmax,h)    # Time steps t
Tlpoints = []                   # Tl-208 atoms at time t
Pbpoints = []                   # Pb-208 atoms at time t

# Main loop
for t in tpoints:
    Tlpoints.append(NTl)
    Pbpoints.append(NPb)
    # Calculate the number of atoms that decay
    decay = 0
    for i in range(NTl):
        if np.random.default_rng().uniform() < p:
            decay += 1
    NTl -= decay
    NPb += decay

# Make the graph
plt.plot(tpoints,Tlpoints,label="Tl-208")
plt.plot(tpoints,Pbpoints,label="Pb-208")
plt.xlabel("Time, seconds")
plt.ylabel("Number of atoms")
plt.title("Simulated decay of Tl-308")
plt.legend()
plt.show()