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

# # 0.5 Particles in Harmonic Well
# 
# Here we show how to run some basic particle simulations in a harmonic well. We make use of other Python libraries, like getting random numbers and making complex plots.

# In[1]:


import hoomd, hoomd.htf, hoomd.md
import numpy as np
import matplotlib.pyplot as plt
import seaborn
import scipy.stats as ss
import tensorflow as tf
tf.get_logger().setLevel('ERROR')


# ## Build the computational graph
# 
# Our graph is relatively simply. We just compute the distance from the origin and square that to get a harmonic potential.

# In[2]:


# set-up graph
directory = 'normal'
graph = hoomd.htf.graph_builder(0)
positions = tf.Variable(tf.zeros_like(graph.positions), validate_shape=False, name='positions', trainable=False)
# save positions
save_pos = positions.assign(graph.positions)
# Energy is distance from origin squared
energy = tf.reduce_sum(tf.norm(graph.positions[:,:2], axis=1))
# make it clear to htf that we want to use forces derived from positions, instead of a neighbor list
# by setting positions=True
forces = graph.compute_forces(energy, positions=True)
graph.save(directory,force_tensor=forces, out_nodes=[save_pos])


# ## Build Initial Config
# 
# We make the starting momentum and positions random for our particles and create s snapshot to use in hoomd-blue.

# In[3]:


c = hoomd.context.initialize('--mode=cpu')
# set-up initial distribution
p0bar = [[1,0], [0.1,0.1]]
q0bar = [[-3, -1.5], [0.05, 0.05]]
N = 512 # number of particles
T = 1000 # length of trajectories
P = 32 # period of saving positions

snapshot = hoomd.data.make_snapshot(N=N,
                                    box=hoomd.data.boxdim(Lx=100,
                                                          Ly=100,
                                                          Lz=1),
                                    particle_types=['A'])
q0 = np.zeros((N,3))
q0[:,0] = ss.norm.rvs(scale=q0bar[1][0], loc=q0bar[0][0], size=N)
q0[:,1] = ss.norm.rvs(scale=q0bar[1][1], loc=q0bar[0][1], size=N)
p0 = np.zeros((N,3))
p0[:,0] = ss.norm.rvs(scale=p0bar[1][0], loc=p0bar[0][0], size=N)
p0[:,1] = ss.norm.rvs(scale=p0bar[1][1], loc=p0bar[0][1], size=N)

snapshot.particles.position[:] = q0
snapshot.particles.velocity[:] = p0
snapshot.particles.typeid[:] = 0
system = hoomd.init.read_snapshot(snapshot)


# ## Run the simulation
# Now we run the simulation in the NVE ensemble. We disable the hoomd-blue sorter, so that the particle ordering doesn't change. This makes it easy to plot the particles

# In[4]:


# run Hoomd-blue simulation
with hoomd.htf.tfcompute(directory, device='/CPU:0') as tfcompute:
    tfcompute.attach(save_period=P)
    hoomd.md.integrate.mode_standard(dt=0.001)
    hoomd.md.integrate.nve(group=hoomd.group.all())                                   
    c.sorter.disable()
    hoomd.run(T * P)


# ## Load and plot
# 
# Now we plot

# In[5]:


# load particle trajectories
paths = np.empty((T, N, 2) )

for i in range(1,T):
    variables = hoomd.htf.load_variables(directory, checkpoint=i * P, names=['positions'])
    paths[i] = variables['positions'][:,:2]


# In[9]:


# plot particle trajectories
def make_segments(data, particle_index):
    points = np.array([data[:,particle_index, 0], data[:,particle_index, 1]]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    return segments

from matplotlib.collections import LineCollection
plt.figure(figsize=(8,8))
plt.style.use('seaborn-dark')
for i in range(N):
    lc = LineCollection(make_segments(paths, i), cmap='cividis', norm=plt.Normalize(0,1), alpha=0.5)
    # Set the values used for colormapping
    lc.set_array(np.linspace(1,0,T))
    lc.set_linewidth(2)
    line = plt.gca().add_collection(lc)
plt.xlim(np.nanmin(paths), np.nanmax(paths))
plt.ylim(np.nanmin(paths), np.nanmax(paths))
plt.xticks([])
plt.yticks([])
plt.savefig('normal.png', dpi=90)
plt.show()