import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from itertools import product

n = 500 #number of particles
P = 9.9*np.random.rand(n, 3) - 9.9/2
Pprior = P
#P = np.ones((n,3))
Pdot = 6*np.random.rand(n,3) - 2.5
#Pdot = 1*np.ones((n,3))
dt = 0.05 #time step
e = 0.05 #threshold for wall positions
door = 2 #width of demon door

walls = 5 #temporary

def getVertices(N, length):
    """inputs: dimension N, side length. outputs list of vertices centered at origin"""
    return np.array(list(product((length/2,-length/2), repeat=N)))

def update(P,Pdot,dt, Pprior, e):
    #check for colisions
    def getCollisionIndex(vect, cond, e):
        return np.transpose(np.nonzero((vect<(-cond)) | (vect>(cond))))
    def getGate(P, Pdot, dt, e):
        return np.transpose(np.nonzero(((np.sign(P[:,0]) != (np.sign(P[:,0]+dt*1.1*Pdot[:,0]))) & 
                                        ((P[:,1]>1) | (P[:,1]<-1)) & 
                                        ((P[:,2]>1) | (P[:,2]<-1)))))
    def getDoor(P, Pdot, dt):
        indices = np.transpose(np.nonzero((np.sign(P[:,0]) != (np.sign(P[:,0]+dt*1.1*Pdot[:,0]))) &
                                          ((P[:,1]<1) | (P[:,1]>-1)) &
                                          ((P[:,2]<1) | (P[:,2]>-1)) &
                                          (((np.linalg.norm(Pdot, axis=1)>2.5) & (P[:,0] > 0)) | ((np.linalg.norm(Pdot, axis=1)<2.5) & (P[:,0] < 0)))
                                          ))
        return indices

    while True:
        Pprior = P

        wallCollisions = getGate(P, Pdot, dt, e) #the middle wall
        Pdot[wallCollisions,0] = -1 * Pdot[wallCollisions,0]

        doorIndices = getDoor(P, Pdot, dt)
        Pdot[doorIndices,0] = -1 * Pdot[doorIndices,0]

        for i in range(3): 
            #box for the particles
            collisionIndex = getCollisionIndex(P[:,i], 5, e)
            Pdot[collisionIndex,i] = -1 * Pdot[collisionIndex,i]
            
        P = P + dt*Pdot
        P[wallCollisions,0] = Pprior[wallCollisions,0]
        P[doorIndices,0] = Pprior[doorIndices,0]
        yield P, Pdot

fig = plt.figure()
ax = fig.add_subplot(projection="3d")

# Create initial points
#points, = ax.plot(P[:,0],P[:,1],P[:,2],linestyle="", marker="o")
p = np.linalg.norm(Pdot, axis=1)
points = ax.scatter(P[:,0],P[:,1],P[:,2], s=2, c=p, cmap="plasma")#seismic
#ax.scatter.viridis()

ori = ax.plot(0,0,0,linestyle="", marker="o", color="black")
box_pts = getVertices(3,10)
box = ax.plot(box_pts[:,0],box_pts[:,1], box_pts[:,2], linestyle="", marker="o", color="black" )

#plot where middle wall is
yy, zz = np.meshgrid(range(-1,2), range(-1,2))
xx = yy*0
#ax.plot_surface(xx, yy, zz, linewidth=0)
ax.plot_wireframe(xx, yy, zz, color="plum")

my_iter = update(P,Pdot,dt,walls,e)
#Setting the axes properties
ax.set(xlim3d=(-10, 10), xlabel='X')
ax.set(ylim3d=(-10, 10), ylabel='Y')
ax.set(zlim3d=(-10, 10), zlabel='Z')
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
ax.set_aspect('equal')

def update_frame(n):
    pts, ptsdot = next(my_iter)
    #print(next(my_iter))
    points._offsets3d = (pts[:,0],pts[:,1],pts[:,2])
    #points.set_data(pts[:,0],pts[:,1])
    #points.set_3d_properties(pts[:,2])
    #print(ptsdot)
    return points

fig.set_facecolor('black')
ax.set_facecolor('black')
ax.grid(False)
ax.xaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
ax.yaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
ax.zaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
#fyi:
#plt.style.use('dark_background')

# Creating the Animation object
ani = animation.FuncAnimation(
    fig, update_frame, interval=20)

plt.show()
#FFwriter = animation.FFMpegWriter(fps=30, extra_args=['-vcodec','libx264'])
#ani.save('plot.mp4', writer=FFwriter)