import numpy as np
import scipy.stats as st


# Bulge distributions
class radialPlummer(st.rv_continuous):
    def _pdf(self, x):#(3M/4πa3)(1+(r/a)2)−5/2
        return 3/(4*np.pi**3) * (1 + x**2)**(-5/2) * 4*np.pi*x**2
PLUMMER = radialPlummer(a=0, b=5, name='rPlummer')

# [TODO: Check Hernquist. Not currently in use]
class radialHernquist(st.rv_continuous):
    def _pdf(self, x):
        return 1/(2*np.pi) * 1/(x**4) * 4*np.pi*x**2
HERNQUIST = radialHernquist(a=0, b=5, name='rHernquist')


# Disk distributions
class radialUniform(st.rv_continuous):
    def _pdf(self, x):
        return 2*x if x<1 else 0
UNIFORM = radialUniform(a=0, b=10, name='rUniform')

class radialExp(st.rv_continuous):
    def _pdf(self, x):
        return x*np.exp(-x)
EXP = radialExp(a=0, b=10, name='rExp')


# Halo distributions
class radialNFW(st.rv_continuous):
    def _pdf(self, x):
        y = 1 / (x * (1 + x)**2) * x**2
        # Normalize pdf in (0, 5) range
        y /= (-5/6 + np.log(6))
NFW = radialNFW(a=0, b=5, name='rNFW')