import sys
sys.path.append("..")
from smfsmodels import cossio
import kinetics

"""
Explanation: We start by loading a few libraries, either external or from the smFS package. Above this point, you should be checking elsewhere in case anything failed. Below this point, we are to blame. First we modify the path, so that we can actually install the package and load the modules both for running Brownian dynamics and for running kinetics.
End of explanation
"""


def pmf2d(xq, qk):
    fig, ax = plt.subplots(figsize=(6,5))
    counts, ybins, xbins, image = ax.hist2d(xk, qk, \
            bins=[np.linspace(-12,12,100), np.linspace(-12,12,100)])
    pmf = -np.log(counts.transpose())
    pmf -= np.min(pmf)
    cs = ax.contourf(pmf, \
            extent=[xbins.min(), xbins.max(), ybins.min(), ybins.max()], \
            cmap=cm.rainbow, levels=np.arange(0,10 ,1))
    cbar = plt.colorbar(cs)
    ax.set_xlim(-12,12)
    ax.set_ylim(-12,12)
    ax.set_xlabel('$x$', fontsize=20)
    ax.set_ylabel('$q$', fontsize=20)
    plt.tight_layout()

def calc_rates(y):
    lifeA, lifeB = kinetics.calc_life([y])
    meanA = 1./np.exp(np.mean(np.log([x for x in lifeA if x>0])))
    meanB = 1./np.exp(np.mean(np.log([x for x in lifeB if x>0])))
    errorA = meanA/np.sqrt(len(lifeA))
    errorB = meanA/np.sqrt(len(lifeB))
    return np.mean([meanA, meanB]), np.mean([errorA, errorB])

def smooth(y, box_pts):
    box = np.ones(box_pts)/box_pts
    y_smooth = np.convolve(y, box, mode='same')
    return y_smooth

"""
Explanation: Now we define a couple of functions, that let us produce 2D-histograms and calculate kinetics.
End of explanation
"""


x = np.linspace(-10, 10, 1000)
fig, ax = plt.subplots(2,1, figsize=(6,5), sharex=True)
Gx = [cossio.Gx(y, barrier=3.) for y in x]
dGqxdx = [cossio.dGqxdx(0, y, barrier=3.) for y in x]
ax[0].plot(x, Gx, lw=3)
ax[0].set_ylabel('$G(x)$', fontsize=20)
ax[0].set_ylim(-1.1*5,0.5*5)

ax[1].plot(x, dGqxdx, lw=3)
ax[1].set_xlabel('$x$', fontsize=20)
ax[1].set_ylabel('$\partial G(x)/\partial x$', fontsize=20)
ax[1].hlines(0, -10, 10, linestyle='dashed', linewidth=0.5)

ax[1].set_xlim(-10,10)
#ax[2].set_ylabel('$\partial^2 G(x)/\partial x^2$', fontsize=20)

plt.tight_layout(h_pad=0.2)

"""
Explanation: Molecular potential of mean force
First we show the molecular potential of mean force on the extension coordinate, $x$, represented by a bistable potential.
End of explanation
"""


# Globals
dt = 5e-4
Dx = 1. # Diffusion coefficient for molecular coordinate

"""
Explanation: Diffusion on molecular coordinate
The next step running Brownian dynamics on this surface. For this we must define a couple of global parameters that define the time-scales in the model, the diffusion coefficient ($D_x$) and the timestep ($dt$).
End of explanation
"""


start = time.time()
x, q = [5., 0.]
t, xk, qk = cossio.run_brownian(x0=x, dt=dt, barrier=3., \
                             Dx=Dx, Dq=Dx, F12=0, numsteps=int(1e8), \
                                fwrite=int(0.01/dt))
end = time.time()
print (end - start)

data = np.column_stack((t, xk, qk))
h5file = "data/cossio_kl%g_Dx%g_Dq%g_dt%g.h5"%(0, Dx, 0, dt)
try:
    os.makedirs("data")
except OSError:
    pass
with h5py.File(h5file, "w") as hf:
    hf.create_dataset("data", data=data)

fig, ax = plt.subplots(figsize=(14,3))
ax.plot(t, xk, '.', ms=0.1)
ax.plot(t, smooth(xk, 1000), linewidth=0.5)
ax.set_ylim(-10,10)
ax.set_xlim(0,10000)
ax.set_yticks([-5.,0,5.])
ax.set_xlabel('time', fontsize=20)
ax.set_ylabel('x', fontsize=20)

rate0 = np.array(calc_rates(xk))
print (1./rate0[0])

lifeA, lifeB = kinetics.calc_life([xk])

fig, ax = plt.subplots(1,3,figsize=(10,4), sharey=True, sharex=True)
lifetimes = lifeA+lifeB
ax[0].hist(lifeA, bins=np.linspace(0,np.max(lifetimes),10), density=True)
ax[1].hist(lifeB, bins=np.linspace(0,np.max(lifetimes),10),density=True)
ax[2].hist(lifetimes, bins=np.linspace(0,np.max(lifetimes),10), density=True)
ax[0].set_xlabel(r'$\tau_A$', fontsize=20)
ax[1].set_xlabel(r'$\tau_B$', fontsize=20)
ax[2].set_xlabel(r'$\tau$', fontsize=20)
ax[0].set_ylabel('Frequency', fontsize=20)
ax[0].set_xlim(0,np.max(lifetimes))
plt.tight_layout(w_pad=0)

fig, ax = plt.subplots()
hist, bin_edges, patches = \
        ax.hist(xk, bins=np.arange(-10,10,0.25), density=True)
ax.set_ylabel('P(x)')
ax.set_xlabel('x')
ax.set_xlim(-10,10)
plt.tight_layout()

bin_centers = [0.5*(bin_edges[i]+bin_edges[i+1]) \
               for i in range(len(bin_edges)-1)]
fig, ax = plt.subplots()
ax.plot(bin_centers, -np.log(hist), 'b', lw=7, label='Sim')
ax.plot(bin_centers, [cossio.Gx(x, barrier=3.) + 4.7   for x in bin_centers], 'red', lw=2, label='Calc')
ax.set_ylabel('$G(x)$', fontsize=20)
ax.set_xlabel('$x$', fontsize=20)
ax.set_ylim(0,10)
ax.legend(loc=1)
plt.tight_layout()

data = np.column_stack((t, xk, qk))
tau_f, tau_u, data_f, data_u, tp_f, tp_u, recrossings = \
    kinetics.lifetimes(data, f_bound=-4, u_bound=4)

fig, ax = plt.subplots(figsize=(14,3))
#ax.plot(t, xk, '.', ms=0.1)
#ax.plot(t, smooth(xk, 1000), linewidth=0.5)
#ax.set_ylim(-10,10)
ax.plot(t, xk, lw=0.2)
for x in data_f:
    ax.plot(x[:,0], x[:,1], color='tab:orange', lw=0.1)
for x in data_u:
    ax.plot(x[:,0], x[:,1], color='tab:green', lw=0.1)
#ax.set_yticks([-5.,0,5.])
ax.set_xlim(0,10000)
ax.set_xlabel('time', fontsize=20)
ax.set_ylabel('x', fontsize=20)

tpt_u = []
tpt_f = []
for td in tp_u:
    tpt_u.append(td[:,0][-1] - td[:,0][0])
for td in tp_f:
    tpt_f.append(td[:,0][-1] - td[:,0][0])

from scipy.stats import expon
# All tps
fig, ax = plt.subplots(2,1, sharex=True, sharey=True)
histogram, bin_edges, patches = \
        ax[0].hist(tpt_f, \
                   bins=np.linspace(0,np.max(tpt_f),20), \
                   density=True, cumulative=False)
histogram, bin_edges, patches = \
        ax[0].hist(tpt_u, \
                   bins=np.linspace(0,np.max(tpt_f),20), \
                   density=True, cumulative=False, alpha=0.5)

# Fitting exponential tail
histogram, bin_edges, patches = \
        ax[1].hist([x for x in tpt_f if x > 5], \
                   bins=np.linspace(0,np.max(tpt_f),20), \
                   density=True, cumulative=False)
param = expon.fit([x for x in tpt_f if x > 5] , floc=5)
x = np.linspace(0,np.max(tpt_f),200)
ax[1].plot(x, expon.pdf(x, loc=param[0], scale=param[1]), color='lightgreen', lw=3)

histogram, bin_edges, patches = \
        ax[1].hist([x for x in tpt_u if x>5], \
                   bins=np.linspace(0,np.max(tpt_f),20), \
                   density=True, cumulative=False, alpha=0.5)
param = expon.fit([x for x in tpt_u if x > 5] , floc=5)
ax[1].plot(x, expon.pdf(x, loc=param[0], scale=param[1]), color='tomato', lw=3)

ax[1].set_xlabel(r'$\tau$')
#lx = np.linspace(0,np.log(np.max(lifetimes)),2000)
#pdf = 1/param[1]*np.exp(lx)*np.exp(-1/param[1]*(np.exp(lx)))
#ax[1].plot(lx, pdf, color='lightblue', lw=3)
ax[1].set_xlim(0,20)
#ax[1].set_ylim(0,0.4)
#ax[1].set_yticks([0,0.2,0.4])
plt.tight_layout(h_pad=0)

"""
Explanation: Next we invoke the cossio module to run the dynamics, whose function run_brownian will do the job. We are passing a number of input parameters, including the size of the free energy barrier and some variables for IO.
End of explanation
"""


kl = 0.25


x = np.linspace(-10,10,100)
G2d = np.ones((50, 100), float)*[cossio.Gx(y, barrier=5.) for y in x]
q = np.linspace(-12,12,50)
for i, j in itertools.product(range(50), range(100)):
    G2d[i,j] += cossio.V(q[i], x[j], kl)
G2d = np.array(G2d)
G2d -= np.min(G2d)

import matplotlib.cm as cm
fig, ax = plt.subplots(figsize=(6,5))
cs = ax.contourf(x, q,  G2d, cmap=cm.rainbow, \
                 levels=np.arange(0,12,0.5), alpha=0.9)
cbar = plt.colorbar(cs)
ax.set_xlim(-12,12)
ax.set_ylim(-12,12)
ax.set_xlabel('$x$', fontsize=20)
ax.set_ylabel('$q$', fontsize=20)

expGM = np.trapz(np.exp(-G2d), q, axis=0)
GM = -np.log(expGM)
expGA = np.trapz(np.exp(-G2d), x, axis=1)
GA = -np.log(expGA)

fig, ax = plt.subplots()
ax.plot(x, GM - np.min(GM), label='$G_M$', lw=4)
ax.plot(q, GA  - np.min(GA), label='$G_A$', lw=4)
#ax.plot(bin_centers, [cossio.Gx(y) for y in bin_centers], '--', c='red', lw=3)
ax.set_xlim(-10,10)
ax.set_ylim(-1,7)
ax.set_xlabel('Extension', fontsize=20)
ax.set_ylabel('Free Energy', fontsize=20)
ax.legend(loc=1)
fig.tight_layout()

x, q = [5., 5.]
t, xk, qk = cossio.run_brownian(x0=x, dt=dt, barrier=5., kl=kl, \
                             Dx=Dx, Dq=Dx, numsteps=int(1e8), fwrite=int(0.1/dt))

fig, ax = plt.subplots(2,1, figsize=(10,5), sharex=True, sharey=True)
ax[0].plot(xk, linewidth=0.2)
ax[1].plot(qk, 'r', linewidth=0.2)
ax[1].set_xlabel('time', fontsize=20)
ax[1].set_ylabel('$q$', fontsize=20)
ax[0].set_ylabel('$x$', fontsize=20)
ax[0].set_ylim(-10, 10)
#ax[0].set_xlim(0, numsteps)

scaling = [5, 2, 1, 0.5, 0.2, 0.1, 0.05, 0.02, 0.01]
nsteps = [1e7, 2e7, 2e7, 2e7, 2e7, 2e7, 2e7, 5e7, 5e7, 5e7]
trajs = {}
for i,sc in enumerate(scaling):
    trajs[sc] = {}
    x, q = [5., 5.]
    numsteps = int(nsteps[i])
    xk, qk = cossio.run_brownian(x0=x, dt=dt, barrier=5., kl=kl, \
                             Dx=Dx, Dq=Dx*sc, numsteps=numsteps, fwrite=100)
    trajs[sc]['x'] = xk
    trajs[sc]['q'] = qk
    
    fig, ax = plt.subplots(2,1, figsize=(10,5), sharex=True, sharey=True)
    ax[0].plot(xk, linewidth=0.2, label='Dx/Dq=%g'%(1./sc))
    ax[1].plot(qk, 'r', linewidth=0.2)
    ax[1].set_xlabel('Time', fontsize=20)
    ax[1].set_ylabel('$q$', fontsize=20)
    ax[0].set_ylabel('$x$', fontsize=20)
    ax[0].set_ylim(-10, 10)
    ax[0].legend(loc=1)

kfq = []
kfx = []
for i, sc in enumerate(scaling):
    kfx.append(calc_rates(trajs[sc]['x']))
    kfq.append(calc_rates(trajs[sc]['q']))

fig, ax = plt.subplots()
rates = np.array(kfq)
ax.errorbar(1./np.array(scaling), rates[:,0]/rate0[0], \
            rates[:,1]/rate0[0], marker='o')
ax.set_xscale('log')
ax.set_xlabel('$D_x/D_q$', fontsize=24)
ax.set_ylabel('$k_{MA}/k_M$', fontsize=24)
ax.set_xlim(2e-2,5e2)
plt.tight_layout()

scaling = [10]
nsteps = [1e8]
for i,sc in enumerate(scaling):
    trajs[sc] = {}
    x, q = [5., 5.]
    numsteps = int(nsteps[i])
    xk, qk = cossio.run_brownian(x0=x, dt=dt/5., barrier=5., kl=kl, \
                             Dx=Dx, Dq=Dx*sc, numsteps=numsteps, fwrite=1000)
    trajs[sc]['x'] = xk
    trajs[sc]['q'] = qk
    
    fig, ax = plt.subplots(2,1, figsize=(10,5), sharex=True, sharey=True)
    ax[0].plot(xk, linewidth=0.2, label='Dx/Dq=%g'%(1./sc))
    ax[1].plot(qk, 'r', linewidth=0.2)
    ax[1].set_xlabel('Time', fontsize=20)
    ax[1].set_ylabel('$q$', fontsize=20)
    ax[0].set_ylabel('$x$', fontsize=20)
    ax[0].set_ylim(-10, 10)
    ax[0].legend(loc=1)

scaling = [10, 5, 2, 1, 0.5, 0.2, 0.1, 0.05, 0.02, 0.01]
kfq = []
kfx = []
for i, sc in enumerate(scaling):
    #kfx.append(calc_rates(trajs[sc]['x']))
    kfq.append(calc_rates(trajs[sc]['q']))

fig, ax = plt.subplots()
rates = np.array(kfq)
ax.errorbar(1./np.array(scaling), rates[:,0]/rate0[0], \
            rates[:,1]/rate0[0], marker='o')
ax.set_xscale('log')
ax.set_xlabel('$D_x/D_q$', fontsize=24)
ax.set_ylabel('$k_{MA}/k_M$', fontsize=24)
ax.set_xlim(2e-2,5e2)
plt.tight_layout()

"""
Explanation: Diffusion in two dimensions
End of explanation
"""


