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

# ##  $N$ interacting masses in $\cos$ potential
# Additional requirements for this example: `scipy`
# 
# ### Sine Gordon equation
# 
# Let $x$ be an array with position of $N$ masses in  $\cos$ potential  linearly coupled to nearest neighbours.  Then the Newton's equations read:
# 
# $$ \begin{cases}
# \displaystyle\frac{dx}{dt} = v \\
# \displaystyle\frac{dv}{dt} = - M x + \sin(x),\;\;\;(1)
# \end{cases} $$
# 
# where M is a discrete Laplacian. 
# 
# On the other hand this is an aproximation to a sine-Gordon equation: 
# 
# $$ x_{tt}- x_{\xi\xi} + \sin x= 0 $$
# 
# which has the following soliton solutions:
# 
# $$x_\text{soliton}(\xi, t) := 4 \arctan \exp\left\{ \frac{ \xi - v t }{\sqrt{1 - v^2}}\right\}$$
# 
# Below we solve numerically the system (1) and animate results in 3D.
# 
# 

# In[ ]:


from scipy.integrate import odeint
import numpy as np 


# In[ ]:


N = 153
M = np.diag((N-1)*[ 1.0],-1)+np.diag((N-1)*[ 1.0],1)+np.diag(N*[-2.0],0)


# In[ ]:


y0 = np.zeros(2*N)
x = y0[:N] # Fist N variables are positions
v = y0[N:]
ksi = np.linspace(-5,25,N, dtype=np.float32)
h = np.diff(ksi)[0]

instanton = lambda x,v,t: 4 * np.arctan(np.exp( (x-v*t)/np.sqrt(1-v**2) ))
v1,v2 = 0.4,0.05
x[:] = instanton(ksi,v1,0)
v[:] = -v1/h*np.gradient( instanton(ksi,v1,0) )

x[:] += instanton(ksi,v2,50)
v[:] += -v2/h*np.gradient( instanton(ksi,v2,50) )


# In[ ]:


def lhs(y_,t):
    y = y_.copy()
    x = y_[:N]
    v = y_[N:]
    y[:N] = v 
    y[N:] = -np.sin(x) + 1/h**2* np.dot(M,x)  
    y[0] = 0
    y[N-1] = 0 
    return y
ts = np.linspace(0,140,50)  
get_ipython().run_line_magic('time', 'xt = odeint(lhs,y0, ts).astype(np.float32)')


# In[ ]:


import k3d
import numpy as np
import time 
plot = k3d.plot()
r = .6
origins = np.vstack([ksi,np.zeros(N),np.zeros(N)]).T.copy().astype(np.float32)
vectors = np.vstack([np.zeros(N),r*np.sin(xt[0,:N]),r*np.cos(xt[0,:N])]  ).T.astype(np.float32)

vector_plot = k3d.vectors(origins, vectors)
line_plot = k3d.line(vectors + origins,color=0xff0000)

plot += vector_plot
plot += line_plot
plot.display()

def update_plot(xx):
    vectors = np.vstack([np.zeros_like(xx),r*np.sin(xx),r*np.cos(xx)]  ).T
    vector_plot.vectors = vectors.copy()
    line_plot.vertices = vectors+origins


# In[ ]:


plot.camera_auto_fit = False
plot.grid_auto_fit = False


# In[ ]:


get_ipython().run_cell_magic('time', '', 'import time\nplot.auto_rendering = False\nfor xx in xt[:,:N]:\n    update_plot(xx)\n    plot.render()\n    time.sleep(0.05)\n')


# In[ ]:


from ipywidgets import widgets,interact

@interact(ith = widgets.IntSlider(min=0,max=(ts.size-1)))
def draw(ith):
    update_plot(xt[ith,:N])
    plot.render()