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

c = np.array([5, 10, 20, 30, 0, 0, 0, 1, -20, 0, -2, -0.1, 2, 0, 0, 0, 0, 0, 0], dtype=np.float16)
force_function = UFuncWrapper(0, c)
interaction2D = UFuncWrapper(1, c)

x_coords = np.array([0, 1], dtype=np.float16)
y_coords = np.array([0, 0], dtype=np.float16)

x_momenta = np.array([0, 0], dtype=np.float16)
y_momenta = np.array([0, 0], dtype=np.float16)

time = np.arange(0, 5, 1, dtype=np.float16)

time_evolution = deque()

for t in time:
	x_momenta, y_momenta = interaction2D(x_coords, y_coords, x_momenta, y_momenta)
	x_coords += x_momenta
	y_coords += y_momenta

	time_evolution.append(copy(x_coords))

time_evolution = np.array(time_evolution)

particle_one_evolution = time_evolution[:, 0]
particle_two_evolution = time_evolution[:, 1]

plt.subplot(2, 1, 1)

r = np.arange(0, 50, 0.1, dtype=np.float16)
plt.title("Force")
plt.xlabel("particle distance")
plt.ylabel("scalar force")
plt.plot(r, force_function(r))

plt.subplot(2, 1, 2)
plt.title("Particle x-positions over time")
plt.xlabel("time")
plt.ylabel("particle x-position")
h0, = plt.plot(time, particle_one_evolution, label="particle 1")
h1, = plt.plot(time, particle_two_evolution, label="particle 2")
plt.legend(handles=[h0, h1])
plt.show()