import numpy as np
from lenstronomy.GalKin.galkin import Galkin
from lenstronomy.Util import util
from lenstronomy.Util import mask_util
from scipy import signal


class GalkinShells(Galkin):
    """
    class to calculate velocity dispersion for radial shells in a fast way
    """

    def __init__(self, kwargs_model, kwargs_aperture, kwargs_psf, kwargs_cosmo, kwargs_numerics=None,
                 analytic_kinematics=False):
        """

        :param kwargs_model: keyword arguments describing the model components
        :param kwargs_aperture: keyword arguments describing the spectroscopic aperture, see Aperture() class
        :param kwargs_psf: keyword argument specifying the PSF of the observation
        :param kwargs_cosmo: keyword arguments that define the cosmology in terms of the angular diameter distances
         involved
        :param kwargs_numerics: numerics keyword arguments
        :param analytic_kinematics: bool, if True uses the analytic kinematic model
        """
        Galkin.__init__(self, kwargs_model=kwargs_model, kwargs_aperture=kwargs_aperture, kwargs_psf=kwargs_psf,
                        kwargs_cosmo=kwargs_cosmo, kwargs_numerics=kwargs_numerics,
                        analytic_kinematics=analytic_kinematics)
        if not self.aperture_type == 'IFU_shells':
            raise ValueError('GalkinShells is not supported with aperture_type %s. Only support with "IFU_shells"'
                             % self.aperture_type)
        self._r_bins = kwargs_aperture['r_bins']
        r_max = np.max(self._r_bins)
        self._num_pix = 100
        # factor of 1.5 to allow outside flux to be convolved into the largest bin
        self._delta_pix = 1.5 * r_max * 2 / self._num_pix
        x_grid, y_grid = util.make_grid(numPix=self._num_pix, deltapix=self._delta_pix)
        self._r_grid = np.sqrt(x_grid**2 + y_grid**2)

    def dispersion_map(self, kwargs_mass, kwargs_light, kwargs_anisotropy, **kwargs):
        """

        :param kwargs_mass: mass model parameters (following lenstronomy lens model conventions)
        :param kwargs_light: deflector light parameters (following lenstronomy light model conventions)
        :param kwargs_anisotropy: anisotropy parameters, may vary according to anisotropy type chosen.
         We refer to the Anisotropy() class for details on the parameters.
        :return: array of velocity dispersion for each IFU shell [km/s]
        """
        I_R_sigma2, IR = self.numerics._I_R_sigma2_interp(self._r_grid, kwargs_mass, kwargs_light,
                                                          kwargs_anisotropy)
        ir_map = util.array2image(IR)
        ir_sigma2_map = util.array2image(I_R_sigma2)
        kernel = self.convolution_kernel(delta_pix=self._delta_pix, num_pix=self._num_pix)
        I_R_sigma2_conv = signal.fftconvolve(ir_sigma2_map, kernel, mode='same')
        I_R_sigma2_conv = util.image2array(I_R_sigma2_conv)
        I_R_conv = signal.fftconvolve(ir_map, kernel, mode='same')
        I_R_conv = util.image2array(I_R_conv)

        # average over radial bins
        vel_disp_array = []
        r_min = self._r_bins[0]
        for r_max in self._r_bins[1:]:
            mask = mask_util.mask_shell(self._r_grid, 0, center_x=0, center_y=0, r_in=r_min, r_out=r_max)
            vel_disp = np.sum(I_R_sigma2_conv * mask) / np.sum(I_R_conv * mask)  # luminosity weighted average
            vel_disp_array.append(vel_disp)
            r_min = r_max
        self.numerics.delete_cache()
        return np.sqrt(np.array(vel_disp_array)) / 1000
