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_x = [-100, 100]
borders_y = [-100, 100]
n_particles = 600
frames = 100
spawn_restriction = 1.1
dt = 0.01
c[-1] = dt

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)


x_momenta = np.zeros(n_particles, dtype=np.float16)
y_momenta = np.zeros(n_particles, dtype=np.float16)



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, = ax.plot(x_coords, y_coords, "b.")
center_of_mass, = ax.plot(x_coords.mean(), y_coords.mean(), "r-")
center_of_mass_history_x = deque([x_coords.mean()])
center_of_mass_history_y = deque([y_coords.mean()])

brown = BrownIterator(-1, c
			, x_coords, y_coords
			, y_momenta, y_momenta
			, borders_x, borders_y
			#, [], []
			, border_dampening=1
			, dt=dt)

u = iter(brown)

def update(i):
	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)