import tellurium as te
import roadrunner
import numpy as np
import lmfit; import math
import copy; import random
import pylab as plt;
import scipy

r = te.loada("""
     $Xo -> S1; k1*Xo;
      S1 -> S2; k2*S1;
      S2 -> S3; k3*S2;
      S1 -> S3; k4*S1;
      S3 -> $X1; k5*S3;
      
     Xo = 10; 
     S1 = 0; S2 = 0; S3 = 0;
     k1 = 0.2;  k2 = 0.26;
     k3 = 0.78; k4 = 0.45;
     k5 = 0.33;
""")

toFit = ['k1', 'k2', 'k3', 'k4', 'k5']; nParameters = len (toFit)

timeToSimulate = 15
nDataPoints = 24
# Create the experimental data
m = r.simulate (0, timeToSimulate, nDataPoints)
r.plot()
groundTruth = copy.deepcopy (m)

# Change this index to use a different variable
SIndexList = [1, 2, 3]
x_data = m[:,0]; 

y_data = []
for i in range (len(SIndexList)):
    y_data.append (m[:,SIndexList[i]])
y_noise = np.empty([nDataPoints])

for k in range (len (SIndexList)):
   for i in range (0, len (y_data[k])):
       y_noise[i] = 0.1 # standard deviation of noise
       y_data[k][i] = y_data[k][i] + np.random.normal (0, y_noise[i]); # Add noise

# Compute the simulation at the current parameter values
# Return the variable indicated by SIndex
def my_ls_func(p, SIndex):
    r.reset()  
    pp = p.valuesdict()
    for i in range(0, nParameters):
       r.model[toFit[i]] = pp[toFit[i]]
    m = r.simulate (0, timeToSimulate, nDataPoints)
    return m[:,SIndex]

# Compute the residuals between objective and experimental data
def residuals(p):
    y1 = (y_data[0] - my_ls_func (p, SIndexList[0]));#/y_noise    
    y1 = np.concatenate ((y1, ))
    for k in range (1, len (SIndexList)):
        y1 = np.concatenate ((y1, (y_data[k] - my_ls_func (p, SIndexList[k]))))
    return y1
   
def unWeightedResiduals(p):
    y1 = (y_data[0] - my_ls_func (p, SIndexList[0]))
    return y1

# Set up the parameters that we will fit
params = lmfit.Parameters()

params.add('k1', value=1, min=0, max=10)
params.add('k2', value=1, min=0, max=10)
params.add('k3', value=1, min=0, max=10)
params.add('k4', value=1, min=0, max=10)
params.add('k5', value=1, min=0, max=10)


minimizer = lmfit.Minimizer(residuals, params)
result = minimizer.minimize(method='differential_evolution')  #'leastsqr')
result = minimizer.minimize()
print ("Results:")
lmfit.report_fit(result.params, min_correl=0.5)
ci = lmfit.conf_interval(minimizer, result)
lmfit.printfuncs.report_ci(ci)

# Assign fitted parameters to the model
r.reset()
for i in range(0, nParameters):
   r.model[toFit[i]] = result.params[toFit[i]].value
m = r.simulate (0, timeToSimulate, 100)

# Set up some convenient font sizes
plt.rcParams.update({'axes.titlesize': 16})
plt.rcParams.update({'axes.labelsize': 14})
plt.rcParams.update({'xtick.labelsize': 13})
plt.rcParams.update({'ytick.labelsize': 13})

plt.figure (figsize=(7,5))
# Plot experimental data

ldata, = plt.plot (x_data, y_data[0], 'dm', markersize=8)
for k in range (1, len (SIndexList)):
    ldata2, = plt.plot (x_data, y_data[k], 'dm', markersize=8)

# Plot the fitted lines for S1, S2 and S3
# Retrive lfit to use in the legend
lfit, = plt.plot (m[:,0], m[:,1], '-b', linewidth=2)
plt.plot (m[:,0], m[:,2], '-g', linewidth=2)
plt.plot (m[:,0], m[:,3], '-r', linewidth=2)

# Plot the residuals
resids = unWeightedResiduals(result.params)
lresids, = plt.plot (x_data, resids, 'bo', markersize=6)
plt.vlines(x_data, [0], resids, color='r', linewidth=2)

theResiduals = copy.deepcopy (resids)
finalFittedData = copy.deepcopy (y_data[0])
#originalYData = copy.deepcopy (y_data[1])

plt.tick_params(axis='both', which='major', labelsize=16)
plt.xlabel('Time')
plt.ylabel("Concentration", fontsize=16)
plt.legend([ldata, lfit, lresids],['Data', 'Best fit', 'Residuals'], loc=0, fontsize=10)
plt.axhline (y=0, color='k')
plt.savefig('fittedCurves.pdf')
plt.show()

print ("Percentage Parameter Errors:")
print ("k1 = ", 100*math.fabs ((0.2 - result.params['k1'].value)/result.params['k1'].value))
print ("k2 = ", 100*math.fabs (0.26 - result.params['k2'].value)/result.params['k2'].value)
print ("k3 = ", 100*math.fabs (0.78 - result.params['k3'].value)/result.params['k3'].value)
print ("k4 = ", 100*math.fabs (0.45 - result.params['k4'].value)/result.params['k4'].value)
print ("k5 = ", 100*math.fabs ( 0.33 - result.params['k5'].value)/result.params['k5'].value)

#print result.params
#print 100*math.fabs ((0.2 - result.params['k1'].value)/result.params['k1'].value)
#print 100*math.fabs (0.26 - result.params['k2'].value)/result.params['k2'].value
#print 100*math.fabs (0.78 - result.params['k3'].value)/result.params['k3'].value
#print 100*math.fabs (0.45 - result.params['k4'].value)/result.params['k4'].value
#print 100*math.fabs ( 0.33 - result.params['k5'].value)/result.params['k5'].value
