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("""
      J1: ->  y1; p0;
      J2: y1 ->;  y1;
      J3: y1 -> y2; p1*y1*(1+p4*y2^4)
      J4: y2 ->; y2*p6;
      
     p0 = 7;
     p1 = 1
     p4 = 1
     p6 = 4.96
""")

m = r.simulate (0, 10, 100)
r.plot()

toFit = ['p0', 'p1', 'p4', 'p6']; nParameters = len (toFit)

timeToSimulate = 10
# Load the experimental data
m = np.loadtxt('expdata.txt')
groundTruth = copy.deepcopy (m)
nDataPoints = np.shape(m)[0]

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

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

for k in range (len (SIndexList)):
   for i in range (0, len (y_data[k])):
       y_data[k][i] = y_data[k][i]

# 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])); 
    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('p0', value=7, min=0, max=20)
params.add('p1', value=7, min=0, max=50)
params.add('p4', value=7, min=0, max=10)
params.add('p6', value=8, min=0, max=20)

minimizer = lmfit.Minimizer(residuals, params)
#result = minimizer.minimize(method='differential_evolution', init='random', popsize=20, maxiter=3000, recombination=0.7)  #'leastsqr')
result = minimizer.minimize(method='leastsqr') 
#result = minimizer.minimize(method='nelder')  
#result = minimizer.minimize(method='ampgo')  
#result = minimizer.minimize(method='brute') 
 
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], 'o', markersize=6)
for k in range (1, len (SIndexList)):
    ldata2, = plt.plot (x_data, y_data[k], 'H', markersize=6)

# 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)

# 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])

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 "p0 = ", 100*math.fabs ((0.2 - result.params['p0'].value)/result.params['p0'].value)
print "p1 = ", 100*math.fabs (0.26 - result.params['p1'].value)/result.params['p1'].value
print "p4 = ", 100*math.fabs (0.78 - result.params['p4'].value)/result.params['p4'].value
print "p6 = ", 100*math.fabs (0.45 - result.params['p6'].value)/result.params['p6'].value

print result.params

