%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt

import sys, time, math
import numpy as np
from numpy import linalg as nplin

"""
Explanation: <font color='red'>HJCFIT- maximum likelihood fit of single-channel data: </font>
<font color='red'>Records at four concentrations fitted simultaneously</font>
Some general settings:
End of explanation
"""


from dcpyps.samples import samples
from dcpyps import dataset, mechanism, dcplots, dcio

# LOAD DATA: Burzomato 2004 example set.
scnfiles = [["./samples/glydemo/A-10.scn"], 
            ["./samples/glydemo/B-30.scn"],
            ["./samples/glydemo/C-100.scn"], 
            ["./samples/glydemo/D-1000.scn"]]
tr = [0.000030, 0.000030, 0.000030, 0.000030]
tc = [0.004, -1, -0.06, -0.02]
conc = [10e-6, 30e-6, 100e-6, 1000e-6]

"""
Explanation: Load data
HJCFIT depends on DCPROGS/DCPYPS module for data input and setting kinetic mechanism:
End of explanation
"""


# Initaialise SCRecord instance.
recs = []
bursts = []
for i in range(len(scnfiles)):
    rec = dataset.SCRecord(scnfiles[i], conc[i], tr[i], tc[i])
    recs.append(rec)
    bursts.append(rec.bursts.intervals())
    rec.printout()

"""
Explanation: Initialise Single-Channel Record from dcpyps. Note that SCRecord takes a list of file names; several SCN files from the same patch can be loaded.
End of explanation
"""


# LOAD FLIP MECHANISM USED in Burzomato et al 2004
mecfn = "./samples/mec/demomec.mec"
version, meclist, max_mecnum = dcio.mec_get_list(mecfn)
mec = dcio.mec_load(mecfn, meclist[2][0])

# PREPARE RATE CONSTANTS.
# Fixed rates.
#fixed = np.array([False, False, False, False, False, False, False, True,
#    False, False, False, False, False, False])
for i in range(len(mec.Rates)):
    mec.Rates[i].fixed = False

# Constrained rates.
mec.Rates[21].is_constrained = True
mec.Rates[21].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[21].constrain_args = [17, 3]
mec.Rates[19].is_constrained = True
mec.Rates[19].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[19].constrain_args = [17, 2]
mec.Rates[16].is_constrained = True
mec.Rates[16].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[16].constrain_args = [20, 3]
mec.Rates[18].is_constrained = True
mec.Rates[18].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[18].constrain_args = [20, 2]
mec.Rates[8].is_constrained = True
mec.Rates[8].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[8].constrain_args = [12, 1.5]
mec.Rates[13].is_constrained = True
mec.Rates[13].constrain_func = mechanism.constrain_rate_multiple
mec.Rates[13].constrain_args = [9, 2]
mec.update_constrains()
# Rates constrained by microscopic reversibility
mec.set_mr(True, 7, 0)
mec.set_mr(True, 14, 1)

# Update constrains
mec.update_constrains()

#Propose initial guesses different from recorded ones 
initial_guesses = [5000.0, 500.0, 2700.0, 2000.0, 800.0, 15000.0, 300.0, 120000, 6000.0,
                   0.45E+09, 1500.0, 12000.0, 4000.0, 0.9E+09, 7500.0, 1200.0, 3000.0, 
                   0.45E+07, 2000.0, 0.9E+07, 1000, 0.135E+08]

#initial_guesses = [3687.69, 6091.43, 2467.35, 32621.5, 7061.15, 129984., 1050.69, 20984., 3387.64,
#                   0.166224E+09, 20783.8, 6308.02, 2258.42, 0.332447E+09, 31335.4, 144.530, 831.686, 
#                   0.620171E+06, 554.457, 0.124034E+07, 277.229, 0.186051E+07]

#initial_guesses = mec.unit_rates()
mec.set_rateconstants(initial_guesses)
mec.update_constrains()
mec.printout()

"""
Explanation: Load demo mechanism (C&H82 numerical example)
End of explanation
"""


# Scale for ideal pdf
def scalefac(tres, matrix, phiA):
    eigs, M = eig(-matrix)
    N = inv(M)
    k = N.shape[0]
    A, w = np.zeros((k, k, k)), np.zeros(k)
    for i in range(k):
        A[i] = np.dot(M[:, i].reshape(k, 1), N[i].reshape(1, k))
    for i in range(k):
        w[i] = np.dot(np.dot(np.dot(phiA, A[i]), (-matrix)), np.ones((k, 1)))
    return 1 / np.sum((w / eigs) * np.exp(-tres * eigs))

from dcprogs.likelihood import QMatrix
from dcprogs.likelihood import missed_events_pdf, ideal_pdf, IdealG, eig, inv

fig, axes = plt.subplots(len(recs), 2, figsize=(12,15))
for i in range(len(recs)):
    mec.set_eff('c', recs[i].conc)
    qmatrix = QMatrix(mec.Q, mec.kA)
    idealG = IdealG(qmatrix)
    
    # Plot apparent open period histogram
    ipdf = ideal_pdf(qmatrix, shut=False) 
    iscale = scalefac(recs[i].tres, qmatrix.aa, idealG.initial_occupancies)
    epdf = missed_events_pdf(qmatrix, recs[i].tres, nmax=2, shut=False)
    dcplots.xlog_hist_HJC_fit(axes[i,0], recs[i].tres, recs[i].opint,
                               epdf, ipdf, iscale, shut=False)
    axes[i,0].set_title('concentration = {0:3f} mM'.format(recs[i].conc*1000))

    # Plot apparent shut period histogram
    ipdf = ideal_pdf(qmatrix, shut=True)
    iscale = scalefac(recs[i].tres, qmatrix.ff, idealG.final_occupancies)
    epdf = missed_events_pdf(qmatrix, recs[i].tres, nmax=2, shut=True)
    dcplots.xlog_hist_HJC_fit(axes[i,1], recs[i].tres, recs[i].shint,
                               epdf, ipdf, iscale, tcrit=math.fabs(recs[i].tcrit))
    axes[i,1].set_title('concentration = {0:6f} mM'.format(recs[i].conc*1000))

fig.tight_layout()

"""
Explanation: Check data histograms and probability densities calculated from initial guesses
Plot dwell-time histograms for inspection. In single-channel analysis field it is common to plot these histograms with x-axis in log scale and y-axis in square-root scale. After such transformation exponential pdf has a bell-shaped form.
Note that to properly overlay ideal and missed-event corrected pdfs ideal pdf has to be scaled (need to renormailse to 1 the area under pdf from $\tau_{res}$).
End of explanation
"""


def dcprogslik(x, lik, m, c):
    m.theta_unsqueeze(np.exp(x))
    l = 0
    for i in range(len(c)):
        m.set_eff('c', c[i])
        l += lik[i](m.Q)
    return -l * math.log(10)

def printiter(theta):
    global iternum, likelihood, mec, conc
    iternum += 1
    if iternum % 100 == 0:
        lik = dcprogslik(theta, likelihood, mec, conc)
        print("iteration # {0:d}; log-lik = {1:.6f}".format(iternum, -lik))
        print(np.exp(theta))

# Import HJCFIT likelihood function
from dcprogs.likelihood import Log10Likelihood

kwargs = {'nmax': 2, 'xtol': 1e-12, 'rtol': 1e-12, 'itermax': 100,
    'lower_bound': -1e6, 'upper_bound': 0}
likelihood = []

for i in range(len(recs)):
    likelihood.append(Log10Likelihood(bursts[i], mec.kA,
        recs[i].tres, recs[i].tcrit, **kwargs))

# Extract free parameters
theta = mec.theta()
print ('\ntheta=', theta)
print('Number of free parameters = ', len(theta))

lik = dcprogslik(np.log(theta), likelihood, mec, conc)
print ("\nInitial likelihood = {0:.6f}".format(-lik))

"""
Explanation: Prepare likelihood function
End of explanation
"""


from scipy.optimize import minimize
print ("\nScyPy.minimize (Nelder-Mead) Fitting started: " +
       "%4d/%02d/%02d %02d:%02d:%02d"%time.localtime()[0:6])
iternum = 0
start = time.clock()
start_wall = time.time()
maxiter = 200
# maxiter = 30000
result = minimize(dcprogslik, np.log(theta), args=(likelihood, mec, conc), method='Nelder-Mead', callback=printiter, 
                  options={'xtol':1e-5, 'ftol':1e-5, 'maxiter': maxiter, 'maxfev': 150000, 'disp': True})
t3 = time.clock() - start
t3_wall = time.time() - start_wall
print ("\nScyPy.minimize (Nelder-Mead) Fitting finished: " +
       "%4d/%02d/%02d %02d:%02d:%02d"%time.localtime()[0:6])
print ('\nCPU time in ScyPy.minimize (Nelder-Mead)=', t3)
print ('Wall clock time in ScyPy.minimize (Nelder-Mead)=', t3_wall)
print ('\nResult   ==========================================\n', result)

print ("\nFinal likelihood = {0:.16f}".format(-result.fun))
mec.theta_unsqueeze(np.exp(result.x))
print ("\nFinal rate constants:")
mec.printout()

"""
Explanation: Run optimisation
To keep execution time of this notebook short we only run the optimization for 200 iterations. Change maxiter below for a more realistic optimisation.
End of explanation
"""


fig, axes = plt.subplots(len(recs), 2, figsize=(12,15))
for i in range(len(recs)):
    mec.set_eff('c', recs[i].conc)
    qmatrix = QMatrix(mec.Q, mec.kA)
    idealG = IdealG(qmatrix)
    
    # Plot apparent open period histogram
    ipdf = ideal_pdf(qmatrix, shut=False) 
    iscale = scalefac(recs[i].tres, qmatrix.aa, idealG.initial_occupancies)
    epdf = missed_events_pdf(qmatrix, recs[i].tres, nmax=2, shut=False)
    dcplots.xlog_hist_HJC_fit(axes[i,0], recs[i].tres, recs[i].opint,
                               epdf, ipdf, iscale, shut=False)
    axes[i,0].set_title('concentration = {0:3f} mM'.format(conc[i]*1000))

    # Plot apparent shut period histogram
    ipdf = ideal_pdf(qmatrix, shut=True)
    iscale = scalefac(recs[i].tres, qmatrix.ff, idealG.final_occupancies)
    epdf = missed_events_pdf(qmatrix, recs[i].tres, nmax=2, shut=True)
    dcplots.xlog_hist_HJC_fit(axes[i,1], recs[i].tres, recs[i].shint,
                               epdf, ipdf, iscale, tcrit=math.fabs(recs[i].tcrit))
    axes[i,1].set_title('concentration = {0:6f} mM'.format(conc[i]*1000))

fig.tight_layout()

"""
Explanation: Plot experimental histograms and predicted pdfs
End of explanation
"""


