#  _________________________________________________________________________
#
#  Kipet: Kinetic parameter estimation toolkit
#  Copyright (c) 2016 Eli Lilly.
#  _________________________________________________________________________

# Sample Problem 3 (From Sawall et.al.)
# Basic extimation of kinetic paramenter using pyomo discretization 
#
#         min \sum_i^{nt}\sum_j^{nl}( D_{i,j} -\sum_{k}^{nc} C_k(t_i)*S(l_j))**2/\delta
#              + \sum_i^{nt}\sum_k^{nc}(C_k(t_i)-Z_k(t_i))**2/\sigma_k       
#
#		\frac{dZ_a}{dt} = -k*Z_a	Z_a(0) = 1
#		\frac{dZ_b}{dt} = k*Z_a		Z_b(0) = 0
#
#               C_a(t_i) = Z_a(t_i) + w(t_i)    for all t_i in measurement points
#               D_{i,j} = \sum_{k=0}^{Nc}C_k(t_i)S(l_j) + \xi_{i,j} for all t_i, for all l_j 


from __future__ import print_function
from __future__ import division
from kipet.library.TemplateBuilder import *
from kipet.library.ParameterEstimator import *
import matplotlib.pyplot as plt

from kipet.library.data_tools import *
import inspect
import sys
import os

if __name__ == "__main__":

    with_plots = True
    if len(sys.argv)==2:
        if int(sys.argv[1]):
            with_plots = False
    
    # read 200x500 spectra matrix D_{i,j}
    # this defines the measurement points t_i and l_j as well
    dataDirectory = os.path.abspath(
        os.path.join( os.path.dirname( os.path.abspath( inspect.getfile(
            inspect.currentframe() ) ) ), '..','..','data_sets'))
    filename = os.path.join(dataDirectory,'Dijsawall.txt')
    D_frame = read_spectral_data_from_txt(filename)

    ##########################################################
    
    builder2 = TemplateBuilder()    
    builder2.add_mixture_component({'A':1,'B':0})

    # note the parameter is not fixed
    builder2.add_parameter('k',bounds=(0.0,1))
    builder2.add_spectral_data(D_frame)

    # define explicit system of ODEs
    def rule_odes(m,t):
        exprs = dict()
        exprs['A'] = -m.P['k']*m.Z[t,'A']
        exprs['B'] = m.P['k']*m.Z[t,'A']
        return exprs

    builder2.set_odes_rule(rule_odes)
    
    pyomo_model = builder2.create_pyomo_model(0.0,200.0)
        
    optimizer = ParameterEstimator(pyomo_model)
    
    optimizer.apply_discretization('dae.collocation',nfe=100,ncp=3,scheme='LAGRANGE-RADAU')

    # Provide good initial guess
    
    p_guess = {'k':0.01}
    raw_results = optimizer.run_lsq_given_P('ipopt',p_guess,tee=False)
    
    optimizer.initialize_from_trajectory('Z',raw_results.Z)
    optimizer.initialize_from_trajectory('S',raw_results.S)
    optimizer.initialize_from_trajectory('C',raw_results.C)
    
    sigmas = {'device':1.94554e-6,
              'A':2.45887e-12,
              'B':2.45887e-12}
    solver_options = {}
    solver_options = {'mu_init': 1e-4, 'bound_push':  1e-6}
    results_pyomo = optimizer.run_opt('ipopt',
                                      tee=True,
                                      variances=sigmas,
                                      solver_opts = solver_options)

    print("The estimated parameters are:")
    for k, v in results_pyomo.P.items():
        print(k, v)

    tol =1e-3
    assert(abs(results_pyomo.P['k']-0.01)<tol)
    # display results
    if with_plots:
        results_pyomo.C.plot.line(legend=True)
        plt.xlabel("time (s)")
        plt.ylabel("Concentration (mol/L)")
        plt.title("Concentration Profile")

        results_pyomo.S.plot.line(legend=True)
        plt.xlabel("Wavelength (cm)")
        plt.ylabel("Absorbance (L/(mol cm))")
        plt.title("Absorbance  Profile")

        plt.show()

