"""Plot the eccentricity of the orbits of the particles in an encounter"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils
from analysis.segmentation import segmentEncounter

# Segment the encounter
data = utils.loadData('200mass', 15000)
(bridgeMask, stolenMask, orbittingMask, tailMask), _, _ = segmentEncounter(data)
# Avoid central masses
orbittingMask, stolenMask = orbittingMask[1:], stolenMask[1:] 

# Calculate the eccentricities
tailE = utils.calculateEccentricity(1.0, data['r_vec'][0], data['v_vec'][0], 
	data['r_vec'][tailMask], data['v_vec'][tailMask])
stolenE = utils.calculateEccentricity(2.0, data['r_vec'][1], data['v_vec'][1], 
	data['r_vec'][stolenMask], data['v_vec'][stolenMask])
orbittingE = utils.calculateEccentricity(1.0, data['r_vec'][0], data['v_vec'][0], 
	data['r_vec'][orbittingMask], data['v_vec'][orbittingMask])

# Plotting
f, ax = plt.subplots(1, 1)
bins = np.linspace(0, 10, 301)
style = {'linestyle':'solid', 'normed':True, 'color':'black'}
ax.hist(tailE, bins=bins, histtype='step', label='tail', **style)
ax.hist(stolenE, bins=bins, histtype='step', label='stolen', **style)
ax.hist(orbittingE, bins=bins, histtype='step', label='orbitting', **style)

# Plotting styling
utils.setSize(ax, x=(0,2))
utils.setAxes(ax, x='Eccentricity', y='Density of particles', 
	xcoords=(.85,-0.08), ycoords=(-0.1,.65))
ax.legend(loc='upper right')
utils.stylizePlot([ax])
ax.axvline(x=1, c='black')
plt.show()

# Summarize the results
print('Not shown, fraction of particles in tail with e>2.0:', 
	np.sum(tailE > 2)/len(tailE))
print('Tail, higher than e>1:', np.sum(tailE > 1)/len(tailE))
print('Stolen, higher than e>1:', np.sum(stolenE > 1)/len(stolenE))
print('Orbitting, higher than e>1:', np.sum(orbittingE > 1)/len(orbittingE))