import numpy as np 
import scipy.stats as st 
import matplotlib.pyplot as plt 

np.random.seed(42)

#defines constant and preallocates
c = np.array([1,5,15,3./4])

x, y, z = 400, 800, 0
xx = np.array([])
yy = np.array([])
zz = np.array([])
tt = np.array([])
t1 = 600000000
t = 0
i = 0
#simmulates dynamics
while t < t1:
    v = np.array([x, x*(x - 1.)/2., y, y])    
    a = c * v
    lam = np.sum(a)
    #samples
    p = a / np.sum(a)
    rand = st.uniform.rvs()
    r = np.where(np.cumsum(p) >= rand)[0].min()
    t = t + st.expon.rvs(scale=lam)
    print(t, lam)
    #changes dynamics according to r
    if r == 0:
        x = x - 1
    elif r == 1:
        x, y = x - 2, y + 1
    elif r == 2:
        x, y = x + 2, y - 1
    elif r == 3:
        y, z = y - 1, z + 1

    xx = np.append(xx,x)
    yy = np.append(yy,y)
    zz = np.append(zz,z)
    tt = np.append(tt,t)
    i = i + 1
#plots
plt.plot(tt, xx, linewidth = 1., label = r'$N^1$')
plt.plot(tt, yy, linewidth = 1., label = r'$N^2$')
plt.plot(tt, zz, linewidth = 1., label = r'$N^3$')
plt.legend()
plt.show()