"""Basic testing to ensure the chosen timestep is sensible and to probe
the symplectic character of the integrator."""

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)
for name, label, factor, color in zip(['dt1E-2', 'dt1E-3'],
	[r'$dt=10^{-2}$', r'$dt=10^{-3}$'],
	[1, 10], 
	['cornflowerblue', 'darkorange']):
	time = []
	median = []
	areas = []

	for n in np.linspace(10, 1000, 100):
		truth = utils.loadData('dt1E-4', int(n*100))
		
		# Load the data
		data = utils.loadData(name, int(n*factor))
		# Find the deviation
		time.append(truth['t'])
		deviations = np.linalg.norm(data['r_vec'] - truth['r_vec'], axis=-1)
		median.append(np.median(deviations))
		areas.append(np.percentile(deviations, [6, 16, 31, 100-31, 100-16, 100-6]))

	# Plot
	areas = np.array(areas)
	areas = savgol_filter(areas, 11, 3, axis=0) #Smoothing
	median = savgol_filter(median, 11, 3)
	ax.fill_between(time, areas[:,0], areas[:,-1], alpha=.25, color=color)
	ax.fill_between(time, areas[:,1], areas[:,-2], alpha=.25, color=color)
	ax.fill_between(time, areas[:,2], areas[:,-3], alpha=.25, color=color)
	ax.plot(time, median, color=color, label=label)
	utils.stylizePlot([ax])
	utils.setAxes(ax, x='Time', y='Displacement', 
		xcoords=(.85,-0.08), ycoords=(-0.1,.75))
	ax.set_yscale('log')
	ax.set_xlim(0, time[-1])
	ax.axvline(x=3.55, linestyle='--', linewidth=1, color='black')

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