# inline plotting/interaction
%pylab inline 

from sympy import * # symbolic python
init_printing() # pretty printing

import numpy as np # numeric python

import time # timing, for performance monitoring

# activate latex text rendering
from matplotlib import rc
rc('text', usetex=True)

# use matplotlib for plotting
import matplotlib.pyplot as plt 
from pylab import meshgrid,cm,imshow,contour,clabel,colorbar,axis,title,show

"""
Explanation: Script for the analysis of kinesin models (See Ref. [1])
References
[1] Fluctuating Currents in Stochastic Thermodynamics II. Energy Conversion and Nonequilibrium Response in Kinesin Models. Altaner, Wachtel and Vollmer (2015)
[2] Non-equilibrium fluctuations and mechanochemical couplings of a molecular motor. Lau, Lacoste, and Mallick, Phys. Rev. Lett., 99, 158102 (2007)
[3] Kinesin's Network of Chemomechanical Motor Cycles. Liepelt and Lipowsky. Phys. Rev. Lett. 98, 258102 (2007)
Set-up jupyter (ipython-notebook) and sympy libraries
End of explanation
"""


# The lambda functions for the expressions discussed and analyzed in Ref. [1]
from lambdification import lau, ll 

# velocity, diffusion and normalized mechanical response for the kinesin model of Ref. [2]
velocityLa, diffusionLa, tmechLa = lau() # 

# velocity, hydrolysis rate, etc. for the 6-state [3] and 4-state [2] kinesin models
# NOTE: quick=False yields expressions that are numerically more stable, but take much more time and memory
vel6, hyd6, dif6, coupling6, invfano6, tmech6, \
vel4, hyd4, dif4, coupling4, invfano4, tmech4, \
vel_relerr, hyd_relerr, dif_relerr = ll(quick=True) 

"""
Explanation: Initialize the kinesin model(s)
Overview of libraries:

models.py: contains the information about several kinesin models
cumulants.py: implements the algorithm described in [1]
lambdification.py: simplifies and transforms sympy expressions obtained for the models to lambda functions which are used to calculate numerical values
End of explanation
"""


[fmin, fmax, mumin, mumax] = [-30,30,-30,30] # boundaries of full phase diagram
resolution = 400 # plot resolution

plotarea = [fmin, fmax, mumin, mumax] # full plotarea

# prepare the plotting grid for kinesin 6 figures
xpts = linspace(fmin, fmax, resolution)
ypts = linspace(mumin, mumax, resolution)
X, Y  = meshgrid(xpts, ypts)


"""
Explanation: Set-up plot range and resolution
Full parameter space (Figs. 5 -- 9 in Ref. [1])
End of explanation
"""


[fmin, fmax, mumin, mumax] = [0,20,5,30] # boundaries of zoomed phase diagram
resolution = 400 # plot resolution

zoomarea = [fmin, fmax, mumin, mumax] # zoomed plotarea

# prepare the plotting grid for kinesin 6 figures
xpts = linspace(fmin, fmax, resolution)
ypts = linspace(mumin, mumax, resolution)
XX, YY  = meshgrid(xpts, ypts)


"""
Explanation: Zoomed parameter space (Fig. 10 in Ref [1])
End of explanation
"""


# Full plot area
V6     = vel6(X,Y)
log_V6 = (np.log( np.abs( V6 ))/np.log(10))
H6     = hyd6(X,Y)
log_H6 = (np.log( np.abs( H6 ))/np.log(10))
D6     = dif6(X,Y)
log_D6 = (np.log( np.abs( D6 ))/np.log(10))
QTC6   = coupling6(X,Y)
F6     = invfano6(X,Y)
T6     = tmech6(X,Y)

V_relerr = vel_relerr(X,Y)
H_relerr = hyd_relerr(X,Y)
D_relerr = dif_relerr(X,Y)

# Zoomed plot area (comparison plot Fig. 10)
VLa_comp = velocityLa(XX,YY)
TLa_comp = tmechLa(XX,YY)

V6_comp = vel6(XX,YY)
T6_comp = tmech6(XX,YY)

"""
Explanation: Use the lambda functions to create data arrays
Lambda function evaluate the expressions at every point in the plotting grid
End of explanation
"""


## Setup for Figures (for instance for talks, slides, posters)

fig_size = (12,10) # in inch
fs = 25 # font size

font = {'family' : 'serif',
        'color'  : 'black',
        'weight' : 'normal',
        'size'   : fs,
        }

figdir = "plotdata/"

def pplot(g, t, logplot=True, highlight=0, crop=[], numerical=False):
    fig = figure(figsize=fig_size)
    
    if(numerical):
        G = g
        GG = G
        ccmm = cm.YlOrBr
    elif(logplot):
        G = np.log(np.abs(g(X,Y)))/np.log(10)
        GG = g(X,Y)
        ccmm = cm.gist_earth
    else:
        G = g(X,Y)
        GG = G
        ccmm = cm.coolwarm

    # the slicing parameter [::-1] reverses the y-axis before plotting
    
    im = imshow( G[::-1], cmap=ccmm, extent=plotarea ) # drawing the function
    if(len(crop)==2):
        im.set_clim(vmin=crop[0], vmax=crop[1])
        
    # adding the Contour lines with labels
    if( not numerical):
        cset1    = contour( X,Y, G, arange(-20,20,1),linewidths=1,linestyles="-",colors='black')
        stalling = contour( X,Y, GG,    [highlight], linewidths=3,linestyles="-",colors='white')
    
    # adding the colobar on the right
    cb = colorbar(im)
    
    # latex fashion title
    title(t, fontdict=font)
    xlabel('$f$', fontdict=font)
    ylabel(r'$\Delta \mu$', fontdict=font) # r'\newcommand' is a raw python string, meaning \n is not replaced by a newline, etc.
    
    #savefig(figdir+t+".png")
    return(fig)


"""
Explanation: General Figure Settings
End of explanation
"""


# Some exemplary figures

pplot(vel6,"Velocity 6-state (Fig. 5b)", crop=[-20,2])
pplot(dif6,"Diffusion 6-state (Fig. 7a)", crop=[-20,2])
pplot(coupling6,"Current ratio 6-state (Fig. 5c)",logplot=False, crop=[-2,2], highlight=1)
pplot(tmech6,"Mechanical Response 6-state (Fig. 8)",logplot=False, crop=[-2,2],highlight=1)

# Efficiency of energy conversion, see Fig. 6
P_mech = X*vel6(X,Y) # mechanical power
P_chem = Y*hyd6(X,Y) # chemical power

inversion = np.sign(P_mech)*np.sign(P_chem)
eff = P_mech/P_chem
EFF = np.minimum((eff)**inversion,(eff)**(-inversion))

pplot(EFF,"Efficiency",crop=[0,1],numerical=True)
print("Done")

"""
Explanation: Create some actual plots
Six-state model of Ref [2]
End of explanation
"""


pplot(vel4,"Velocity 4-state", crop=[-20,2])
pplot(dif4,"Diffusion 4-state", crop=[-20,2])
pplot(coupling4,"Current ratio 4-state",logplot=False, crop=[-2,2], highlight=1)
pplot(tmech4,"Normalized mechanical response 4-state",logplot=False, crop=[-2,2],highlight=1)
print("Done")

"""
Explanation: Corresponding plots in four-state model of Ref [1] (Figures are not in the paper, only relative errors are shown in Fig. 9)
End of explanation
"""


pplot(velocityLa,"Velocity Lau et. al. (2007)", crop=[-20,2])
pplot(diffusionLa,"Diffusion Lau et. al. (2007)", crop=[-20,2])
pplot(tmechLa,"Mechanical Response Lau et. al. (2007)",logplot=False, crop=[-2,2],highlight=1)
print("Done")

"""
Explanation: Some plots for the model of Ref [3] (Note the deviation to the models above in physically unrealistic parameter regions)
End of explanation
"""


# Colors
tango_chameleon3 = "#4e9a06"
tango_sky3 = "#204a87"

#Fonts
overlay = {'family' : 'serif',
        'color'  : 'white',
        'weight' : 'normal',
        'size'   : 55,
        }

a=1.5 # scaling factor: figure size/font size

f = figure(figsize=(3.5*a,1.8*a))
gs = GridSpec(100,100)


axLL = f.add_subplot(gs[:,0:43])
im = axLL.imshow( T6_comp[::-1], cmap=cm.coolwarm, extent=zoomarea, aspect=1 ) # drawing the function
# adding the Contour lines with labels
stalling = contour( XX,YY, V6_comp, [0], linewidths=1.4,linestyles="-",colors='white')

FDR = contour( XX,YY, T6_comp , [1], linewidths=1,linestyles="--",colors=tango_chameleon3)
for c in FDR.collections:
    c.set_dashes([(0, (6, 3))])
NDR = contour( XX,YY, T6_comp , [0], linewidths=1,linestyles="--",colors=tango_sky3)
for c in NDR.collections:
    c.set_dashes([(0, (1.0, 1.0))])

#axLL.set_title(r'Liepelt, Lipowsky,' '\n' r'Phys. Rev. Lett. \textbf{98} (2007)', fontsize=5)
text(.5, 27, r"Liepelt \& Lipowsky" '\n' r"Phys.~Rev.~Lett.~\textbf{98} (2007)", fontdict=overlay,fontsize=6,color='black')  # u'Unícôdè ∫triŋs'


xlabel('$f$', labelpad = 2)
ylabel(r'$\Delta\mu$', labelpad=2) #labelpad: move label closer to axis

im.set_clim(vmin=-2, vmax=2)


axLa = f.add_subplot(gs[:,32:100])
im = axLa.imshow( TLa_comp [::-1], cmap=cm.coolwarm, extent=zoomarea, aspect=1 ) # drawing the function
# adding the Contour lines with labels
stalling = contour( XX,YY, VLa_comp, [0], linewidths=1.4,linestyles="-",colors='white')
FDR = contour( XX,YY, TLa_comp , [1], linewidths=1,linestyles="--",colors=tango_chameleon3)
for c in FDR.collections:
    c.set_dashes([(0, (6, 3))])
NDR = contour( XX,YY, TLa_comp , [0], linewidths=1,linestyles="--",colors=tango_sky3)
for c in NDR.collections:
    c.set_dashes([(0, (1.0, 1.0))])
text(.5, 27, r"Lau, Lacoste \& Mallick" '\n' r"Phys.~Rev.~Lett.~\textbf{99} (2007)", fontdict=overlay,fontsize=6,color='black')  # u'Unícôdè ∫triŋs'

## Modify labels
setp( axLa.get_yticklabels(), visible=False) # hide y labels on La plot
#axLa.set_xticks(arange(5,25,5)) # don't display zero
axLL.set_xticks(arange(0,20,5)) # don't display "20"

xlabel('$f$', labelpad = 2)
im.set_clim(vmin=-2, vmax=2)


cb=colorbar(im,ticks=arange(-2,2.1,.5))

"""
Explanation: Code that generates Fig. 10 in Ref. [1]
End of explanation
"""


