from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
import numpy as np

__all__ = ["FluxRatioLikelihood"]


class FluxRatioLikelihood(object):
    """Likelihood class for magnification of multiply lensed images."""

    def __init__(
        self,
        lens_model_class,
        flux_ratios,
        flux_ratio_errors,
        source_type="INF",
        window_size=0.1,
        grid_number=100,
        polar_grid=False,
        aspect_ratio=0.5,
    ):
        """

        :param lens_model_class: LensModel class instance
        :param flux_ratios: ratio of fluxes of the multiple images (relative to the first appearing in
         same order as the images)
        :param flux_ratio_errors: errors in the flux ratios (relative to the first appearing). Alternatively
         a log-normal covariance matrix. Note: in the case of a covariance matrix, the errors are assumed
         to be log-normal, i.e. the logarithms of the flux ratios, ln(F[i]/F[0]) are assumed to have a multivariate
         Gaussian distribution, with the given covariance matrix.
        :param source_type: string, type of source, 'INF' specifies a point source, while 'GAUSSIAN' specifies a
         finite-size source modeled as a Gaussian
        :param window_size: size of window to compute the finite flux
        :param grid_number: number of grid cells per axis in the window to numerically compute the flux
        """
        self._lens_model_class = lens_model_class
        self._flux_ratios = np.array(flux_ratios)
        self._flux_ratio_errors = np.array(flux_ratio_errors)
        self._lens_model_extensions = LensModelExtensions(lensModel=lens_model_class)
        self._source_type = source_type
        self._window_size = window_size
        self._gird_number = grid_number
        self._polar_grid = polar_grid
        self._aspect_ratio = aspect_ratio

    def logL(self, x_pos, y_pos, kwargs_lens, kwargs_special):
        """

        :param kwargs_lens: lens model keyword argument list
        :param kwargs_special: dictionary of 'special' keyword parameters
        :return: log likelihood of the measured flux ratios given a model
        """
        if self._source_type == "INF":
            mag = np.abs(
                self._lens_model_class.magnification(x_pos, y_pos, kwargs_lens)
            )
        else:
            source_sigma = kwargs_special["source_size"]
            mag = self._lens_model_extensions.magnification_finite(
                x_pos,
                y_pos,
                kwargs_lens,
                source_sigma=source_sigma,
                window_size=self._window_size,
                grid_number=self._gird_number,
                polar_grid=self._polar_grid,
                aspect_ratio=self._aspect_ratio,
            )
        if len(mag) - 1 != len(self._flux_ratios):
            return -(10**15)
        mag_ratio = mag[1:] / mag[0]
        return self._logL(mag_ratio)

    def _logL(self, flux_ratios):
        """


        :param flux_ratios: flux ratios from the model
        :return: log likelihood fo the measured flux ratios given a model
        """
        if not np.isfinite(flux_ratios).any():
            return -(10**15)
        if self._flux_ratio_errors.ndim <= 1:
            dist = (
                (flux_ratios - self._flux_ratios) ** 2
                / self._flux_ratio_errors**2
                / 2
            )
            logL = -np.sum(dist)
        elif self._flux_ratio_errors.ndim == 2:
            # Assume covariance matrix is in ln units!
            D = np.log(flux_ratios) - np.log(self._flux_ratios)
            logL = (
                -1 / 2 * D @ np.linalg.inv(self._flux_ratio_errors) @ D
            )  # TODO: only calculate the inverse once
        else:
            raise ValueError(
                "flux_ratio_errors need dimension of 1 or 2. Current dimensions are %s"
                % self._flux_ratio_errors.ndim
            )
        if not np.isfinite(logL):
            return -(10**15)
        return logL

    @property
    def num_data(self):
        """

        :return: integer, number of data points associated with the flux ratios
        """
        return len(self._flux_ratios)
