"""Plot the mass that, at large times, is stolen, in the tail
or still orbitting the main mass for different perturbing masses
"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils
from analysis.segmentation import segmentEncounter

# Masses are given relative to 100, with respect to the main mass
masses = np.array([20, 33, 50, 75, 100, 150, 200, 300, 400, 500, 600])
tailFractions, stolenFractions = [], []

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

# For each encounter calculate the fraction of stolen and tail mass
for mass in masses:
    data = utils.loadData('{}mass'.format(mass), 10000)

    masks, _, _ = segmentEncounter(data)
    tailFractions.append(len(masks[3])/len(data['r_vec']))
    stolenFractions.append(len(masks[1])/len(data['r_vec']))

# Plotting
ax.plot(masses/100, tailFractions, marker='+', color='black', 
	markersize=10, label='Tail')
ax.plot(masses/100, stolenFractions, marker='o', color='black', 
	markersize=8, linestyle='dashed', label='Stolen', markerfacecolor='none')

ax.legend()
utils.setAxes(ax, x='Companion mass / main mass', y='Mass fraction', 
	xcoords=(.7,-0.08) ,ycoords=(-0.08,.8))
utils.stylizePlot([ax])
ax.set_ylim(0,.45)
plt.show()