"""Basic testing to ensure a pair of galaxies follows the correct orbit."""

import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import savgol_filter

from analysis import utils

f, ax = plt.subplots(1, 1)

time = []
deviations = []

for n in np.linspace(100, 10000, 100):
	# Load the data
	data = utils.loadData('orbit', int(n))
	# Find the deviation
	# The analytical orbit is r = 1/(1 + cos(phi))
	time.append(data['t'])
	phi = np.arctan2(data['r_vec'][0,1], data['r_vec'][0,0])
	truthr = 1/(1 + np.cos(phi))
	print(truthr, np.linalg.norm(data['r_vec'][0,:2]))
	deviation = np.abs(np.linalg.norm(data['r_vec'][0,:2]) - truthr)
	deviations.append(deviation)

# Plot
deviations = savgol_filter(deviations, 11, 3)
ax.plot(time, deviations, color='darkorange')
utils.stylizePlot([ax])
utils.setAxes(ax, x='Time', y='Error in radius', 
	xcoords=(.85,-0.08), ycoords=(-0.1,.75))
ax.set_yscale('log')
ax.set_xlim(0, time[-1])
ax.set_ylim(1E-5, None)

#plt.legend(loc='upper left')
plt.show()