import sys
import os.path as op
import math
sys.path.insert(0, op.abspath('..'))

import numpy as np
import scipy as sp
from scipy import stats, linalg

import matplotlib.pyplot as plt
%matplotlib inline
from mpl_toolkits.mplot3d import proj3d

import vlgp
from vlgp import util, hyper
from vlgp import simulation as sim

def lorenz(n, dt=0.01, s=10, r=28, b=8/3, x0=(0.0, 1.0, 1.05)):
    """Lorenz system
    dx/dt = s*(y - x)
    dy/dt = x*(r - z) - y
    dz/dt = x*y - b*z
    """
    from numpy import empty, inf
    from numpy.linalg import norm

    def dlorenz(x, y, z):
        x_dot = s * (y - x)
        y_dot = x * (r - z) - y
        z_dot = x * y - b * z
        return x_dot, y_dot, z_dot

    xs = empty((n, 3), dtype=float)
    xs[0, :] = x0
    
    # Euler intergration
    for i in range(n - 1):
        dx, dy, dz = dlorenz(xs[i, 0], xs[i, 1], xs[i, 2])
        xs[i + 1, 0] = xs[i, 0] + dx * dt
        xs[i + 1, 1] = xs[i, 1] + dy * dt
        xs[i + 1, 2] = xs[i, 2] + dz * dt

    return xs

"""
Explanation: Multiscale Lorenz Examples
End of explanation
"""


ntrial = 20  # number of trials
obs_ndim = 30  # number of neurons (spike trains)
dyn_ndim = 3  # latent dimensions
binwidth = 1  # in ms 
tau = 100 # in ms
nbin = 1000 # 1s

"""
Explanation: We used the standard Lorenz system, $\sigma = 10$, $\rho = 28$ and $\beta = 10/3$. To show vLGP works for various latent dynamics, we use 3 different sampling rate $\Delta t$, a moderate one (0.006), a slow one (0.001) and a fast one (0.01).
Here we only simulated spike trains of 30 neurons under one condition (random initial state). The synthetic dataset consist of 20 trials per condition. We fix each trial as 1 sec and the time binwidth is 1 ms.
End of explanation
"""


np.random.seed(1)
bias = np.log(5 / nbin)  # log base firing rate
a = 5 * ((np.random.rand(dyn_ndim, obs_ndim) + 1) * np.sign(np.random.randn(dyn_ndim, obs_ndim))) # loading matrix
b = np.vstack((bias * np.ones(obs_ndim)))

"""
Explanation: We use the same random weights and 5 spike/s baseline firing rate for all the data.
End of explanation
"""


dt = 0.006
t = 100
n = math.ceil(t / dt)
np.random.seed(0)
x0 = np.random.randn(3)
x = lorenz(n * 2, dt, s=10, r=28, b=8/3, x0=x0)
xs = np.array_split(x[n:, :], t // 4, axis=0)
i = 200 * 19 + 100
x_k = x[i : i + nbin, :]
x_k = x_k - x_k.mean(axis=0)
x_k /= linalg.norm(x_k, axis=0, ord=np.inf)

dyn = np.tile(x_k, (ntrial, 1, 1))

"""
Explanation: We keep everything the same except the $\Delta t$ when generating different latent dynamics. 
We use the same inference options for all three types of dynamics.
Moderate Speed Lorenz Dynamics
End of explanation
"""


plt.figure(figsize=(6, 6))
ax1 = plt.subplot(311)
ax1.plot(x_k[:, 0], lw=5, color='k')
ax1.set_axis_off()

ax2 = plt.subplot(312, sharex=ax1)
ax2.plot(x_k[:, 1], lw=5, color='k')
ax2.set_axis_off()

ax3 = plt.subplot(313, sharex=ax1)
ax3.plot(x_k[:, 2], lw=5, color='k')
ax3.set_axis_off()

plt.figure(figsize=(5, 5))
ax = plt.gca(projection='3d')
ax.plot(*x_k.T, color='k')
ax.set_axis_off()

y = np.stack([sim.spike(dyn[i, ...], a, b)[0] for i in range(ntrial)]).squeeze()

"""
Explanation: We plot the three dimensions separately and together.
End of explanation
"""


plt.imshow(y[0, ...].T, vmin=0, vmax=1, interpolation='none', aspect='auto', cmap=plt.cm.Greys);
plt.gca().set_axis_off()

np.random.seed(0)
mod_fit, _ = vlgp.fit(y, ['spike']*y.shape[-1], 
                     dyn_ndim=3,
                     sigma=np.ones(dyn_ndim) * (1 - 1e-3), 
                     omega=np.ones(dyn_ndim)*(2*(binwidth/tau)**2),
                     lag=0, rank=500, 
                     x=dyn, alpha=a, beta=b,  
                     method='VB',  
                     niter=100, tol=1e-5, verbose=False, 
                     learn_param=True, learn_post=True, e_niter=5, m_niter=5, adjust_hessian=False, decay=0,Adam=False,
                     learn_hyper=True, nhyper=5, subsample_size=200, hyper_obj='ELBO', 
                     gp_noise=1e-3, successive=False)
util.save(mod_fit, 'mod_fit.h5')

"""
Explanation: The following is one trial of spike trains. We did not sort the neurons.
End of explanation
"""


x_hat = util.rotate(util.add_constant(mod_fit['mu'][0, ...]), mod_fit['x'][0, ...])

plt.figure(figsize=(8, 9))
ax1 = plt.subplot(311)
ax1.plot(mod_fit['x'][0, :, 0], lw=5, color='k')
ax1.plot(x_hat[:, 0], lw=2, color='r')
ax1.set_axis_off()

ax2 = plt.subplot(312, sharex=ax1)
ax2.plot(mod_fit['x'][0, :, 1], lw=5, color='k')
ax2.plot(x_hat[:, 1], lw=2, color='r')
ax2.set_axis_off()

ax3 = plt.subplot(313, sharex=ax1)
ax3.plot(mod_fit['x'][0, :, 2], lw=5, color='k')
ax3.plot(x_hat[:, 2], lw=2, color='r')
ax3.set_axis_off()

"""
Explanation: We plot one trial of the true (black) and inferred (red) trajectory.
End of explanation
"""


dt = 0.001
t = 100
n = math.ceil(t / dt)
np.random.seed(0)
x0 = np.random.randn(3)
x = lorenz(n * 2, dt, s=10, r=28, b=8/3, x0=x0)
xs = np.array_split(x[n:, :], t // 4, axis=0)
i = 200 * 19 + 100
x_k = x[i : i + nbin, :]
x_k = x_k - x_k.mean(axis=0)
x_k /= linalg.norm(x_k, axis=0, ord=np.inf)

dyn = np.tile(x_k, (ntrial, 1, 1))
nbin = dyn.shape[1]

y = np.stack([sim.spike(dyn[i, ...], a, b)[0] for i in range(ntrial)]).squeeze()
plt.imshow(y[0, ...].T, vmin=0, vmax=1, interpolation='none', aspect='auto', cmap=plt.cm.Greys);
plt.gca().set_axis_off()

np.random.seed(0)
slow_fit, _ = vlgp.fit(y, ['spike']*y.shape[-1], 
                     dyn_ndim=3,
                     sigma=np.ones(dyn_ndim) * (1 - 1e-3), 
                     omega=np.ones(dyn_ndim)*(2*(binwidth/tau)**2),
                     lag=0, rank=500, 
                     x=dyn, alpha=a, beta=b,  
                     method='VB',  
                     niter=100, tol=1e-5, verbose=False, 
                     learn_param=True, learn_post=True, e_niter=5, m_niter=5, adjust_hessian=False, decay=0,Adam=False,
                     learn_hyper=True, nhyper=5, subsample_size=200, hyper_obj='ELBO', 
                     gp_noise=1e-3, successive=False)
util.save(slow_fit, 'slow_fit.h5')

x_hat = util.rotate(util.add_constant(slow_fit['mu'][0, ...]), slow_fit['x'][0, ...])

plt.figure(figsize=(8, 9))
ax1 = plt.subplot(311)
ax1.plot(slow_fit['x'][0, :, 0], lw=5, color='k')
ax1.plot(x_hat[:, 0], lw=2, color='r')
ax1.set_axis_off()

ax2 = plt.subplot(312, sharex=ax1)
ax2.plot(slow_fit['x'][0, :, 1], lw=5, color='k')
ax2.plot(x_hat[:, 1], lw=2, color='r')
ax2.set_axis_off()

ax3 = plt.subplot(313, sharex=ax1)
ax3.plot(slow_fit['x'][0, :, 2], lw=5, color='k')
ax3.plot(x_hat[:, 2], lw=2, color='r')
ax3.set_axis_off()

"""
Explanation: Slow Lorenz Dynamics
End of explanation
"""


dt = 0.01
t = 100
n = math.ceil(t / dt)
np.random.seed(0)
x0 = np.random.randn(3)
x = lorenz(n * 2, dt, s=10, r=28, b=8/3, x0=x0)
xs = np.array_split(x[n:, :], t // 4, axis=0)
i = 200 * 19 + 100
x_k = x[i : i + nbin, :]
x_k = x_k - x_k.mean(axis=0)
x_k /= linalg.norm(x_k, axis=0, ord=np.inf)

dyn = np.tile(x_k, (ntrial, 1, 1))
nbin = dyn.shape[1]

y = np.stack([sim.spike(dyn[i, ...], a, b)[0] for i in range(ntrial)]).squeeze()
plt.imshow(y[0, ...].T, vmin=0, vmax=1, interpolation='none', aspect='auto', cmap=plt.cm.Greys);
plt.gca().set_axis_off()

np.random.seed(0)
fast_fit, _ = vlgp.fit(y, ['spike']*y.shape[-1], 
                     dyn_ndim=3,
                     sigma=np.ones(dyn_ndim) * (1 - 1e-3), 
                     omega=np.ones(dyn_ndim)*(2*(binwidth/tau)**2),
                     lag=0, rank=500, 
                     x=dyn, alpha=a, beta=b,  
                     method='VB',  
                     niter=100, tol=1e-5, verbose=False, 
                     learn_param=True, learn_post=True, e_niter=5, m_niter=5, adjust_hessian=False, decay=0,Adam=False,
                     learn_hyper=True, nhyper=5, subsample_size=200, hyper_obj='ELBO', 
                     gp_noise=1e-3, successive=False)
util.save(fast_fit, 'fast_fit.h5')

x_hat = util.rotate(util.add_constant(fast_fit['mu'][0, ...]), fast_fit['x'][0, ...])

plt.figure(figsize=(8, 9))
ax1 = plt.subplot(311)
ax1.plot(fast_fit['x'][0, :, 0], lw=5, color='k')
ax1.plot(x_hat[:, 0], lw=2, color='r')
ax1.set_axis_off()

ax2 = plt.subplot(312, sharex=ax1)
ax2.plot(fast_fit['x'][0, :, 1], lw=5, color='k')
ax2.plot(x_hat[:, 1], lw=2, color='r')
ax2.set_axis_off()

ax3 = plt.subplot(313, sharex=ax1)
ax3.plot(fast_fit['x'][0, :, 2], lw=5, color='k')
ax3.plot(x_hat[:, 2], lw=2, color='r')
ax3.set_axis_off()

"""
Explanation: Fast Lorenz Dynamics
End of explanation
"""


