"""Plot the evolution of the length of the tail as a function
of time for different companion masses
"""

import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import savgol_filter

from analysis import utils
from analysis.segmentation import segmentEncounter

# Masses are given relative to 100, with respect to the main mass
masses = [50, 100, 200, 400]
labels = ['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
    lengths = [] # length of the tail

    for t in ts:
        data = utils.loadData('{}mass'.format(masses[j]), int(t))
        _, _, length = segmentEncounter(data)
        lengths.append(length)

    # For some masses we manually omit some misbehaved points
    if masses[j]==50: ax.plot(ts[20:], savgol_filter(lengths[20:], 21, 3), 
    	label=labels[j], c='black', linestyle=style)
    else: ax.plot(ts, savgol_filter(lengths, 21, 3), 
    	label=labels[j], c='black', linestyle=style)

ax.legend(title='main:companion', fontsize=14)
utils.setAxes(ax, x='Time', y='Tail length', 
	xcoords=(.9,-0.08), ycoords=(-0.08,.7))
utils.setSize(ax, x=(5000,10000))
utils.stylizePlot([ax])
plt.show()