"""Compares the matter distribution of the merger to the de Vacouleurs' law"""

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import linregress

from analysis import utils

#https://stackoverflow.com/questions/39418380/histogram-with-equal-number-of-points-in-each-bin
def equalNHistogramBins(x, nbin):
    """Calculates the bins necessary for a histogram 
    with equal number of points in each bin
    """
    return np.interp(np.linspace(0, len(x), nbin + 1), np.arange(len(x)), np.sort(x))

def plotRadialMagnitude(data, ax):
    """Plots the magnitude (log(areal density)) as a function of radius"""
    xs, ys = [], []
    
    # Project onto all three planes calculate the areal density and plot the graph
    for dim, style, label in zip([[1,2],[0,2],[0,1]], ['o','s','x'], ['Onto y-z plane', 'Onto x-z plane', 'Onto x-y plane']):
        # Measure radial distance with respect to COM
        COM = np.sum(data['r'] * data['mass'][:,np.newaxis], axis=0) / np.sum(data['mass'])
        r = np.linalg.norm((data['r'] - COM)[:,dim], axis=-1)
        hist, bin_edges = np.histogram(r[data['types']=='disk'], 
            histedges_equalN(r[data['types']=='disk'], 30), density=True)
        bin_edges = (bin_edges[:-1] + bin_edges[1:])/2
        M = np.log(hist) #Magnitude

        if style == 'x':
            plt.scatter((bin_edges**(1/4)), M, marker=style, 
                c='black', label=label)
        else:
            plt.scatter((bin_edges**(1/4)), M, marker=style, 
                facecolors='none', edgecolors='black', label=label)
        
        # Omit points on both ends for the fit, where the law does not hold
        xs.extend((bin_edges**(1/4))[5:-2])
        ys.extend(M[5:-2])
    
    utils.setSize(ax, x=(.5, 1.4), y=(-3.5, 1))
    utils.setAxes(ax, x=r'$r^{1/4}$', y=r'$M$', xcoords=(.9,-0.08), ycoords=(-0.08,.9))
    utils.stylizePlot([ax])
    
    #Fit a line to the data
    slope, intercept, r_value, _, _ = linregress(xs, ys)
    x = np.linspace(.55, 1.35, 10)
    fit = np.poly1d((slope, intercept))(x)
    plt.plot(x, fit, linewidth=.5, label='Global fit', c='black')
    plt.text(.6, -3, r'$R^2 = {:.2f}$'.format(r_value**2), fontsize=14)
    
    plt.legend()
    
data = utils.loadData('hyperbolic_halo', 30000)
f, ax = plt.subplots(1, 1)
plotRadialMagnitude(data, None, ax)