# import modules
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import math
import numpy as np
import scipy.io
#define where the output will be saved. Final part of string is filename
filepath='/home/gln09/Desktop/GN_Ubiquitination/unified/unified_sampling'
# Define the terms of the equations
# X = [protein]
# t = time

# knumber = rate of that reaction
kon = []
k = []
koff = []
for n in range(0,6):
    koff.append(1)
for i in range(1, 6):
    kon.append(0)
    k.append(0)
for n in range(0,5):
    kon[n]=1
    k[n]=1

ks=[]
ks=np.arange(0,26,1).tolist()

# L = [E3 ligase]
L = 1
parameters = (k,kon,koff,L)

iX0 = 5
iX0L = 0
iXuL = 0
iXu = 0
iXuuL = 0
iXuu = 0
iXuuuL = 0
iXuuu = 0
iXuuuuL = 0
iXuuuu = 0
iXuuuuuL = 0
iXuuuuu = 0

y = []
for n in range(1, 13):
    y.append(0)
y[0] = iX0
y[1] = iX0L
y[2] = iXuL
y[3] = iXu
y[4] = iXuuL
y[5] = iXuu
y[6] = iXuuuL
y[7] = iXuuu
y[8] = iXuuuuL
y[9] = iXuuuu
y[10] = iXuuuuuL
y[11] = iXuuuuu

ys=[]
zs=[]

for g in range(0,len(ks)):
    koff[0]=ks[g]
    parameters=(k,kon,koff,L)
    zs.append(ks[g])
# The derivative function:
    def f(y, t, k, kon, koff, L):
        return (
               ((koff[0]*y[1])-(kon[0]*y[0]*L)),#dX0/dt
            ((kon[0]*y[0]*L)-(koff[0]*y[1])-(k[0]*y[1])),#dX0L/dt
            ((k[0]*y[1])+(kon[1]*y[3]*L)-(k[1]*y[2])-(koff[1]*y[2])),#dXuL/dt
            ((koff[1]*y[2])-(kon[1]*y[3]*L)),#dXu/dt
            ((k[1]*y[2])+(kon[2]*y[5]*L)-(k[2]*y[4])-(koff[2]*y[4])),#dXuuL/dt
            ((koff[2]*y[4])-(kon[2]*y[5]*L)),#dXuu/dt
            ((k[2]*y[4])+(kon[3]*y[7]*L)-(k[3]*y[6])-(koff[3]*y[6])),#dXuuuL/dt
            ((koff[3]*y[6])-(kon[3]*y[7]*L)),#dXuuu/dt
            ((k[3]*y[6])+(kon[4]*y[9]*L)-(k[4]*y[8])-(koff[4]*y[8])),#dXuuuuL/dt
            ((koff[4]*y[8])-(kon[4]*y[9]*L)),#dXuuuu/dt
            ((k[4]*y[8])-(koff[5]*y[10])),#dXuuuuuL/dt
            ((koff[5]*y[10]))#dXuuuuu/dt
            )
    iC = []
    for i in range(0,len(y)):
        iC.append(0)
    for i in range(0,len(y)):
        iC[i]=y[i]
    t = range(0, 41, 1)
    res = odeint(f,iC,t,parameters)
    transres = np.transpose(res)
    toappend = []
    toappend.append(transres[0][-1])
    toappend.append(transres[3][-1])
    toappend.append(transres[5][-1])
    toappend.append(transres[7][-1])
    toappend.append(transres[9][-1])
    toappend.append(transres[11][-1])
    ys.append(toappend)
    toappend=[]
#output to MATLAB    
scipy.io.savemat(filepath, mdict={'ys':ys})

#Plot
'''
#initial = sum(res[0])
#print "initial [protein]", initial
#final = sum(res[-1])
#print "final [protein]", final
x = zs
y = ys

plt.plot(x,y)
plt.title(r"Fraction of X$_0$ Present as k$_{off}$ Varies in a Unified Model of Polyubiquitination")
plt.ylabel(r"Fraction of X$_0$")
plt.xlabel(r"k$_{off}$")
plt.ylim(0,1)
plt.show()
'''
