from ODESolver import *
import numpy as np
import matplotlib.pyplot as plt

class Decay:
    def __init__(self,tau_a,tau_b):
        self.tau_a = tau_a
        self.tau_b = tau_b

    def __call__(self, u, t):
        u_a = u[0]
        u_b = u[1]
        tau_a = self.tau_a
        tau_b = self.tau_b

        du_a = u_b/tau_b - u_a/tau_a
        du_b = u_a/tau_a - u_b/tau_b
        return du_a, du_b

    def terminate(self,u,t,k):
        #return True if u_a/u_b ~ tau_a/tau_b
        u_k = u[k,:] #u at current time step
        u_a = u_k[0]
        u_b = u_k[1]
        tol = 1e-3
        return abs(u_a/u_b - self.tau_a/self.tau_b) < tol


problem = Decay(tau_a = 8, tau_b = 40)
solver = RungeKutta4(problem)

dt = 10.0/60
T_stop = 150
n = round(T_stop/dt)
time = np.linspace(0,T_stop,n+1)

U0 = (1,1)
solver.set_initial_condition(U0)

u,t = solver.solve(time,terminate=problem.terminate)

plt.plot(t,u[:,0],label='$u_a$')
plt.plot(t,u[:,1],label='$u_b$')
plt.legend()
plt.show()

"""
Terminal> python radioactive_decay2.py

"""