"""Plot the evolution of the amount of matter in the bridge as 
a function of time for varying companion 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 = [20, 50, 100, 200]
labels = ['1:0.25', '1:0.5', '1:1', '1:2', '1:4'] # main:companion

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

for j, style in zip(range(len(masses)), [':', '-.','-','--']):
    ts = np.linspace(4000, 10000, 61) # time range
    fractions = [] # fractional mass in bridge

    for t in ts:
        data = utils.loadData('{}mass'.format(masses[j]), int(t))
        masks, _, _ = segmentEncounter(data, mode='bridge')
        fractions.append(np.sum(masks[0])/len(data['r_vec']))

    ax.plot(ts, fractions, label=labels[j], c='black', linestyle=style)

ax.legend(title='main:companion', fontsize=14)
utils.setAxes(ax, x='Time', y='Fractional mass in bridge', 
	xcoords=(.9,-0.08), ycoords=(-0.08,.6))
utils.setSize(ax, x=(5000,10000), y=(0,0.1))
utils.stylizePlot([ax])
plt.show()