"""Plot the encounters for the halo analysis."""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

def plotDisk(data, ax, axisOn, zoomOut=False):
    """Plot the disk of an encounter

    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.
        zoomOut (bool): scale of the plotting
    """
    r_vec = data['r_vec']
    theta, phi = data['CONFIG']['galaxy1']['orientation']
    
    utils.plotTracks(ax, data['tracks'])
    utils.plotCenterMasses(ax, data)
    mask = data['types'][:,0]=='disk'
    ax.scatter(r[mask][:,0], r[mask][:,1], c='#c40f4c', s=2, marker='.')
    if zoomOut: utils.setSize(x=(-10, 10), y=(-10, 10), mode='square')
    else: utils.setSize(x=(-5, 5), y=(-5, 5), mode='square')
    
    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])

# Plot all four encounters described in the report at 3 different times
f, axs = plt.subplots(2, 3,  figsize=(15, 10), sharey=False, sharex=False)

data = utils.loadData('parabolic_nohalo', 6800)
plotDisk(data, ax=axs[0,0], axisOn=True)
data = utils.loadData('parabolic_nohalo', 8800)
plotDisk(data, ax=axs[0,1], axisOn=False)
data = utils.loadData('parabolic_nohalo', 24800)
plotDisk(data, ax=axs[0,2], axisOn=True, zoomOut=True)

data = utils.loadData('parabolic_halo', 9400)
plotDisk(data, ax=axs[1,0], axisOn=True)
data = utils.loadData('parabolic_halo', 12800)
plotDisk(data, ax=axs[1,1], axisOn=False)
data = utils.loadData('parabolic_halo', 20800)
plotDisk(data, ax=axs[1,2], axisOn=True, zoomOut=True)
    
f.subplots_adjust(hspace=0, wspace=0)
plt.tight_layout()
plt.show()