"""Plot the trajectories followed by the central masses
to illustrate orbital decay due to dark matter halos"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

def plotTrajectory(data, ax, axisOn, trajectory):
    """Plot the trajectory followed by the central masses

    Parameters:
        ax: the matplotlib axis where the data should be plotted
        data: a data dictionary in the format saved by the simulation
        axisOn (bool): whether to include the y axis in this subplot.
            It can be omitted when it is shared among subplots.
        trajectory (string): type of theoretical trajectory 
            ('ellipse' or 'hyperbola') to superpose
    """
    # Plot trajectory   
    ax.plot(np.array(data['tracks'])[0,:,0], np.array(data['tracks'])[0,:,1],
        c='black', alpha=1.0, linewidth=1)
    ax.plot(np.array(data['tracks'])[1,:,0], np.array(data['tracks'])[1,:,1],
        c='black', alpha=0.6, linewidth=1)
    ax.axis('square')

    # Superpose theoretical trajectory
    if trajectory=='ellipse':
        traj = patches.Arc((-.5,0), width=2, height=np.sqrt(3), 
            theta1=-180, theta2=100, linewidth=2, edgecolor='black', 
            facecolor='none', alpha=.7, linestyle='dashed')
        ax.add_patch(traj)
        utils.setSize(ax, x=(-1.5, 1.5), y=(-1.5, 1.5))
    elif trajectory=='ellipse': #x = 1/2 - y^2/2
        y = np.linspace(-2.5, 3, 1000)
        x = 1/2 - y**2/2
        ax.plot(x, y, linewidth=2, color='black', 
            alpha=.7, linestyle='dashed')
        utils.setSize(ax, x=(-1.5, 1.5), y=(-1.5, 1.5))
    else: raise Exception('Unknown trajectory')
    
    # Styling
    utils.setAxes(ax, x=r'x\' / $r_{min}$', y=r'y\' / $r_{min}$', 
        ycoords=(-0.2, .8), mode=('bottomleft' if axisOn else 'bottom'))
    utils.stylizePlot([ax])
 
f, axs = plt.subplots(1, 4,  figsize=(15, 4), sharey=False, sharex=False)

# Plot the four encounters described in the report
data = utils.loadData('elliptical_nohalo', 30000)
plotTrajectory(data, ax=axs[0], axisOn=True, trajectory='ellipse')
data = utils.loadData('elliptical_halo', 30000)
plotTrajectory(data, ax=axs[1], axisOn=True, trajectory='ellipse')
data = utils.loadData('hyperbolic_nohalo', 50000)
plotTrajectory(data, ax=axs[2], axisOn=True, trajectory='ellipse')
data = utils.loadData('hyperbolic_halo', 100000)
plotTrajectory(data, ax=axs[3], axisOn=True, trajectory='ellipse')

f.subplots_adjust(hspace=0, wspace=0)
plt.tight_layout()
plt.show()