"""Make a plot of the shape of the tails as a function of
mass of the perturbing mass, time and ring size
"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils
from analysis.segmentation import segmentEncounter

def plotShape(ax, distances, angle, style, label):
    """Plot the shape curves, normalized to maximum radius = 1

    Parameters:
        ax: the matplotlib axis where the data should be plotted
        distances (arr): the distances of each point in the shape of the tail
        angle (arr): the angle of each point in the shape curve of the tail
        style (string): style (dotted, dashed...) for the line, passed to
            matplotlib.
        label (string): label of the encounter. Used for the legend.
    """
    ax.plot(-angles+np.min(angles), distances/np.max(distances), 
        label=label, linestyle=style, c='k', alpha=0.7)


f, axs = plt.subplots(1, 3,  figsize=(12, 3), sharey=True)

# Varying the mass of the encounter
for m, st in zip([50, 100, 200], [':', '-','--']):
    data = utils.loadData('{}mass70radius'.format(m), 14000)
    _, (distances, angles), _ = segmentEncounter(data, plot=False)
    plotShape(axs[0], distances, angles, st, m/100)
    axs[0].legend(title=r'$m_{companion}$')
utils.setAxes(axs[0], x=r'Angle / rad', y=r'Radius / $r_{max}$', 
    ycoords=(-0.1,.7))

# Varying the time since the start of the encounter
for t, st in zip([4400, 7400, 10400], [':', '-','--']):
    data = utils.loadData('200mass70radius'.format(200), t)
    _, (distances, angles), _ = segmentEncounter(data, plot=False)
    plotShape(axs[1], distances, angles, st, t)
    axs[1].legend(title=r'$t$')
utils.setAxes(axs[1], x=r'Angle / rad')

# Varying the radius of the ring
for r, st in zip([50, 60, 70, 80, 90], [':', '-.','-','--']):
    data = utils.loadData('200mass{}radius'.format(r), 10000)
    _, (distances, angles), _ = segmentEncounter(data, plot=False)
    plotShape(axs[2], distances, angles, st, r)
    axs[2].legend(title=r'Disk radius')
utils.setAxes(axs[2], x=r'Angle / rad')

# Some styling
utils.stylizePlot(axs)
f.subplots_adjust(hspace=0, wspace=0)
plt.show()