"""Plot a survey of the inclination angle"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

def plotEdgeView(data, ax, axisOn):
    """Plot the encounter as viewed from the edge of the galaxy

    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.
    """
    r_vec = data['r_vec']
    theta, phi = data['CONFIG']['galaxy1']['orientation']

    # Obtain the edge view, rotate 90º around y axis
    spin = np.array([np.sin(theta)*np.cos(phi), 
            np.sin(theta)*np.sin(phi), 
            np.cos(theta)])
    phi = np.pi/2
    M = np.array([[np.cos(-phi), 0,np.sin(-phi)], 
        [0, 1, 0], 
        [-np.sin(-phi), 0, np.cos(-phi)]])
    r_vec = np.einsum('ai,ij->aj', r_vec, M)
    spin = np.einsum('i,ij->j', spin, M)
    data['tracks'] = np.einsum('abi,ij->abj', data['tracks'], M)
    data['tracks'] -= r_vec[0,:] #Relative to center of galaxy
    r_vec[:,:] -= r_vec[0,:] 

    # Flip y to be consistent with Toomre and Toomre
    r_vec[:,1] = -r_vec[:,1]
    data['tracks'][:,:,1] *= -1

    # Plotting
    utils.plotCenterMasses(ax, {'r_vec': r_vec, 'type':data['type']})
    ax.plot(data['tracks'][1,:,0], data['tracks'][1,:,1], 
        c='black', alpha=1.0, linewidth=1)
    ax.scatter(r_vec[:,0], r_vec[:,1], s=0.01, marker='.', 
        c='#c40f4c', alpha=.6)
    utils.setSize(ax, x=(-1.5, 1.5), y=(-1.5,3.5), mode='square')
    utils.stylizePlot([ax])
    utils.setAxes(ax, x=r'x / $r_{min}$', y=r'y / $r_{min}$', 
        ycoords=(-0.2, .8), mode=('bottomleft' if axisOn else 'bottom'))
    
def plotFrontalView(data, ax, axisOn):
    """Plot the encounter as viewed normal to the galactic plane

    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.
    """
    r_vec = data['r_vec']
    theta, phi = data['CONFIG']['galaxy1']['orientation']
    
    # Rotate the position vectors into the appropiate view
    M2 = np.array([[np.cos(-phi), np.sin(-phi), 0],
        [-np.sin(-phi), np.cos(-phi), 0],
        [0, 0, 1]])
    M1 = np.array([[1, 0, 0],
        [0, np.cos(-theta), np.sin(-theta)],
        [0, -np.sin(-theta), np.cos(-theta)]])
    r_vec = np.einsum('ai,ij->aj', np.einsum('ai,ij->aj', r_vec, M2), M1)
    data['tracks'] = np.einsum('abi,ij->abj', data['tracks'], M2)
    data['tracks'] = np.einsum('abi,ij->abj', data['tracks'], M1)
    data['tracks'] -= data['tracks'][0]
    r_vec[:,:] -= r_vec[0,:]

    #Plotting
    utils.plotCenterMasses(ax, {'r_vec': r_vec, 'type':data['type']})
    utils.plotTracks(ax, data['tracks'])
    ax.scatter(r_vec[:,0], r_vec[:,1], c='#c40f4c', s=0.01, marker='.')
    utils.setSize(ax, x=(-3, 2), y=(-1.5, 1.5), mode='square')
    utils.stylizePlot([ax])
    utils.setAxes(ax, x=r'x / $r_{min}$', y=r'y / $r_{min}$', 
        ycoords=(-0.2, .8), mode=('bottomleft' if axisOn else 'bottom'))

#########################################################

# Select the inclinations
inclinations = [15, 45, 60, 90, 135]
f, axs = plt.subplots(2, len(inclinations),  figsize=(11, 6), 
    sharey=False, sharex=False, gridspec_kw = {'height_ratios':[5/3, 1]})
t = 5000 # time of plots

# Plot both views for each encounter
for i in range(len(inclinations)):
    data = utils.loadData('i{}'.format(inclinations[i]), t)
    plotEdgeView(data, ax=axs[0,i], axisOn=True if i==0 else False)
    data = utils.loadData('i{}'.format(inclinations[i]), t)
    plotFrontalView(data, ax=axs[1,i], axisOn=True if i==0 else False)
    
f.subplots_adjust(hspace=0, wspace=0)
plt.tight_layout()
plt.show()