import numpy as np
from pyHalo.defaults import lenscone_default
from scipy.interpolate import interp1d
from pyHalo.Cosmology.cosmology import Cosmology
from scipy.integrate import quad
from pyHalo.Halos.lens_cosmo import LensCosmo
from scipy.integrate import simpson as simps
from pyHalo.concentration_models import preset_concentration_models
from lenstronomy.LensModel.lens_model import LensModel

class ITSampling(object):

    def __init__(self, x, cdf):
        """
        This class performs inverse transform sampling of a distribution given samples from the distribution
        :param samples: samples from the distribution
        """
        cdf = cdf / float(np.max(cdf))
        self._cdf_inverse = interp1d(cdf, x)
        self._umin = cdf[0]
        self._umax = cdf[-1]

    @classmethod
    def from_samples(cls, samples):
        """
        samples from the target distribution
        :param samples:
        :return:
        """
        ran = (np.min(samples), np.max(samples))
        h, x = np.histogram(samples, range=ran, bins=150)
        x = x[0:-1] + (x[1] - x[0]) / 2
        cdf = np.cumsum(h)
        cdf = cdf / float(np.max(cdf))
        return ITSampling(x, cdf)

    def __call__(self, n_samples):
        """
        Generates samples from the distribution
        :param n_samples: number of samples to draw
        :return: the samples
        """
        u = np.random.uniform(self._umin, self._umax, int(n_samples))
        samples_out = self._cdf_inverse(u)
        if n_samples == 1:
            return float(samples_out)
        else:
            return np.squeeze(samples_out)

def inverse_transform_sampling(x, function, args, n_samples):
    """

    :param x: the domain of the function across which you want to obtain samples
    :param function: the function or probability density you want to sample from
    :param args: arguments passed to function after x
    :param n_samples: number of samples to draw
    :return: samples from the probability density described by function
    """
    y = function(x, *args)
    cdf = np.cumsum(y)
    return inverse_transform_sampling_from_cdf(x, cdf, n_samples)

def inverse_transform_sampling_from_cdf(x, cdf, n_samples):
    """

    :param x: the domain of the function across which you want to obtain samples
    :param cdf: the cumulative distribution function of the pdf
    :param n_samples: number of samples to draw
    :return: samples from the probability density that corresponds to cdf
    """
    cdf = cdf * float(np.max(cdf)) ** -1
    cdf_inverse = interp1d(cdf, x)
    u = np.random.uniform(cdf[0], cdf[-1], n_samples)
    return cdf_inverse(u)

def generate_lens_plane_redshifts(zlens, zsource):
    """
    This routine sets up the redshift planes along the line of sight in the lens system
    :param zlens: main deflector plane redshift (if there is no main lens plane, then this can be set as None)
    :param zsource: source plane redshift
    :return: lens plane redshifts and the thickness of each slice
    """
    zmin = lenscone_default.default_zstart
    zstep = lenscone_default.default_z_step
    if zlens is None:
        redshifts = np.arange(zmin, zsource, zstep)
    else:
        front_z = np.arange(zmin, zlens, zstep)
        back_z = np.arange(zlens, zsource, zstep)
        redshifts = np.append(front_z, back_z)
    delta_zs = []
    for i in range(0, len(redshifts) - 1):
        delta_zs.append(redshifts[i + 1] - redshifts[i])
    delta_zs.append(zsource - redshifts[-1])
    return list(np.round(redshifts, 2)), np.round(delta_zs, 2)

def interpolate_ray_paths(x_coordinates, y_coordinates, lens_model, kwargs_lens, zsource,
                          terminate_at_source=False, source_x=None, source_y=None, evaluate_at_mean=False,
                          cosmo=None):

    """
    :param x_coordinates: x coordinates to interpolate (arcsec) (list)
    :param y_coordinates: y coordinates to interpolate (arcsec) (list)
    Typically x_coordinates/y_coordinates would be four image positions, or the coordinate of the lens centroid
    :param lens_model: instance of LensModel (lenstronomy)
    :param kwargs_lens: keyword arguments for lens model
    :param zsource: source redshift
    :param terminate_at_source: fix the final angular coordinate to the source coordinate
    :param source_x: source x coordinate (arcsec)
    :param source_y: source y coordinate (arcsec)
    :param evaluate_at_mean: if True, returns two single interp1d instances (one for each of x/y) that return the
    average of each individual x/y coordinate evaluated at each lens plane. For example, if you pass in four images positions
    the output would be an interpolation of the average x/y coordinate along the path traversed by the light
    (This is useful for aligning realizations with a background source significantly offset from the lens centroid)

    :return: Instances of interp1d (scipy) that return the angular coordinate of a ray given a
    comoving distance
    """

    angle_x = []
    angle_y = []

    if cosmo is None:
        cosmo = Cosmology()

    for i, (xpos, ypos) in enumerate(zip(x_coordinates, y_coordinates)):

        theta_x = [xpos]
        theta_y = [ypos]

        ray_x, ray_y, d = compute_comoving_ray_path(xpos, ypos, lens_model, kwargs_lens, zsource,
                                                          terminate_at_source, source_x, source_y, cosmo=cosmo)

        for rx, ry, di in zip(ray_x[1:], ray_y[1:], d[1:]):
            theta_x.append(rx / di)
            theta_y.append(ry / di)

        distances = [0.] + list(d[1:])
        distances = np.array(distances)
        theta_x = np.array(theta_x)
        theta_y = np.array(theta_y)

        angle_x.append(interp1d(distances, theta_x))
        angle_y.append(interp1d(distances, theta_y))

    if evaluate_at_mean:

        zrange = np.linspace(0., zsource, 100)
        comoving_distance_calc = cosmo.D_C_transverse
        distances = [comoving_distance_calc(zi) for zi in zrange]

        angular_coordinates_x = []
        angular_coordinates_y = []
        for di in distances:
            x_coords = [ray_x(di) for ray_x in angle_x]
            y_coords = [ray_y(di) for ray_y in angle_y]
            x_center = np.mean(x_coords)
            y_center = np.mean(y_coords)
            angular_coordinates_x.append(x_center)
            angular_coordinates_y.append(y_center)

        angle_x = [interp1d(distances, angular_coordinates_x)]
        angle_y = [interp1d(distances, angular_coordinates_y)]

    return angle_x, angle_y

def compute_comoving_ray_path(x_coordinate, y_coordinate, lens_model, kwargs_lens, zsource,
                              terminate_at_source=False, source_x=None, source_y=None, cosmo=None):

        """
        :param x_coordinate: x coordinates to interpolate (arcsec) (float)
        :param y_coordinate: y coordinates to interpolate (arcsec) (float)
        Typically x_coordinates/y_coordinates would be four image positions, or the coordinate of the lens centroid
        :param lens_model: instance of LensModel (lenstronomy)
        :param kwargs_lens: keyword arguments for lens model
        :param zsource: source redshift
        :param terminate_at_source: fix the final angular coordinate to the source coordinate
        :param source_x: source x coordinate (arcsec)
        :param source_y: source y coordinate (arcsec)
        :return: Instance of interp1d (scipy) that returns the angular coordinate of a ray given a
        comoving distance
        """

        if cosmo is None:
            cosmo = Cosmology()

        redshift_list = lens_model.redshift_list + [zsource]
        zstep = lenscone_default.default_z_step
        finely_sampled_redshifts = np.linspace(zstep, zsource - zstep, 50)
        all_redshifts = np.unique(np.append(redshift_list, finely_sampled_redshifts))

        all_redshifts_sorted = all_redshifts[np.argsort(all_redshifts)]

        comoving_distance_calc = cosmo.D_C_transverse

        x_start, y_start = 0., 0.
        z_start = 0.

        x_list = [0.]
        y_list = [0.]
        distances = [0.]

        alpha_x_start, alpha_y_start = x_coordinate, y_coordinate
        for zi in all_redshifts_sorted:

            x_start, y_start, alpha_x_start, alpha_y_start = lens_model.lens_model.ray_shooting_partial(x_start, y_start,
                                                                                alpha_x_start, alpha_y_start,
                                                                                z_start, zi, kwargs_lens)
            d = float(comoving_distance_calc(zi))
            x_list.append(x_start)
            y_list.append(y_start)
            distances.append(d)
            z_start = zi

        if terminate_at_source:
            d_src = comoving_distance_calc(zsource)
            x_list[-1] = source_x * d_src
            y_list[-1] = source_y * d_src

        return np.array(x_list), np.array(y_list), np.array(distances)

def de_broglie_wavelength(log10_m_uldm,v):
    '''
    Returns de Broglie wavelength of the ultra-light axion in kpc.

    :param log10_m_uldm: log(axion mass) in eV
    :param v: velocity in km/s
    '''
    m_axion=10**log10_m_uldm
    return 1.2*(1e-22/m_axion)*(100/v)

def delta_kappa(z_lens, z_source, m, rein, de_Broglie_wavelength):
    '''
    Returns standard deviation of the density fluctuations in projection in convergence units

    :param z_lens,z_source: lens and source redshifts
    :param m: main deflector halo mass in M_solar
    :param rein: Einstein radius in kpc
    :param de_Broglie_wavelength: de Broglie wavelength of axion in kpc
    '''
    l = LensCosmo(z_lens,z_source)
    sigma_crit = l.get_sigma_crit_lensing(z_lens, z_source) * (1e-3) ** 2
    ds = delta_sigma(m, z_lens, rein, de_Broglie_wavelength)
    delta_kappa = ds / sigma_crit
    return delta_kappa

def delta_sigma(m, z, rein, de_Broglie_wavelength):
    """
    Returns the mean ULDM fluctuation ampltiude of the host dark matter halo in units M_sun/kpc^2
    :param m:
    :param z:
    :param rein:
    :param de_Broglie_wavelength:
    :return:
    """
    cosmo = Cosmology()
    l = LensCosmo(z, 2.0, cosmo)
    model, _ = preset_concentration_models('DIEMERJOYCE19')
    cmodel = model(cosmo.astropy, scatter=False)
    c = cmodel.nfw_concentration(m, z)
    rhos, rs, _ = l.NFW_params_physical(m, c, z)
    nfw_rho_squared = projected_density_squared(rein, rhos, rs, c)
    delta_sigma = (np.sqrt(np.pi) * nfw_rho_squared * de_Broglie_wavelength)**0.5
    return delta_sigma

def projected_density_squared(R_ein, rhos, rs, concentration):

    '''
    Returns the integral along the line of sight of the square of the NFW profile
    i.e. integral rho_nfw(z)^2 dz

    :param R_ein: Einstein radius in kpc
    :param rhos: scale radius density of main deflector halo
    :param rs: scale radius of main deflector halo
    :param concentration: concentration of main deflector halo
    '''

    r200 = concentration * rs
    zmax = np.sqrt(r200 ** 2 - R_ein**2)

    x = lambda z: np.sqrt(R_ein ** 2 + z ** 2)/rs
    nfw_density_square = lambda z: rhos**2 / (x(z) * (1+x(z))**2)**2

    return 2 * quad(nfw_density_square, 0, zmax)[0]

def nfw_velocity_dispersion(rhos, rs, c, x=1):
    """
    Computes the central velocity dispersion of an NFW profile integrating from r = x * rs
    :param rhos:
    :param rs:
    :param c:
    :return:
    """
    G = 4.3e-6  # kpc/solMass * (km/sec)**2
    prefactor = 4 * np.pi * G * rhos ** 2 * rs ** 2
    density_at_r = rhos / (x * (1 + x) ** 2)
    _integrand = lambda x: (np.log(1 + x) - x / (1 + x)) / (x * (1 + x) ** 2) / x ** 2
    x = np.linspace(x, c, 150)
    y = _integrand(x)
    return np.sqrt(prefactor * simps(y, x) / density_at_r)

class MinHaloMassULDM(object):

    def __init__(self, log10_m_uldm, astropy_instance, log_mlow):
        """
        This class implements the minimum halo mass for ultra-light dark matter, given the particle mass and a cosmology

        The call method returns the maximum of log_mlow, and the resulting minimum ULDM halo mass to ensure that one does
        not inadvertently generate halos down to extremely low masses for heavy particles
        :param log10_m_uldm: particle mass
        :param astropy_instance: an instannce of astropy
        :param log_mlow: minimum halo mass to render, regardless of ULDM computation
        """
        m22 = 10**(log10_m_uldm + 22)
        self._Mmin0 = 4.4e7 * m22**(-3/2)
        self._astropy_instance = astropy_instance
        self._log_mlow = log_mlow

    def __call__(self, z):
        m_min = self.m_min(z)
        log10_m_min = np.log10(m_min)
        return max(log10_m_min, self._log_mlow)

    def m_min(self, z):
        return self._a(z) ** (-3 / 4) * (self._zeta(z) / self._zeta(0)) ** (1 / 4) * self._Mmin0

    def _a(self, z):
        return (1 + z) ** -1

    def _Om(self, z):
        return self._astropy_instance.Om(z)

    def _zeta(self, z):
        return (18 * np.pi ** 2 + 82 * (self._Om(z) - 1) - 39 * (self._Om(z) - 1) ** 2) / self._Om(z)

def multiplane_convergence(realization, cone_opening_angle_arcsec=2.5, npix=100, lens_model_list_macro=None,
                           kwargs_lens_macro=None, redshift_list_macro=None):
    """
    This function computes the effective multiplane convergence from a realization of dark matter halos.
    If no macromodel is specified (through lens_model_list_macro, kwargs_lens_macro, redshift_list_macro) then
    a typical one is assumed.

    :param realization: an instance of Realization
    :param cone_opening_angle_arcsec: the window size in arcsec (each axis spans 2*window_size)
    :param npix: number of pixels per axis
    :param lens_model_list_macro: list of lens models for the main deflector
    :param kwargs_lens_macro: keyword arguments for the macro lens models
    :param redshift_list_macro: list of redshifts for the macro lens models
    :return: the effective multiplane convergence, the full lens model used to create it, the keyword arguments for it,
    the macromodel lens model, and the keyword arguments for it
    """

    lens_model_list_halos, redshfit_array_halos, kwargs_lens_halos, _ = realization.lensing_quantities()
    if lens_model_list_macro is None:
        lens_model_list_macro = ['EPL', 'SHEAR']
        kwargs_lens_macro = [{'theta_E': 1.0, 'center_x': 0.0, 'center_y': 0.0, 'e1': 0.1, 'e2': -0.1,
                              'gamma': 2.0}, {'gamma1': 0.05, 'gamma2': 0.02}]
        redshift_list_macro = [realization.lens_cosmo.z_lens] * len(lens_model_list_macro)

    lens_model_macro = LensModel(lens_model_list_macro)
    lens_model = LensModel(lens_model_list_macro + lens_model_list_halos,
                           z_source=realization.lens_cosmo.z_source, multi_plane=True,
                           lens_redshift_list=redshift_list_macro + list(redshfit_array_halos),
                           cosmo=realization.lens_cosmo.cosmo.astropy)

    _r = np.linspace(-cone_opening_angle_arcsec/2, cone_opening_angle_arcsec/2, npix)
    xx, yy = np.meshgrid(_r, _r)
    kappa_macro = lens_model_macro.kappa(xx.ravel(), yy.ravel(), kwargs_lens_macro)
    kappa = lens_model.kappa(xx.ravel(), yy.ravel(), kwargs_lens_macro + kwargs_lens_halos)
    delta_kappa = (kappa - kappa_macro).reshape(npix, npix)
    return delta_kappa, lens_model, kwargs_lens_macro + kwargs_lens_halos, lens_model_macro, kwargs_lens_macro


def mask_annular(center_x, center_y, x_grid, y_grid, r_min, r_max = None):
    """
    An annular mask of specified inner and outer radii

    :param center_x: x-coordinate of center position of circular mask
    :param center_y: y-coordinate of center position of circular mask
    :param x_grid: x-coordinate grid
    :param y_grid: y-coordinate grid
    :param r_min: inner radius of mask in units of grid coordinates
    :param r_max: outer radius of mask in units of grid coordinates
    :return: mask array of shape x_grid with =0 inside the radius and =1 outside
    """
    x_shift = x_grid - center_x
    y_shift = y_grid - center_y
    R = np.sqrt(x_shift*x_shift + y_shift*y_shift)
    mask = np.ones(x_shift.shape)
    if r_max == None:
        mask[R <= r_min] = 0
    else:
        mask[R <= r_min] = 0
        mask[R >= r_max] = 0

    return mask
