"""Plot a profile of the radial density"""

import matplotlib.pyplot as plt
import numpy as np

from analysis import utils

def plotRadialDensity(data):
    """Plots a profile of the radial density, divided into components.

    Parameters:
        data: a data dictionary in the format saved by the simulation
    """
    f, ax = plt.subplots(1, 1)

    # We measure the radial distance from the Center Of Mass
    COM = np.sum(data['r_vec'] * data['mass'][:,np.newaxis], axis=0) / np.sum(data['mass'])
    r = np.linalg.norm((data['r'] - COM), axis=-1)

    # Divide the particles into groups
    diskMask = data['types'][:,1] == 'disk'
    bulgeMask = data['types'][:,1] == 'bulge'
    luminousMask = np.logical_or(diskMask, bulgeMask)
    darkMask = data['types'][:,1] == 'dark'

    # Plot each group separately
    style = {'histtype':'step', 'cumulative':True, 'bins':np.logspace(-2, 2, 1000),
        'color':'black', 'linestyle':'dashed', 'label':'disk'}
    plt.hist(r[diskMask], weights=data['mass'][diskMask], 
        label='disk', **style)
    plt.hist(r[bulgeMask], weights=data['mass'][bulgeMask], 
        label='bulge', **style)
    plt.hist(r[luminousMask], weights=data['mass'][luminousMask], 
        label='luminous', **style)
    plt.hist(r[darkMask], weights=data['mass'][darkMask], 
        label='dark', **style)
    
    # Plotting styling
    plt.legend(loc='upper left')
    plt.yscale('log'); plt.xscale('log')
    utils.setAxes(ax, x=(3E-2, 1E2), y=(1E-2, 2))
    utils.setAxes(ax, x=r'$r$', y=r'Cumulative mass', 
        xcoords=(.9,-0.08), ycoords=(-0.1,.7))
    utils.stylizePlot([ax])
    plt.show()
    
data = utils.loadData('hyperbolic_halo', 30000)
plotRadialDensity(data)