"""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


def calculatePotentialEnergy(r_vec, v_vec, mass, soft=0.):
    """Calculates the potential energy of a set of masses.
       It is an adaptation of the acceleration routines.
    """
    mask = mass!=0
    massMassive = mass[mask]
    rMassive_vec = r_vec[mask]
    mass_mat = massMassive.reshape(1, -1, 1)
    r_ten = rMassive_vec.reshape(1, -1, 3) - r_vec.reshape(-1, 1, 3)
    r_mat = np.linalg.norm(r_ten, axis=-1, keepdims=True)
    # Avoid self interactions and divisions by 0
    mass_mat = np.where(r_mat == 0, np.zeros_like(mass_mat), mass_mat)
    r_mat = np.where(r_mat == 0, np.ones_like(r_mat), r_mat)
    PE = -mass_mat/(r_mat + soft)
    # Add all contributions but do not sum them twice
    PE = np.sum(PE, axis=1)[:,0]/2 * mass
    energy = PE + 1/2 * np.linalg.norm(v_vec, axis=-1)**2 * mass
    return np.sum(energy)
    '''PE = np.sum(PE, axis=1)[:,0]
                print((mass==0)[:5])
                print('PE', PE[:5], (1/2 * np.linalg.norm(v_vec, axis=-1)**2)[:5])
                energy = PE + 1/2 * np.linalg.norm(v_vec, axis=-1)**2
                return np.sum(energy[mass==0])'''

#calculatePotentialEnergy(np.array([[0,0,0],[1,0,0], [2,0,0]]), np.array([1,1, 0]), soft=0.)

f, ax = plt.subplots(1, 1)
for name, label, factor, color in zip(['dt1E-2', 'dt1E-3', 'dt1E-4'],
	[r'$dt=10^{-2}$', r'$dt=10^{-3}$', r'$dt=10^{-4}$'],
	[1, 10, 100], 
	['cornflowerblue', 'darkorange', '#75d2aa']):
	time = []
	errors = []

	# Use start of the simulation as ground truth
	truth = utils.loadData(name, 10*factor)
	truthEnergy = calculatePotentialEnergy(truth['r_vec'], truth['v_vec'], truth['mass'], soft=0.1)

	for n in np.linspace(10, 1000, 100):
		# Load the data
		data = utils.loadData(name, int(n*factor))
		# Find the deviation
		time.append(data['t'])
		# Energies per unit mass
		energy = calculatePotentialEnergy(data['r_vec'], data['v_vec'], data['mass'], soft=0.1)
		errors.append(np.abs(energy - truthEnergy)/np.abs(truthEnergy))

	# Plot
	errors = savgol_filter(errors, 11, 3) # Smooth
	ax.plot(time, errors, color=color, label=label)
	# Styling
	utils.stylizePlot([ax])
	utils.setAxes(ax, x='Time', y='Relative energy error', 
		xcoords=(.85,-0.08), ycoords=(-0.1,.75))
	ax.set_yscale('log')
	ax.set_xlim(0, time[-1])
	ax.set_ylim(1E-4, None)
	ax.axvline(x=3.55, linestyle='--', linewidth=1, color='black')

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