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

# *Examples are given here in abbreviated form. We invite you to look over our other examples which explain these calculations in greater detail, as well as read the freud documentation to better understand these examples*

# In[1]:


# import bokeh for plotting; use the inline (local) files
# so that an internet connection is not needed
from bokeh.resources import INLINE
from bokeh.io import output_notebook
output_notebook(resources=INLINE)
from bokeh.plotting import figure, output_file, show
from bokeh.layouts import gridplot
# import numpy
import numpy as np
# import freud and set number of parallel threads
from freud import parallel
parallel.setNumThreads(4)
# import ipywidgets for display purposes
from ipywidgets import IntProgress
from IPython.display import display
# import time to display computation time
import time
    
# from aps_helpers import *

verts = [[-0.25, -0.5],
         [0.25, -0.5],
         [0.25, 0.5], 
         [-0.25, 0.5]]
c_list = ["#30A2DA", "#FC4F30", "#E5AE38", "#6D904F", "#9757DB",
    "#188487", "#FF7F00", "#9A2C66", "#626DDA", "#8B8B8B"]
c_dict = dict()
c_dict[6] = c_list[0]
c_dict[5] = c_list[1]
c_dict[4] = c_list[2]
c_dict[3] = c_list[7]
c_dict[7] = c_list[4]


# In[2]:


def default_bokeh(p):
    """
    wrapper which takes the default bokeh outputs and changes them to more sensible values
    """
    p.title.text_font_size = "18pt"
    p.title.align = "center"

    p.xaxis.axis_label_text_font_size = "14pt"
    p.yaxis.axis_label_text_font_size = "14pt"

    p.xaxis.major_tick_in = 10
    p.xaxis.major_tick_out = 0
    p.xaxis.minor_tick_in = 5
    p.xaxis.minor_tick_out = 0

    p.yaxis.major_tick_in = 10
    p.yaxis.major_tick_out = 0
    p.yaxis.minor_tick_in = 5
    p.yaxis.minor_tick_out = 0

    p.xaxis.major_label_text_font_size = "12pt"
    p.yaxis.major_label_text_font_size = "12pt"


# In[3]:


def local_to_global(verts, positions, orientations):
    """
    Take a list of vertices, positions, and orientations and create
    a list of vertices in the "global coordinate system" for plotting
    in bokeh
    """
    num_particles = len(positions)
    num_verts = len(verts)
    # create list of vertices in the "local reference frame" i.e.
    # centered at (0,0)
    l_verts = np.zeros(shape=(num_particles, num_verts, 2), dtype=np.float32)
    l_verts[:] = verts
    # create array of rotation matrices
    rot_mat = np.zeros(shape=(num_particles, 2, 2), dtype=np.float32)
    for i, theta in enumerate(orientations):
        rot_mat[i] = [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]
    # rotate; uses einsum for speed; please see numpy documentation
    # for more information
    r_verts = np.einsum("lij,lkj->lki", rot_mat, l_verts)
    # now translate to global coordinates
    # need to create a position array with same shape as vertex array
    l_pos = np.zeros(shape=(num_particles, num_verts, 2), dtype=np.float32)
    for i in range(num_particles):
        for j in range(len(verts)):
            l_pos[i,j] = positions[i]
    # translate
    output_array = np.add(r_verts, l_pos)
    return output_array


# # Create Example System and visualize
# 
# We will create a three particle system of rectangles and visualize

# In[4]:


# import freud box class
from freud import box
# create data
fbox = box.Box(Lx=10, Ly=10, is2D=True)
pos = np.array([[-1, 0, 0], [0, 0, 0], [1, 0, 0]], dtype=np.float32)
ang = np.array([0.0, 0.0, 0.0], dtype=np.float32)

side_length = max(fbox.Lx, fbox.Ly)
l_min = -side_length / 2.0
l_min *= 1.1
l_max = -l_min

p = figure(title="System Visualization", x_range=(l_min, l_max), y_range=(l_min, l_max))
patches = local_to_global(verts, pos[:,0:2], ang)
for i in range(3):
    fill_color = np.array([c_list[i] for _ in range(1)])
    l_patches = np.zeros(shape=(1,patches.shape[1],patches.shape[2]), dtype=patches.dtype)
    l_patches[0] = patches[i]
    p.patches(xs=l_patches[:,:,0].tolist(),
              ys=l_patches[:,:,1].tolist(),
              fill_color=fill_color.tolist(),
              line_color="black",
              line_width=1.5,
              legend="{}".format(i))
p.patches(xs=[[-fbox.Lx/2, fbox.Lx/2, fbox.Lx/2, -fbox.Lx/2]],
          ys=[[-fbox.Ly/2, -fbox.Ly/2, fbox.Ly/2, fbox.Ly/2]],
          fill_color=(0,0,0,0), line_color="black", line_width=2)
p.legend.location='bottom_center'
p.legend.orientation='horizontal'
default_bokeh(p)
show(p)


# # Compute Pairing
# 
# The 2D pairing module is found in the order module. Set the max distance to search for neighbors, the number of neighbors to search, and the tolerance for the dot product used in the computation

# In[5]:


# pair stuff
from freud import order
myPair = order.Pairing2D(rmax=1.1, k=2, compDotTol=0.1)


# ## Create array of angles to search for pairs
# 
# Not all shapes have one orientation where a pair could form. In this system a paired rectangle could be either to the left or right i.e. $\theta=0$ or $\theta=\pi$. To account for this, we create an array that gives every particle a set of local orientations to search for a pair, in this case $\theta=0$ or $\theta=\pi$.

# In[6]:


c_ang = np.zeros(shape=(pos.shape[0],2), dtype=np.float32)
c_ang[:,1] = np.pi
myPair.compute(fbox, pos, ang, c_ang)


# # Investigate output
# 
# Two arrays are provided as output: the match array and the pair array.
# 
# ## Match array ##
# 
# The match array is a boolean array that is $1$ if a particle is matched and $0$ if not.
# 
# ## Pair array ##
# 
# The pair array gives the index of the paired particle. Note that if a particle could be paired to multiple particles, only one pair is returned, the lower of the two index values.
# 
# In this example, particle 0 is paired with particle 1, particle 1 is paired with particle 0 (and technically 2), and particle 2 is paired with particle 1.

# In[7]:


x = myPair.getMatch()
print(x)
y = myPair.getPair()
print(y)


# # Other examples

# In[8]:


pos = np.array([[-1.5, 0, 0], [0, 0, 0], [1, 0, 0]], dtype=np.float32)
ang = np.array([0.0, 0.0, 0.0], dtype=np.float32)

side_length = max(fbox.Lx, fbox.Ly)
l_min = -side_length / 2.0
l_min *= 1.1
l_max = -l_min

p = figure(title="System Visualization", x_range=(l_min, l_max), y_range=(l_min, l_max))
patches = local_to_global(verts, pos[:,0:2], ang)
for i in range(3):
    fill_color = np.array([c_list[i] for _ in range(1)])
    l_patches = np.zeros(shape=(1,patches.shape[1],patches.shape[2]), dtype=patches.dtype)
    l_patches[0] = patches[i]
    p.patches(xs=l_patches[:,:,0].tolist(),
              ys=l_patches[:,:,1].tolist(),
              fill_color=fill_color.tolist(),
              line_color="black",
              line_width=1.5,
              legend="{}".format(i))
p.patches(xs=[[-fbox.Lx/2, fbox.Lx/2, fbox.Lx/2, -fbox.Lx/2]],
          ys=[[-fbox.Ly/2, -fbox.Ly/2, fbox.Ly/2, fbox.Ly/2]],
          fill_color=(0,0,0,0), line_color="black", line_width=2)
p.legend.location='bottom_center'
p.legend.orientation='horizontal'
default_bokeh(p)
show(p)

myPair.compute(fbox, pos, ang, c_ang)
x = myPair.getMatch()
print(x)
y = myPair.getPair()
print(y)


# Now that we move particle 0 farther away, it is no longer paired, and particles 1 and 2 are paired.

# In[9]:


pos = np.array([[-1.0, 0, 0], [0, 0, 0], [1, 0, 0]], dtype=np.float32)
ang = np.array([0.1, 0.0, 0.0], dtype=np.float32)

side_length = max(fbox.Lx, fbox.Ly)
l_min = -side_length / 2.0
l_min *= 1.1
l_max = -l_min

p = figure(title="System Visualization", x_range=(l_min, l_max), y_range=(l_min, l_max))
patches = local_to_global(verts, pos[:,0:2], ang)
for i in range(3):
    fill_color = np.array([c_list[i] for _ in range(1)])
    l_patches = np.zeros(shape=(1,patches.shape[1],patches.shape[2]), dtype=patches.dtype)
    l_patches[0] = patches[i]
    p.patches(xs=l_patches[:,:,0].tolist(),
              ys=l_patches[:,:,1].tolist(),
              fill_color=fill_color.tolist(),
              line_color="black",
              line_width=1.5,
              legend="{}".format(i))
p.patches(xs=[[-fbox.Lx/2, fbox.Lx/2, fbox.Lx/2, -fbox.Lx/2]],
          ys=[[-fbox.Ly/2, -fbox.Ly/2, fbox.Ly/2, fbox.Ly/2]],
          fill_color=(0,0,0,0), line_color="black", line_width=2)
p.legend.location='bottom_center'
p.legend.orientation='horizontal'
default_bokeh(p)
show(p)

myPair.compute(fbox, pos, ang, c_ang)
x = myPair.getMatch()
print(x)
y = myPair.getPair()
print(y)


# Now we rotate particle 0 so that it's no longer within the dot product tolerance for being bound

# In[10]:


# set a higher tolerance
myPair = order.Pairing2D(rmax=1.1, k=2, compDotTol=0.5)
pos = np.array([[-1.0, 0, 0], [0, 0, 0], [1, 0, 0]], dtype=np.float32)
ang = np.array([0.1, 0.0, 0.0], dtype=np.float32)

side_length = max(fbox.Lx, fbox.Ly)
l_min = -side_length / 2.0
l_min *= 1.1
l_max = -l_min

p = figure(title="System Visualization", x_range=(l_min, l_max), y_range=(l_min, l_max))
patches = local_to_global(verts, pos[:,0:2], ang)
for i in range(3):
    fill_color = np.array([c_list[i] for _ in range(1)])
    l_patches = np.zeros(shape=(1,patches.shape[1],patches.shape[2]), dtype=patches.dtype)
    l_patches[0] = patches[i]
    p.patches(xs=l_patches[:,:,0].tolist(),
              ys=l_patches[:,:,1].tolist(),
              fill_color=fill_color.tolist(),
              line_color="black",
              line_width=1.5,
              legend="{}".format(i))
p.patches(xs=[[-fbox.Lx/2, fbox.Lx/2, fbox.Lx/2, -fbox.Lx/2]],
          ys=[[-fbox.Ly/2, -fbox.Ly/2, fbox.Ly/2, fbox.Ly/2]],
          fill_color=(0,0,0,0), line_color="black", line_width=2)
p.legend.location='bottom_center'
p.legend.orientation='horizontal'
default_bokeh(p)
show(p)

myPair.compute(fbox, pos, ang, c_ang)
x = myPair.getMatch()
print(x)
y = myPair.getPair()
print(y)


# When we increase the dot product tolerance, particle 0 and particle 1 are paired once again

# In[ ]: