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

# ## Molecular dynamics 
# 
# * *simple visualization* of live MD simulation visualisation 
# 

# In[ ]:


# Run if you need extra packages: numba and scipy
get_ipython().system('pip install scipy numba')


# In[ ]:


from IPython.display import clear_output

import numpy as np
from MD_utils import lennardjones, simple_molecule_vis
from scipy.constants import N_A
epsilon = 0.9959
sigma = 0.3405
k_B = 0.008311
M = 39.948


# ### initial condition

# In[ ]:


plt = simple_molecule_vis(bs=1)


# In[ ]:


plt.plot.display()


# In[ ]:


import numpy  as np
# bs jest rozmiarem kostki w której znajdują się cząstki.  
nx = 8
dx = sigma *1.2
bs = dx*nx 
box = np.array([bs,bs,bs])
try:
    plt.update_box(bs=bs)
    print("no plot object")
except:
    pass

Nparticles = nx**3
print(Nparticles)

U = np.zeros((3,Nparticles))
l = 0
for i in range(nx):
    for j in range(nx):
        for k in range(nx):
            U[:,l] = (dx*i-dx*(nx/2-0.50),dx*j-dx*(nx/2-0.5),dx*k-dx*(nx/2-0.5))
            l+=1
V = np.zeros_like(U)
plt.pkts.positions = U.T.copy()
plt.pkts.point_size = .3
plt.pkts.colors = 0xff0000*np.ones(Nparticles)


# In[ ]:


box = np.array([bs,bs,bs])
get_ipython().run_line_magic('time', 'Epot, F, Vir = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)')


# In[ ]:


V = 0.1*(np.random.randn(3,Nparticles)-0.5)


# In[ ]:


get_ipython().run_cell_magic('time', '', '# Start simulation\n\nimport time \nn_steps = 1520\ndt = 0.005\nN = Nparticles\n\nbox = np.array([bs,bs,bs])\nplt.update_box(bs=bs)\n\n(epot,F,Vir) = lennardjones(U,box,sigma = 0.3405, epsilon=0.9959)\ntraj = []\nfor i in range(n_steps):\n    U += V * dt + 0.5 * F/M * dt * dt\n    F0 = F[:]\n\n    (epot, F, Vir) = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)\n    \n    V += 0.5 * (F + F0)/M * dt\n    U -= bs*np.rint(U/bs)\n    \n    traj.append(U[:,233].copy())\n    \n    if i%10==0:\n        T = M*np.sum(V**2)/(k_B*(3*N-6))  \n        P = 1/bs**3*(  N*k_B*T + 1/3.*Vir )\n        Ek = 0.5*M*np.sum(V**2)\n        plt.pkts.positions = U.T.copy()\n        plt.pkts.colors = (np.sum((V**2),axis=0)/10.0* 0xffffff).astype(np.int64)\n        plt.update_box(bs=bs)\n        clear_output(wait=True)\n     \n        print(i,epot,Ek,epot+Ek,"T=",T,"P=",P)\n        print("Vir:",Vir)\n        print("bs:",bs)\n    \ntraj = np.array(traj)\n')


# ### The trajectory of a single atom 

# In[ ]:


import k3d


# In[ ]:


plt_traj = k3d.points(traj.copy(), color=0xffff00, point_size=.05)
plt.plot += plt_traj


# In[ ]: