from abc import ABC, abstractmethod
import numpy as np
from pyHalo.utilities import ITSampling
from scipy.integrate import quad

_log10_rperi_bins = np.array([-2.   , -1.885, -1.77 , -1.655, -1.54 , -1.425, -1.31 , -1.195,
       -1.08 , -0.965, -0.85 , -0.735, -0.62 , -0.505, -0.39 , -0.275,
       -0.16 , -0.045,  0.07 ,  0.185])
_log10_prob = np.array([  18,   23,   37,   33,   42,   78,  139,  181,  321,  538,  749,
       1018, 1268, 1501, 1577, 1232,  569,  207,   23,    7])
_cdf = np.cumsum(_log10_prob)
_log10_rpericenter_sampling = ITSampling(_log10_rperi_bins, _cdf)

class Halo(ABC):

    def __init__(self, mass=None, x=None, y=None, r3d=None, mdef=None, z=None,
                 sub_flag=None, lens_cosmo_instance=None, args={}, unique_tag=None, fixed_position=False):

        """
        This is the main class for objects rendered in the lens volume. It keeps track of stuff like the position,
        mass, redshift, and structural properties (e.g. concentration, core radius, etc.)

        :param mass: halo mass in M_sun
        :param x: angular coordinate x in arcsec
        :param y: angular coordinate y in arcsec
        :param r3d: 3D position of halo in kpc (used to compute the truncation radius for subhalos)
        For field halos this is not relevant and is set to None
        :param mdef: mass definition for the halo
        :param z: halo redshift
        :param sub_flag: bool; if True, the halo is treated as a main deflector subhalo
        :param lens_cosmo_instance: an instance of LensCosmo
        :param args: keyword arguments that include default settings for the halo
        :param unique_tag: a random number with 16 decimal places that uniquely identifies each halo
        :param fixed_position: determiens whether halos can be moved around when aligning a realization with
        the rendering volume
        """

        self.lens_cosmo = lens_cosmo_instance
        self.mass = mass
        # x and y in arcsec
        self.x = x
        self.y = y
        self.r3d = r3d
        self.mdef = mdef
        self.z = z
        self.is_subhalo = sub_flag
        self._args = args
        self.unique_tag = unique_tag
        self._rescale_norm = 1.
        self.fixed_position = fixed_position
        self._rescaled_once = False
        assert z > 0, 'Negative redshifts are unphysical for halos'

    def rescale_normalization(self, factor):
        """
        Sets the rescaling factor for the normalization (only can do this once)
        :param factor:
        :return:
        """
        if self._rescaled_once:
            pass
        else:
            self._rescaled_once = True
            self._rescale_norm = factor
            if hasattr(self, '_params_physical'):
                delattr(self, '_params_physical')
            if hasattr(self, '_kwargs_lenstronomy'):
                delattr(self, '_kwargs_lenstronomy')
            if hasattr(self, '_nfw_params'):
                delattr(self, '_nfw_params')

    @property
    @abstractmethod
    def profile_args(self):
        """
        This routine computes properties of the halo required to specify it
        """
        ...

    @property
    def params_physical(self):
        raise Exception('this halo class does not have attribute params_physical')

    @property
    @abstractmethod
    def lenstronomy_ID(self):
        """
        Returns a list of profile names recognized by lenstronomy

        Example:
            a truncated NFW profile would return ['TNFW']
            a hybrid NFW and point mass profile would return ['NFW', 'POINT_MASS']

        """
        ...

    @property
    @abstractmethod
    def lenstronomy_params(self):
        """
        Returns a list of keyword arguments for the profile, must be the same length as lenstronomy_ID
        :return:
        """
        ...

    @property
    def z_infall(self):

        """
        Evaluate the infall redshift using a PDF generated by galacticus.
        Note: This routine is meaningless and therefore not used for LOS halos
        :return: the infall redshift of a halo assuming it is in a host halo at redshift self.z
        """

        if not hasattr(self, '_z_infall'):
            self._z_infall = self.lens_cosmo.z_accreted_from_zlens(self.mass, self.z)
            assert self._z_infall >= self.z, 'infall redshifts less than current redshift are unphysical'
        return self._z_infall

    @property
    def time_since_infall(self):
        """
        This routine calculates the time in Gyr since infall for subhalos using the infall redshift as predicted by
        Galacticus
        :return: the time since the subhalo was accreted onto thehost [Gyr]
        """
        if not self.is_subhalo:
            print("time since infall is a meaningless concept for field halos")
            return None
        if not hasattr(self, '_time_since_infall'):
            astropy = self.lens_cosmo.cosmo.astropy
            self._time_since_infall = astropy.age(self.z).value - astropy.age(self.z_infall).value
            assert self._time_since_infall > 0
        return self._time_since_infall

    @property
    def rperi_units_r200(self):
        """
        Returns the orbital pericenter of a subhalo in units of the host halo virial radius. This method
        uses output from the semi-analytic model Galacticus
        :return:
        """
        if not self.is_subhalo:
            print("Orbital pericenter is a meaningless concept for field halos. It is possible you assigned a tidal "
                  "truncation model that requires this information to field halos.")
            return None
        if not hasattr(self, '_rperi_units_r200'):
            self._rperi_units_r200 = 10**float(_log10_rpericenter_sampling(n_samples=1.0))
        return self._rperi_units_r200

    @property
    def bound_mass(self):
        raise Exception('this halo class does not have a bound mass attribute because the profile does not have '
                        'a tidal truncation radius')

    @property
    def halo_age(self):
        """Computes the age of a halo assuming it collapsed at z=10
        TODO: Make this a function of halo mass
        """
        if not hasattr(self, '_halo_age'):
            self._halo_age = self.lens_cosmo.cosmo.halo_age(self.z, zform=10)
        return self._halo_age

    def density_profile_3d(self, *args, **kwargs):
        raise Exception('This profile does not have a density_profile_3d method defined in the class')

    @property
    def c(self):
        """
        Computes the halo concentration (once)
        """

        raise Exception('this class does not have a well-defined concentration parameter')

    @property
    def z_eval(self):
        """
        Returns the redshift at which to evaluate the concentration-mass relation
        """
        if not hasattr(self, '_zeval'):

            if self.is_subhalo:
                if 'evaluate_mc_at_zlens' in self._args.keys() and self._args['evaluate_mc_at_zlens']:
                    self._zeval = self.z
                else:
                    self._zeval = self.z_infall
            else:
                self._zeval = self.z

        return self._zeval

    @property
    def nfw_params(self):
        """
        Computes the nfw profile parameters (rs,rho_s) from mass and concentration
        :return: rs, r200 and rho_s in units kpc, kpc, and M_sun / kpc^3
        """
        if not hasattr(self, '_nfw_params'):
            if self.mdef in ['TNFWC', 'GNFW']:
                pseudo_nfw = True
            else:
                pseudo_nfw = False
            rhos, rs, r200 = self.lens_cosmo.NFW_params_physical(self.mass, self.c, self.z, pseudo_nfw)
            self._nfw_params = [rhos, rs, r200]
        return self._nfw_params[0], self._nfw_params[1], self._nfw_params[2]

    def mass_3d(self, rmax, profile_args=None):
        """
        Computes the mass enclosed inside a sphere of radius rmax
        :param rmax: the maximum radius
        :return: mass enclosed inside rmax
        """
        if rmax == 'r200':
            rmax = self.nfw_params[2]
        _integrand = lambda r: 4 * np.pi * r ** 2 * self.density_profile_3d(r, profile_args)
        return quad(_integrand, 0, rmax)[0]
