from brown.interaction import UFuncWrapper
from brown.brown import BrownIterator
import numpy as np
from collections import deque
from copy import copy
import matplotlib.pyplot as plt
import matplotlib.animation as ani

from coefficients import c

#force_function = UFuncWrapper(0, c)
#interaction2D = UFuncWrapper(1, c)

# Borders for both the plot and the boundary condition
# (the boundary condition might be deactivated, when creating
# the BrownIterator).
borders_x = [-100, 100]
borders_y = [-100, 100]
n_particles = 600
# Idk, seems to not do anyting.
frames = 100
# Only spawn in 1/x of the borders.
spawn_restriction = 3
# Time resolution. Note that setting this to a too
# high value (i.e. low resolution) will lead to 
# erratic behaviour, because potentials can be skipped.
dt = 0.1
c[-1] = dt

# Initial positions.
x_coords = np.random.uniform(borders_x[0] / spawn_restriction, borders_x[1] / spawn_restriction, n_particles).astype(np.float16)
y_coords = np.random.uniform(borders_y[0] / spawn_restriction, borders_y[1] / spawn_restriction, n_particles).astype(np.float16)


# Initial momenta are 0.
x_momenta = np.zeros(n_particles, dtype=np.float16)
y_momenta = np.zeros(n_particles, dtype=np.float16)



# Prepare the plot, remove axis & stuff.
fig = plt.figure(figsize=(7, 7))
ax = fig.add_axes([0, 0, 1, 1], frameon=False)
ax.set_xlim(*borders_x)
ax.set_xticks([])
ax.set_ylim(*borders_y)
ax.set_yticks([])

# Plot the initial values.
plot, = ax.plot(x_coords, y_coords, "b.")
center_of_mass, = ax.plot(x_coords.mean(), y_coords.mean(), "r-")
# Keep track of the center of mass.
center_of_mass_history_x = deque([x_coords.mean()])
center_of_mass_history_y = deque([y_coords.mean()])

brown = BrownIterator(-1, c # Max iterations, simulation parameters.
			, x_coords, y_coords
			, y_momenta, y_momenta
			# The boundary condition: reflect at the borders,
			, borders_x, borders_y
			# or just let propagate to infinity.
			#, [], []
			# Let the border dampen the system, border_dampening < 1 => energy is absorbed.
			, border_dampening=1
			, dt=dt)

u = iter(brown)

def update(i):
	# Get the next set of positions.
	data = next(u)
	center_of_mass_history_x.append(x_coords.mean())
	center_of_mass_history_y.append(y_coords.mean())

	plot.set_data(*data)
	center_of_mass.set_data(center_of_mass_history_x, center_of_mass_history_y)

animation = ani.FuncAnimation(fig, update, range(frames), interval=1)
plt.show()
#animation.save("animation.mp4", fps=30)