#Classes to import
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx
import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.animation as animation
from ast2000tools.solar_system import SolarSystem
import ast2000tools.constants

##########
#ADD GLOBAL VARIABLES HERE (Boltzmann constant, solar mass, satellite mass etc.)
##########

def simulate_box(Nparticles, Nstep, dt, L, T, measure_means = False, measure_pressure = False, holesize=None):
    # Method that takes in number of particles, number of timesteps, length of timesteps, width of box, temperature and some other parameted
    # and returns position array (pos) for each partivle in all 3 dimesons at each timestep (shape: Nparticles x 3 x Nstep) along with some other variables 

    
    #Definng arrays for position and velocity for each particle:
    x = ... #array with shape (Nparticles, 3), some given distribution
    v = ... #array with shape (Nparticles, 3), some given distribution
    pos = np.zeros((int(Nparticles), 3, Nstep))  #array for position of each particles in all 3 directions at each timestep
    pos[:,:,0] = x[:,:] #initial condition


    for i in range(Nstep-1):
        ...
        # pos array should be filled during this loop


    ...
    #Do any other necessary computations
        
    return pos

def plot_box(positions, i,  Nparticles, dt, L, ax, colorVal=None):
    #plot particles in box at any given time (i) given that fig and ax is already defined
    #note: positions array must here be of dimensions (Nparticles, 3). 

    if len(colorVal)!=0:
        im = ax.scatter(positions[:,0,i], positions[:,1,i], positions[:,2,i], 'o', color=colorVal)
    else:
        im = ax.scatter(positions[:,0,i], positions[:,1,i], positions[:,2,i], 'o')
        
    # Set the axes properties
    ax.set_xlim3d([0.0, L])
    ax.set_xlabel('X')
    
    ax.set_ylim3d([0.0, L])
    ax.set_ylabel('Y')
    
    ax.set_zlim3d([0.0, L])
    ax.set_zlabel('Z')

    t = float(i)*dt
    ax.set_title('Particle distribution, t=%3.2g s' %t)

    return im

def plot_box_snapshot(positions, i, Nparticles, dt, L, filename=None):
    #plot snapshot of the box at given time

    #Set up plot environment
    fig = plt.figure()
    ax = p3.Axes3D(fig)
    
    #Define color array from color_map
    jet = cm = plt.get_cmap('jet')
    values = np.random.uniform(size=Nparticles)
    cNorm  = colors.Normalize(vmin=0, vmax=1)
    scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=jet)
    colorVal = scalarMap.to_rgba(values)
    
    #Plot
    im = plot_box(positions, i, Nparticles, dt, L, ax, colorVal=colorVal)
    plt.show()
    
def animate_box(pos, Nparticles, Nstep, dt, L, filename):
    #Animating the box

    #Set up plot environment
    fig = plt.figure()
    ax = p3.Axes3D(fig)
    
    #Define color array from color_map
    jet = cm = plt.get_cmap('jet')
    values = np.random.uniform(size=Nparticles)
    cNorm  = colors.Normalize(vmin=0, vmax=1)
    scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=jet)
    colorVal = scalarMap.to_rgba(values)

    #Plot init snapshot forst
    im = plot_box(pos, 0, Nparticles, dt, L, ax, colorVal=colorVal)

    def init():
        im = plot_box(pos, 0, Nparticles, dt, L, ax, colorVal=colorVal)
        return [im]
    def animate_next_frame(i):
        print("Plotting timestep %d" %i)
        ax.cla()
        im = plot_box(pos, i, Nparticles, dt, L, ax, colorVal=colorVal)
        return [im]

    anim = animation.FuncAnimation(fig, animate_next_frame, init_func=init, frames=Nstep)
    print('Save %s.mp4' % filename)
    anim.save('%s.mp4' % filename, fps=4, extra_args=['-vcodec', 'libx264'])


pos = simulate_box(100, 1001, 1e-12, 1e-6, 1e4)
animate_box(pos, 100, 1001, 1e-12, 1e-6, "boxanimation")
plot_box_snapshot(pos, 0, 100, 1e-12, 1e-6)