"""Group together various number density descriptors."""
import random
import numpy as np
import frbpoppy.precalc as pc


class NumberDensity:
    """Class for cosmic number density."""

    def __init__(self,
                 model='vol_co',
                 z_max=6.0,
                 H_0=67.74,
                 W_m=0.3089,
                 W_v=0.6911,
                 alpha=-1.5):
        """Draw from particular number density distributions.

        Args:
            model (str): Which number density to follow.
            z_max (float): Maximum redshift to which to draw.
            vol_co_max (float, optional): Maximum comoving redshift [Gpc^3].
            dist_co_max (float, optional): Maximum comoving distance [Gpc].
            alpha (float, optional): Desired log N log S slope for a perfect,
                non-cosmological population.
        """
        self.z_max = z_max

        # Convert various maximum distance values
        self.dt = pc.DistanceTable(H_0=H_0, W_m=W_m, W_v=W_v).lookup
        m = self.dt(z=np.array([z_max]))
        self.dist_co_max = m[1]
        self.vol_co_max = m[2]
        self.cdf_sfr_max = m[-2]
        self.cdf_smd_max = m[-1]
        self.dist_co_max = self.dist_co_max[0]
        self.vol_co_max = self.vol_co_max[0]
        self.alpha = alpha

        # Determine from which type of distribution to draw
        if model == 'vol_co':
            self.draw = self.from_vol_co
        elif model == 'sfr':
            self.draw = self.from_sfr
        elif model == 'smd':
            self.draw = self.from_smd

        # Allow for the steepness of log N log S to be adapted
        if alpha != -1.5:
            self.draw = self.sloped_dist
            self.power = -self.alpha/1.5
            self.maxi = self.vol_co_max**self.power

    def sloped_dist(self, n_gen=1):
        """Draw from a sloped distribution to create a logNlogS slope."""
        vol_co = (self.maxi*np.random.random(n_gen))**(1/self.power)  # [Gpc]
        d = self.dt(vol_co=vol_co)
        z = d[0]
        dist_co = d[1]
        return z, dist_co

    def from_vol_co(self, n_gen=1):
        """Use constant number density of sources per comoving volume.

        Can be influenced by changing alpha.
        """
        vol_co = self.vol_co_max*np.random.random(n_gen)  # [Gpc]
        d = self.dt(vol_co=vol_co)
        z = d[0]
        dist_co = d[1]
        return z, dist_co

    def from_sfr(self, n_gen=1):
        """Get sources to follow star forming rate.

        Return a random redshift for sources following the Star Formation Rate.

        Follows Madau & Dickinson (2014), eq. 15. For more info see
        https://arxiv.org/pdf/1403.0007.pdf
        """
        sampling = np.random.uniform(0., self.cdf_sfr_max, size=n_gen)
        d = self.dt(cdf_sfr=sampling)
        z = d[0]
        dist_co = d[1]
        return z, dist_co

    def from_smd(self, n_gen=1):
        """
        Return a random redshift for sources following Stellar Mass Density.

        Follows Madau & Dickinson (2014), eq. 2 & 15. For more info see
        https://arxiv.org/pdf/1403.0007.pdf
        """
        sampling = np.random.uniform(0., self.cdf_smd_max, size=n_gen)
        d = self.dt(cdf_smd=sampling)
        z = d[0]
        dist_co = d[1]
        return z, dist_co

# def z_from_sfr(z_max=2.5, n_gen=1):
#     """
#     Return a random redshift for sources following the Star Formation Rate.
#
#     Follows Madau & Dickinson (2014), eq. 15. For more info see
#     https://arxiv.org/pdf/1403.0007.pdf
#
#     """
#     dt = pc.DistanceTable(H_0=H_0, W_m=W_m, W_v=W_v).lookup
#
#     def sfr(z):
#         _, _, _, dvol_co = dt(z=z)
#         return (1+z)**2.7/(1+((1+z)/2.9)**5.6) * dvol_co
#
#     def sample(n_gen):
#         return np.random.uniform(0, z_max, (n_gen,))
#
#     def accept(x):
#         return np.random.rand(x.size)*9.0 <= sfr(x)
#
#     z = sample(n_gen)
#     mask = accept(z)
#     reject, = np.where(~mask)
#     while reject.size > 0:
#         fill = sample(reject.size)
#         mask = accept(fill)
#         z[reject[mask]] = fill[mask]
#         reject = reject[~mask]
#
#     return z
#
#
# def z_from_csmd(z_max=6.0, n_gen=1):
#     """
#     Return a random redshift for sources following the Stellar Mass Density.
#
#     Follows Madau & Dickinson (2014), eq. 2 & 15. For more info see
#     https://arxiv.org/pdf/1403.0007.pdf
#
#     """
#     csmd = pc.CSMDTable().lookup
#
#     def sample(n_gen):
#         return np.random.uniform(0, z_max, (n_gen,))
#
#     def accept(x):
#         return np.random.rand(x.size)*0.00065 <= csmd(x)
#
#     # Accept - rejct algorithm
#     z = sample(n_gen)
#     mask = accept(z)
#     reject, = np.where(~mask)
#     while reject.size > 0:
#         fill = sample(reject.size)
#         mask = accept(fill)
#         z[reject[mask]] = fill[mask]
#         reject = reject[~mask]
#
#     return z

@@ -3311,1612 +3311,4 @@
 _co%0A
-%0A# def z_from_sfr(z_max=2.5, n_gen=1):%0A#     %22%22%22%0A#     Return a random redshift for sources following the Star Formation Rate.%0A#%0A#     Follows Madau & Dickinson (2014), eq. 15. For more info see%0A#     https://arxiv.org/pdf/1403.0007.pdf%0A#%0A#     %22%22%22%0A#     dt = pc.DistanceTable(H_0=H_0, W_m=W_m, W_v=W_v).lookup%0A#%0A#     def sfr(z):%0A#         _, _, _, dvol_co = dt(z=z)%0A#         return (1+z)**2.7/(1+((1+z)/2.9)**5.6) * dvol_co%0A#%0A#     def sample(n_gen):%0A#         return np.random.uniform(0, z_max, (n_gen,))%0A#%0A#     def accept(x):%0A#         return np.random.rand(x.size)*9.0 %3C= sfr(x)%0A#%0A#     z = sample(n_gen)%0A#     mask = accept(z)%0A#     reject, = np.where(~mask)%0A#     while reject.size %3E 0:%0A#         fill = sample(reject.size)%0A#         mask = accept(fill)%0A#         z%5Breject%5Bmask%5D%5D = fill%5Bmask%5D%0A#         reject = reject%5B~mask%5D%0A#%0A#     return z%0A#%0A#%0A# def z_from_csmd(z_max=6.0, n_gen=1):%0A#     %22%22%22%0A#     Return a random redshift for sources following the Stellar Mass Density.%0A#%0A#     Follows Madau & Dickinson (2014), eq. 2 & 15. For more info see%0A#     https://arxiv.org/pdf/1403.0007.pdf%0A#%0A#     %22%22%22%0A#     csmd = pc.CSMDTable().lookup%0A#%0A#     def sample(n_gen):%0A#         return np.random.uniform(0, z_max, (n_gen,))%0A#%0A#     def accept(x):%0A#         return np.random.rand(x.size)*0.00065 %3C= csmd(x)%0A#%0A#     # Accept - rejct algorithm%0A#     z = sample(n_gen)%0A#     mask = accept(z)%0A#     reject, = np.where(~mask)%0A#     while reject.size %3E 0:%0A#         fill = sample(reject.size)%0A#         mask = accept(fill)%0A#         z%5Breject%5Bmask%5D%5D = fill%5Bmask%5D%0A#         reject = reject%5B~mask%5D%0A#%0A#     return z%0A

Remove commented code