import numpy as np
from lenstronomy.PointSource.Types.base_ps import (
    PSBase,
    _expand_to_array,
    _shrink_array,
)

__all__ = ["SourcePositions"]


class SourcePositions(PSBase):
    """Class of a single point source defined in the original source coordinate position
    that is lensed. The lens equation is solved to compute the image positions for the
    specified source position.

    Name within the PointSource module: 'SOURCE_POSITION'
    parameters: ra_source, dec_source, source_amp, mag_pert (optional)
    If fixed_magnification=True, than 'source_amp' is a parameter instead of 'point_amp'
    mag_pert is a list of fractional magnification pertubations applied to point source images
    """

    def image_position(
        self,
        kwargs_ps,
        kwargs_lens=None,
        magnification_limit=None,
        kwargs_lens_eqn_solver=None,
        **kwargs
    ):
        """On-sky image positions.

        :param kwargs_ps: keyword arguments of the point source model
        :param kwargs_lens: keyword argument list of the lens model(s), only used when
            requiring the lens equation solver
        :param magnification_limit: float >0 or None, if float is set and additional
            images are computed, only those images will be computed that exceed the
            lensing magnification (absolute value) limit
        :param kwargs_lens_eqn_solver: keyword arguments specifying the numerical
            settings for the lens equation solver see LensEquationSolver() class for
            details
        :return: image positions in x, y as arrays
        """
        if kwargs_lens_eqn_solver is None:
            kwargs_lens_eqn_solver = {}
        ra_source, dec_source = self.source_position(kwargs_ps)
        ra_image, dec_image = self._solver.image_position_from_source(
            ra_source,
            dec_source,
            kwargs_lens,
            magnification_limit=magnification_limit,
            **kwargs_lens_eqn_solver
        )
        return ra_image, dec_image

    def source_position(self, kwargs_ps, **kwargs):
        """Original source position (prior to lensing)

        :param kwargs_ps: point source keyword arguments
        :return: x, y position (as numpy arrays)
        """
        ra_source = kwargs_ps["ra_source"]
        dec_source = kwargs_ps["dec_source"]
        return np.array(ra_source), np.array(dec_source)

    def image_amplitude(
        self,
        kwargs_ps,
        kwargs_lens=None,
        x_pos=None,
        y_pos=None,
        magnification_limit=None,
        kwargs_lens_eqn_solver=None,
    ):
        """Image brightness amplitudes.

        :param kwargs_ps: keyword arguments of the point source model
        :param kwargs_lens: keyword argument list of the lens model(s), only ignored
            when providing image positions directly
        :param x_pos: pre-computed image position (no lens equation solver applied)
        :param y_pos: pre-computed image position (no lens equation solver applied)
        :param magnification_limit: float >0 or None, if float is set and additional
            images are computed, only those images will be computed that exceed the
            lensing magnification (absolute value) limit
        :param kwargs_lens_eqn_solver: keyword arguments specifying the numerical
            settings for the lens equation solver see LensEquationSolver() class for
            details
        :return: array of image amplitudes
        """
        if self._fixed_magnification:
            if x_pos is not None and y_pos is not None:
                ra_image, dec_image = x_pos, y_pos
            else:
                if kwargs_lens_eqn_solver is None:
                    kwargs_lens_eqn_solver = {}
                ra_image, dec_image = self.image_position(
                    kwargs_ps,
                    kwargs_lens=kwargs_lens,
                    magnification_limit=magnification_limit,
                    **kwargs_lens_eqn_solver
                )
            mag = self._lens_model.magnification(ra_image, dec_image, kwargs_lens)
            point_amp = kwargs_ps["source_amp"] * np.abs(mag)
        else:
            point_amp = kwargs_ps["point_amp"]
            if x_pos is not None:
                point_amp = _expand_to_array(point_amp, len(x_pos))
        mag_pert = kwargs_ps.get("mag_pert", 1)
        mag_pert = _shrink_array(mag_pert, len(point_amp))
        point_amp *= np.array(mag_pert)
        return np.array(point_amp)

    def source_amplitude(self, kwargs_ps, kwargs_lens=None):
        """Intrinsic brightness amplitude of point source When brightnesses are defined
        in magnified on-sky positions, the intrinsic brightness is computed as the mean
        in the magnification corrected image position brightnesses.

        :param kwargs_ps: keyword arguments of the point source model
        :param kwargs_lens: keyword argument list of the lens model(s), used when
            brightness are defined in magnified on-sky positions
        :return: brightness amplitude (as numpy array)
        """
        if self._fixed_magnification:
            source_amp = kwargs_ps["source_amp"]
        else:
            ra_image, dec_image = self.image_position(kwargs_ps, kwargs_lens)
            mag = self._lens_model.magnification(ra_image, dec_image, kwargs_lens)
            point_amp = kwargs_ps["point_amp"]
            source_amp = np.mean(np.array(point_amp) / np.array(mag))
        return np.array(source_amp)
