import astropy
from scipy.integrate import quad
if float(astropy.__version__[0]) < 5.0:
    Warning('This routines are only supported for astropy version >=5. Current version is %s.'
            % astropy.__version__)
else:
    from astropy.cosmology.utils import vectorize_redshift_method


class CosmoInterp(object):
    """
    class which interpolates the comoving transfer distance and then computes angular diameter distances from it
    This class is modifying the astropy.cosmology routines
    """
    def __init__(self, cosmo):
        """

        :param cosmo: astropy.cosmology instance (version 4.0 as private functions need to be supported)
        """
        self._cosmo = cosmo

    def _integral_comoving_distance_z1z2(self, z1, z2):
        """
        Comoving line-of-sight distance in Mpc between objects at redshifts
        ``z1`` and ``z2``. The comoving distance along the line-of-sight
        between two objects remains constant with time for objects in the
        Hubble flow.

        Parameters
        ----------
        z1, z2 : Quantity-like ['redshift'] or array-like
            Input redshifts.

        Returns
        -------
        d : `~astropy.units.Quantity` ['length']
            Comoving distance in Mpc between each input redshift.
        """

        return self._cosmo._hubble_distance * self._integral_comoving_distance_z1z2_scalar(z1, z2)

    @vectorize_redshift_method(nin=2)
    def _integral_comoving_distance_z1z2_scalar(self, z1, z2):
        """
        Comoving line-of-sight distance between objects at redshifts ``z1`` and
        ``z2``. Value in Mpc.

        The comoving distance along the line-of-sight between two objects
        remains constant with time for objects in the Hubble flow.

        Parameters
        ----------
        z1, z2 : Quantity-like ['redshift'], array-like, or `~numbers.Number`
            Input redshifts.

        Returns
        -------
        d : float or ndarray
            Comoving distance in Mpc between each input redshift.
            Returns `float` if input scalar, `~numpy.ndarray` otherwise.
        """
        return quad(self._cosmo._inv_efunc_scalar, z1, z2, args=self._cosmo._inv_efunc_scalar_args)[0]
