#!/usr/bin/env python
# coding: utf-8

# # Fitting kinetic parameters to experimental data
# Author: Björn Dahlgren.
# 
# Let us consider the reaction:
# 
# $$
# Fe^{3+} + SCN^- \rightarrow FeSCN^{2+}
# $$
# 
# the product is strongly coloured and we have experimental data (from a stopped-flow appartus) of the absorbance as function of time after mixing for several replicates. The experiment was performed at 7 different temperatures and for one temperature, 7 different ionic strengths. For each set of conditions the experiment was re-run 7 times (replicates). In this notebook, we will determine the activation enthalpy and entropy through regressions analysis, we will also look at the ionic strength dependence.

# In[ ]:


import bz2, codecs, collections, functools, itertools, json
import numpy as np
import matplotlib.pyplot as plt
import chempy
import chempy.equilibria
from chempy.electrolytes import ionic_strength
from chempy.kinetics.arrhenius import fit_arrhenius_equation
from chempy.printing import number_to_scientific_latex, as_per_substance_html_table
from chempy.properties.water_density_tanaka_2001 import water_density
from chempy.units import rescale, to_unitless, default_units as u
from chempy.util.regression import least_squares, irls, avg_params, plot_fit, plot_least_squares_fit, plot_avg_params
from chempy._solution import QuantityDict
get_ipython().run_line_magic('matplotlib', 'inline')
print(chempy.__version__)


# Experimental conditions, the two solutions which were mixed in 1:1 volume ratio in a stopped flow apparatus:

# In[ ]:


sol1 = QuantityDict(u.mM, {'SCN-': 3*u.mM, 'K+': 3*u.mM, 'Na+': 33*u.mM, 'H+': 50*u.mM, 'ClO4-': (33+50)*u.mM})
sol2 = QuantityDict(u.mM, {'Fe+3': 6*u.mM, 'H+': 50*u.mM, 'ClO4-': (3*6+50)*u.mM})


# In[ ]:


sol = (sol1 + sol2)/2  # 1:1 volume ratio at mixing
sol.quantity_name = 'concentration'
Ibase = ionic_strength(rescale(sol/water_density(293*u.K, units=u), u.molal))
print(Ibase)
sol


# In[ ]:


ionic_strength_keys = 'abcd'
ionic_strengths = dict(zip(ionic_strength_keys, [0, 20, 40, 60]*u.molal*1e-3))
temperature_keys = '16.5 18.5 20.5 22.5 24.5'.split()
T0C = 273.15*u.K
temperatures = {k: T0C + float(k)*u.K for k in temperature_keys}
nrep = 7
indices = lambda k: (ionic_strength_keys.index(k[0]), temperature_keys.index(k[1]))


# We will read the data from a preprocessed file:

# In[ ]:


transform = np.array([[1e-3, 0], [0, 1e-4]])  # converts 1st col: ms -> s and 2nd col to absorbance
_reader = codecs.getreader("utf-8")
_dat = {tuple(k): np.dot(np.array(v), transform) for k, v in json.load(_reader(bz2.BZ2File('specdata.json.bz2')))}
data = collections.defaultdict(list)
for (tI, tT, tR), v in _dat.items():
    k = (tI, tT)  # tokens for ionic strength and temperatures
    data[k].append(v)
assert len(data) == len(ionic_strengths)*len(temperatures) and all(len(serie) == nrep for serie in data.values())


# Let's plot the data:

# In[ ]:


def mk_subplots(nrows=1, subplots_adjust=True, **kwargs):
    fig, axes = plt.subplots(nrows, len(ionic_strengths), figsize=(15,6), **kwargs)
    if subplots_adjust:
        plt.subplots_adjust(hspace=0.001, wspace=0.001)
    return axes

def _set_axes_titles_to_ionic_strength(axes, xlim=None, xlabel=None):
    for tI, ax in zip(ionic_strength_keys, axes):
        ax.set_title(r'$I\ =\ %s$' % number_to_scientific_latex(Ibase + ionic_strengths[tI], fmt=3))
        if xlabel is not None:
            ax.set_xlabel(xlabel)
        if xlim is not None:
            ax.set_xlim(xlim)


# In[ ]:


def plot_series(series):
    axes = mk_subplots(sharey=True, sharex=True)
    colors = 'rgbmk'
    for key in itertools.product(ionic_strengths, temperatures):
        idx_I, idx_T = indices(key)
        for serie in series[key]:
            axes[idx_I].plot(serie[:, 0], serie[:, 1], c=colors[idx_T], alpha=0.15)
    _set_axes_titles_to_ionic_strength(axes, xlim=[0, 3.4], xlabel='Time / s')
    axes[0].set_ylabel('Absorbance')
    for c, tT in zip(reversed(colors), reversed(temperature_keys)):
        axes[0].plot([], [], c=c, label=tT + ' °C')
    axes[0].legend(loc='best')
plot_series(data)


# We see that one data series is off: 16.5 ℃ and 0.0862 molal. Let's ignore that for now and perform the fitting, let's start with a pseudo-first order assumption (poor but simple):

# In[ ]:


def fit_pseudo1(serie, ax=None):
    plateau = np.mean(serie[2*serie.shape[0]//3:, 1])
    y = np.log(np.clip(plateau - serie[:, 1], 1e-6, 1))
    x = serie[:, 0]
    # irls: Iteratively reweighted least squares
    res = irls(x, y, irls.gaussian)
    if ax is not None:
        plot_least_squares_fit(x, y, res)
    return res

beta, vcv, info = fit_pseudo1(data['a', '16.5'][0], ax=True)


# In[ ]:


def fit_all(series, fit_cb=fit_pseudo1, plot=False):
    if plot:
        axes = mk_subplots(nrows=len(temperatures), sharex=True, sharey=True)#, subplots_adjust=False)
        _set_axes_titles_to_ionic_strength(axes[0, :])
    avg = {}
    for key in itertools.product(ionic_strengths, temperatures):
        idx_I, idx_T = indices(key)
        opt_params, cov_params = [], []
        for serie in series[key]:
            beta, vcv, nfo = fit_cb(serie)
            opt_params.append(beta)
            cov_params.append(vcv)
        ax = axes[idx_T, idx_I] if plot else None
        avg[key] = avg_params(opt_params, cov_params)
        plot_avg_params(opt_params, cov_params, avg[key], ax=ax, flip=True, nsigma=3)
    for tk, ax in zip(temperature_keys, axes[:, 0]):
        ax.set_ylabel('T = %s C' % tk)
    return avg

result_pseudo1 = fit_all(data, plot=True)


# In[ ]:


def pseudo_to_k2(v):
    unit = 1/sol['Fe+3']/u.second
    k_val = -v[0][1]*unit
    k_err = v[1][1]*unit
    return k_val, k_err

k_pseudo1 = {k: pseudo_to_k2(v) for k, v in result_pseudo1.items()}
k_pseudo1['a', '16.5']


# In[ ]:


k2_unit = 1/u.M/u.s
axes = mk_subplots(sharex=True, sharey=True)
for idxI, (tI, I) in enumerate(ionic_strengths.items()):
    series = np.empty((len(temperatures), 3))
    for idxT, (tT, T) in enumerate(temperatures.items()):
        kval, kerr = [to_unitless(v, k2_unit) for v in k_pseudo1[tI, tT]]
        lnk_err = (np.log(kval + kerr) - np.log(kval - kerr))/2
        series[idxT, :] = to_unitless(1/T, 1/u.K), np.log(kval), 1/lnk_err**2
    x, y, w = series.T
    res = b, vcv, r2 = least_squares(x, y, w)
    plot_least_squares_fit(x, y, res, w**-0.5, plot_cb=functools.partial(
            plot_fit, ax=axes[idxI], kw_data=dict(ls='None', marker='.'), nsigma=3))
    axes[idxI].get_lines()[-1].set_label(r'$E_{\rm a} = %.5g\ kJ/mol$' % (-b[1]*8.314511e-3))
    axes[idxI].legend()
_set_axes_titles_to_ionic_strength(axes)