from symfit import variables, Parameter, Fit, D, ODEModel
import numpy as np
import matplotlib.pyplot as plt

tdata = np.array([0, 2, 4, 8, 12, 24, 48])
concentration = np.array([1, 1.15, 1.6, 1.8, 1.55, 1.2, 1])
A, t, P, bP, bI, RB, RP, RN, bRN, G = variables('A, t, P, bP, bI, RB, RP, RN, bRN, G')

f_P = Parameter('f_P', 10, min=0)
c_N = Parameter('c_N', 3, min=0)
c_bN = Parameter('c_bN', 0.6, min=0)
d_N = Parameter('d_N', 0.7, min=0)

# supply rates
s_P = 1.5
s_B = 0.5
s_A = 0.5
s_G = 0.2

# degradation/inactivation rates
d_P = 1
d_bP = 0.5
d_bI = 0.5
d_R = 0.5
d_bN = 0.5
d_A = 0.5
d_G = 0.5

# reaction/binding rates
c_P = 0.8
c_I = 0.5
c_B = 4
c_A = 0.8

# proliferation rate
p_G = 2.9       

# mortality rate
m_G = 2      

model = ODEModel({
    D(A, t): s_A + c_A * bRN - d_A * A,
    D(P, t): s_P - c_P * P * G - d_P * P,
    D(bP, t): c_P * P * G - d_bP * bP,
    D(bI, t): c_I * (1-bI) * bP - d_bI * bI,    
    D(RB, t): s_B - c_B * RB * bI - d_R * RB,
    D(RP, t): c_B * RB * bI - f_P * RP - d_R * RP,
    D(RN, t): f_P * RP - c_N * RN + c_bN * bRN - d_N * RN,
    D(bRN, t): c_N * RN - c_bN * bRN - d_bN * bRN,  
    D(G, t): p_G * G - m_G * G * A - d_G * G
    },
    initial={t: tdata[0], A: concentration[0], P: 1, bP: 0, bI: 0, RB: 1, RP: 0, RN: 0, 
             bRN: 0, G: 1}
)

fit = Fit(model, t=tdata, A=concentration, P=None, bP=None, bI=None, RB=None, RP=None, RN=None, 
          bRN=None, G=None)
fit_result = fit.execute()

print(fit_result)

taxis = np.linspace(0, 50)
A_fit, P_fit, bP_fit, bI_fit, RB_fit, RP_fit, RN_fit, bRN_fit, G_fit, = model(t=taxis, **fit_result.params)
plt.scatter(tdata, concentration)
plt.plot(taxis, A_fit, label='[A]')
plt.plot(taxis, P_fit, label='[P]')
plt.plot(taxis, bP_fit, label='[bP]')
plt.plot(taxis, bI_fit, label='[bI]')
plt.plot(taxis, RB_fit, label='[RB]')
plt.plot(taxis, RP_fit, label='[RP]')
plt.plot(taxis, RN_fit, label='[RN]')
plt.plot(taxis, bRN_fit, label='[bRN]')
plt.plot(taxis, G_fit, label='[G]')
plt.xlabel('Hours')
plt.ylabel('[AMP]')
plt.ylim(0, 4)
plt.xlim(0, 50)
plt.legend()
plt.show()