import random
from matplotlib import pyplot as plt
import numpy as np
import math

# number of simulation runs
n = 10000

# time bound in the property
t = 100

# rates for different methods (you can add any vector of
# positive real numbers of size 6)
ssa = [1, 1, 0.1, 1, 1, 0.1]
# dwssa = [1.000, 1.003, 0.320, 1.003, 0.993, 3.008]
wssa = [1, 1, 0.05, 1, 1, 0.2]
opt = [1, 1, 0.06, 1, 1, 0.16]
new_rates = [0.47866, 0.44150, 0.01025, 1.53989, 0.06135, 14.65988]

# the effect of reactions on the state of system
R1 = np.array([-1, -1, 1, 0, 0, 0])
R2 = np.array([1, 1, -1, 0, 0, 0])
R3 = np.array([1, 0, -1, 0, 1, 0])
R4 = np.array([0, 0, 0, -1, -1, 1])
R5 = np.array([0, 0, 0, 1, 1, -1])
R6 = np.array([0, 1, 0, 1, 0, -1])

# original rate of transitions and initial state of the system
k = [1, 1, 0.1, 1, 1, 0.1]
molecules_initial_count = [1, 50, 0, 1, 50, 0]

# updated rates for importance sampling (select the one you want to simulate with)
updated_k = wssa

# counter for number(in this case weight) of successful runs
count = 0

progress = np.zeros(n)

for i in range(n):
    time = 0
    x = np.array(molecules_initial_count)
    w = 1
    a = np.array([k[0] * x[0] * x[1],
                  k[1] * x[2],
                  k[2] * x[2],
                  k[3] * x[3] * x[4],
                  k[4] * x[5],
                  k[5] * x[5]])
    a0 = a.sum()

    b = np.array([updated_k[0] * x[0] * x[1],
                  updated_k[1] * x[2],
                  updated_k[2] * x[2],
                  updated_k[3] * x[3] * x[4],
                  updated_k[4] * x[5],
                  updated_k[5] * x[5]])
    b0 = b.sum()

    while time <= t:
        if x[4] <= 40:
            count = count + w
            break
        r1 = random.uniform(0, 1)
        r2 = random.uniform(0, b0)

        tau = float(1.0 / float(a0)) * math.log(float(1.0 / r1))
        j = 0
        temp_sum = 0

        while temp_sum <= r2:
            temp_sum = temp_sum + b[j]
            j = (j + 1)

        j = j - 1
        w = w * (a[j] / b[j]) * (b0 / a0)
        time = time + tau

        if j == 0:
            x = x + R1
        elif j == 1:
            x = x + R2
        elif j == 2:
            x = x + R3
        elif j == 3:
            x = x + R4
        elif j == 4:
            x = x + R5
        elif j == 5:
            x = x + R6
        else:
            print("crucial error!!")

        a = np.array([k[0] * x[0] * x[1],
                      k[1] * x[2],
                      k[2] * x[2],
                      k[3] * x[3] * x[4],
                      k[4] * x[5],
                      k[5] * x[5]])
        a0 = a.sum()
        b = np.array([updated_k[0] * x[0] * x[1],
                      updated_k[1] * x[2],
                      updated_k[2] * x[2],
                      updated_k[3] * x[3] * x[4],
                      updated_k[4] * x[5],
                      updated_k[5] * x[5]])
        b0 = b.sum()

    progress[i] = float(count / float(i + 1))

print("probability: " + str(float(count/float(n))))
plt.xscale("log")
y_axis = np.array(progress)
np.save("opt1", y_axis)
x_axis = np.arange(n)
plt.plot(x_axis, progress)
plt.xlabel("# simulation runs")
plt.ylabel("probability")
plt.show()
# plt.savefig("dwssa2.png")