import numpy as np
import matplotlib.pyplot as plt

from py3d3v.pic3d3v import PIC3DPM, Species

"""
Explanation: Quick Tutorial
This notebook covers the basics of the Particle Mesh (PM) method using py3d3v.
End of explanation
"""


# Dimensions of the simulation domain
L = (2*np.pi, 2*np.pi, 2*np.pi) 
# Number of grid points in each dimension
n = (32, 32, 32)
# Time steps
nt = 100
dt = 0.1

# Build the species class with random initial x and v
size = 1024
xp = np.random.uniform(low=0.0, high=L[0], size=size)
yp = np.random.uniform(low=0.0, high=L[1], size=size)
zp = np.random.uniform(low=0.0, high=L[2], size=size)

vx = np.random.uniform(low=-.5, high=.5, size=size)
vy = np.random.uniform(low=-.5, high=.5, size=size)
vz = np.random.uniform(low=-.5, high=.5, size=size)

q = np.ones_like(xp)
m = np.ones_like(xp)

species = Species(size, q, m, 
                  z0=zp, y0=yp, x0=xp,
                  vz0=vz, vy0=vy, vx0=vx)

"""
Explanation: Data Layouts
Particle positions, velocities, charge and mass are stored in seperate linear 1-D arrays. In general quantities related to particles are stored in 1-D arrays. 
The initial particle information is stored in the Species class. The primary purpose of this class is to prep the arrays for the simulation classes
End of explanation
"""


# Construct the PIC class
pic = PIC3DPM([species], L, n)
# Always call init_run before anything else
pic.init_run(dt)

# Call time step to advance the simulation
for it in range(nt):
    pic.single_step(dt)

"""
Explanation: The PIC3DPM class performs all of the simulation steps for the particle mesh method. Species are passed in a list.
End of explanation
"""


plt.hist(pic.vx)
plt.xlabel("$v_x$")

"""
Explanation: The particle positions and velocities are stored in the obvious array names.
End of explanation
"""


