"""Analysis of the transfer of angular momentum from luminous components
for encounters with a dark matter halo"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

def luminousAngularMomentum(data):
    """Calculate angular momentum of the luminous components

    Parameters:
        data: a data dictionary in the format saved by the simulation

    Returns:
        The norm of the angular momentum vector of the luminous components
    """
    maskDisk = data['type'][:,0]=='disk'
    maskBulge = data['type'][:,0]=='bulge'
    # Plot luminous components
    mask = np.logical_or(maskDisk, maskBulge)
    return np.linalg.norm(np.sum(data['mass'][mask][:,np.newaxis] 
        * np.cross(data['r_vec'][mask], data['v_vec'][mask]), axis=0))

# Plotting
f, ax = plt.subplots(1, 1)

# No halo
j1 = []
ts = np.linspace(1, 50001, 251)
for t in ts:
    data = utils.loadData('hyperbolic_nohalo', int(t))
    j1.append(luminousAngularMomentum(data))
ax.plot(ts, j1/j1[0], c='black', linestyle='dashed', label='1:1:0')

# Halo
j2 = []
ts = np.linspace(1, 150001, 751)
for t in ts:
    data = utils.loadData('hyperbolic_halo', int(t))
    j2.append(luminousAngularMomentum(data))
ax.plot(ts, j2/j2[0], c='black', linestyle='solid', label='1:3:16')

# Plotting styling
utils.setSize(ax, y=(0, 1.1))
utils.setAxes(ax, x=r'$t$', y=r'Angular momentum (luminous)', 
    ycoords=(.9,-0.08))
ax.legend(title=r'bulge:disk:halo')
utils.stylizePlot([ax])
plt.show()

# Quantitaive analysis
print('Angular momentum is conserved in encounter 1 to within', 
    (np.max(j1)-np.min(j1))/j1[0])