__author__ = ['nataliehogg', 'pierrefleury', 'danjohnson98']

from lenstronomy.LensModel.single_plane import SinglePlane
import numpy as np
import copy

__all__ = ['SinglePlaneLOS']


class SinglePlaneLOS(SinglePlane):
    """
    This class is based on the 'SinglePlane' class, modified to include
    line-of-sight effects as presented by Fleury et al. in 2104.08883.

    Are modified:
    - init (to include a new attribute, self.los)
    - fermat potential
    - alpha
    - hessian

    Are unchanged (inherited from SinglePlane):
    - ray_shooting, because it calls the modified alpha
    - mass_2d, mass_3d, density which refer to the main lens without LOS
    corrections.
    """

    def __init__(self, lens_model_list, index_los,
                 numerical_alpha_class=None,
                 lens_redshift_list=None,
                 z_source_convention=None,
                 kwargs_interp=None,
                 kwargs_synthesis=None
                 ):
        """
        Instance of SinglePlaneLOS() based on the SinglePlane(), except:
        - argument "index_los" indicating the position of the LOS model in the
        lens_model_list (for correct association with kwargs)
        - attribute "los" containing the LOS model.
        :param kwargs_synthesis: keyword arguments for the 'SYNTHESIS' lens model, if applicable
        """

        super(SinglePlaneLOS, self).__init__(lens_model_list)
        # NB: It is important to run that init first, in order to create a
        # list_func for the entire model, before splitting it between a main
        # lens and the LOS corrections

        # Extract the los model and import its class
        self._index_los = index_los
        self._los_model = lens_model_list[index_los]
        self.los = self._import_class(self._los_model, custom_class=None, kwargs_interp=None, kwargs_synthesis=kwargs_synthesis)

        # Define a separate class for the main lens
        lens_model_list_wo_los = [
            model for i, model in enumerate(lens_model_list)
            if i != index_los]
        self._main_lens = SinglePlane(lens_model_list_wo_los,
                                      numerical_alpha_class=numerical_alpha_class,
                                      lens_redshift_list=lens_redshift_list,
                                      z_source_convention=z_source_convention,
                                      kwargs_interp=kwargs_interp,
                                      kwargs_synthesis=kwargs_synthesis)

    def split_lens_los(self, kwargs):
        """
        This function splits the list of key-word arguments given to the lens
        model into those that correspond to the lens itself (kwargs_main), and
        those that correspond to the line-of-sight corrections (kwargs_los).

        :param kwargs: the list of key-word arguments passed to lenstronomy
        :return: a list of kwargs corresponding to the lens and a list of kwargs corresponding to the LOS effects
        """

        kwargs_los = copy.deepcopy(kwargs[self._index_los])
        # if 'LOS_MINIMAL' is at play, we set Gamma_os = Gamma_los
        # and Gamma_ds = Gamma_od
        if self._los_model == 'LOS_MINIMAL':
            kwargs_los['kappa_os'] = kwargs_los.pop('kappa_los')
            kwargs_los['gamma1_os'] = kwargs_los.pop('gamma1_los')
            kwargs_los['gamma2_os'] = kwargs_los.pop('gamma2_los')
            kwargs_los['omega_os'] = kwargs_los.pop('omega_los')
            kwargs_los['kappa_ds'] = kwargs_los['kappa_od']
            kwargs_los['gamma1_ds'] = kwargs_los['gamma1_od']
            kwargs_los['gamma2_ds'] = kwargs_los['gamma2_od']
            kwargs_los['omega_ds'] = kwargs_los['omega_od']

        kwargs_main = [kwarg for i, kwarg in enumerate(kwargs)
                       if i != self._index_los]

        return kwargs_main, kwargs_los

    def fermat_potential(self, x_image, y_image, kwargs_lens, x_source=None, y_source=None, k=None):
        """
        Calculates the Fermat Potential with LOS corrections in the tidal regime

        :param x_image: image position
        :param y_image: image position
        :param x_source: source position
        :param y_source: source position
        :param kwargs_lens: list of keyword arguments of lens model parameters matching the lens model classes
        :return: fermat potential in arcsec**2 as a list
        """

        kwargs_main, kwargs_los = self.split_lens_los(kwargs_lens)

        # the amplification matrices
        A_od = np.array([[1-kwargs_los['kappa_od']-kwargs_los['gamma1_od'],
                          -kwargs_los['gamma2_od']+kwargs_los['omega_od']],
                         [-kwargs_los['gamma2_od']-kwargs_los['omega_od'],
                          1-kwargs_los['kappa_od']+kwargs_los['gamma1_od']]])
        A_os = np.array([[1-kwargs_los['kappa_os']-kwargs_los['gamma1_os'],
                          -kwargs_los['gamma2_os']+kwargs_los['omega_os']],
                         [-kwargs_los['gamma2_os']-kwargs_los['omega_os'],
                          1-kwargs_los['kappa_os']+kwargs_los['gamma1_os']]])
        A_ds = np.array([[1-kwargs_los['kappa_ds']-kwargs_los['gamma1_ds'],
                          -kwargs_los['gamma2_ds']+kwargs_los['omega_ds']],
                         [-kwargs_los['gamma2_ds']-kwargs_los['omega_ds'],
                          1-kwargs_los['kappa_ds']+kwargs_los['gamma1_ds']]])

        # the inverse and transposed amplification matrices
        A_od_tsp = np.transpose(A_od)
        A_ds_inv = np.linalg.inv(A_ds)
        A_os_inv = np.linalg.inv(A_os)

        # the composite amplification matrices
        A_LOS = np.dot(np.dot(A_od_tsp, A_ds_inv), A_os)

        # Angular position where the ray hits the deflector's plane
        x_d, y_d = self.los.distort_vector(x_image, y_image,
                                           kappa=kwargs_los['kappa_od'],
                                           omega=kwargs_los['omega_od'],
                                           gamma1=kwargs_los['gamma1_od'],
                                           gamma2=kwargs_los['gamma2_od'])

        # Evaluating the potential of the main lens at this position
        effective_potential = self._main_lens.potential(x_d, y_d, kwargs=kwargs_main, k=k)

        # obtaining the source position
        if x_source is None or y_source is None:
            x_source, y_source = self.ray_shooting(x_image, y_image, kwargs_lens, k=k)

        # the source position, modified by A_os_inv
        b_x = A_os_inv[0][0]*x_source + A_os_inv[0][1]*y_source
        b_y = A_os_inv[1][0]*x_source + A_os_inv[1][1]*y_source

        # alpha'
        f_x = x_image - b_x
        f_y = y_image - b_y

        # alpha' must then be further distorted by A_LOS
        a_x = A_LOS[0][0]*f_x + A_LOS[0][1]*f_y
        a_y = A_LOS[1][0]*f_x + A_LOS[1][1]*f_y

        # we can then obtain the geometrical term
        geometry = (f_x*a_x + f_y*a_y) / 2

        return geometry - effective_potential

    def alpha(self, x, y, kwargs, k=None):
        """
        Displacement angle including the line-of-sight corrections

        :param x: x-position (preferentially arcsec)
        :type x: numpy array
        :param y: y-position (preferentially arcsec)
        :type y: numpy array
        :param kwargs: list of keyword arguments of lens model parameters
         matching the lens model classes, including line-of-sight corrections
        :param k: only evaluate the k-th lens model
        :return: deflection angles in units of arcsec
        """

        kwargs_main, kwargs_los = self.split_lens_los(kwargs)

        # Angular position where the ray hits the deflector's plane
        x_d, y_d = self.los.distort_vector(x, y,
                                           kappa=kwargs_los['kappa_od'],
                                           omega=kwargs_los['omega_od'],
                                           gamma1=kwargs_los['gamma1_od'],
                                           gamma2=kwargs_los['gamma2_od'])

        # Displacement due to the main lens only
        f_x, f_y = self._main_lens.alpha(x_d, y_d, kwargs=kwargs_main, k=k)

        # Correction due to the background convergence, shear and rotation
        f_x, f_y = self.los.distort_vector(f_x, f_y,
                                           kappa=kwargs_los['kappa_ds'],
                                           omega=kwargs_los['omega_ds'],
                                           gamma1=kwargs_los['gamma1_ds'],
                                           gamma2=kwargs_los['gamma2_ds'])

        # Perturbed position in the absence of the main lens
        x_os, y_os = self.los.distort_vector(x, y,
                                             kappa=kwargs_los['kappa_os'],
                                             omega=kwargs_los['omega_os'],
                                             gamma1=kwargs_los['gamma1_os'],
                                             gamma2=kwargs_los['gamma2_os'])

        # Complete displacement
        f_x += x - x_os
        f_y += y - y_os

        return f_x, f_y

    def hessian(self, x, y, kwargs, k=None):
        """
        Hessian matrix

        :param x: x-position (preferentially arcsec)
        :type x: numpy array
        :param y: y-position (preferentially arcsec)
        :type y: numpy array
        :param kwargs: list of keyword arguments of lens model parameters matching the lens model classes
        :param k: only evaluate the k-th lens model
        :return: f_xx, f_xy, f_yx, f_yy components
        """

        kwargs_main, kwargs_los = self.split_lens_los(kwargs)

        # Angular position where the ray hits the deflector's plane
        x_d, y_d = self.los.distort_vector(x, y,
                                           kappa=kwargs_los['kappa_od'],
                                           omega=kwargs_los['omega_od'],
                                           gamma1=kwargs_los['gamma1_od'],
                                           gamma2=kwargs_los['gamma2_od'])

        # Hessian matrix of the main lens only
        f_xx, f_xy, f_yx, f_yy = self._main_lens.hessian(x_d, y_d, kwargs=kwargs_main, k=k)

        # Multiply on the left by (1 - Gamma_ds)
        f_xx, f_xy, f_yx, f_yy = self.los.left_multiply(
                                    f_xx, f_xy, f_yx, f_yy,
                                    kappa=kwargs_los['kappa_ds'],
                                    omega=kwargs_los['omega_ds'],
                                    gamma1=kwargs_los['gamma1_ds'],
                                    gamma2=kwargs_los['gamma2_ds'])

        # Multiply on the right by (1 - Gamma_od)
        f_xx, f_xy, f_yx, f_yy = self.los.right_multiply(
                                    f_xx, f_xy, f_yx, f_yy,
                                    kappa=kwargs_los['kappa_od'],
                                    omega=kwargs_los['omega_od'],
                                    gamma1=kwargs_los['gamma1_od'],
                                    gamma2=kwargs_los['gamma2_od'])

        # LOS contribution in the absence of the main lens
        f_xx += kwargs_los['kappa_os'] + kwargs_los['gamma1_os']
        f_xy += kwargs_los['gamma2_os'] - kwargs_los['omega_os']
        f_yx += kwargs_los['gamma2_os'] + kwargs_los['omega_os']
        f_yy += kwargs_los['kappa_os'] - kwargs_los['gamma1_os']

        return f_xx, f_xy, f_yx, f_yy

    def mass_3d(self, r, kwargs, bool_list=None):
        """
        Computes the mass within a 3d sphere of radius r *for the main lens only*

        :param r: radius (in angular units)
        :param kwargs: list of keyword arguments of lens model parameters matching the lens model classes
        :param bool_list: list of bools that are part of the output
        :return: mass (in angular units, modulo epsilon_crit)
        """

        print("Note: The computation of the 3d mass ignores the LOS corrections.")

        kwargs_main, kwargs_los = self.split_lens_los(kwargs)
        mass_3d = self._main_lens.mass_3d(r=r, kwargs=kwargs_main, bool_list=bool_list)

        return mass_3d

    def mass_2d(self, r, kwargs, bool_list=None):
        """
        Computes the mass enclosed a projected (2d) radius r *for the main lens only*

        The mass definition is such that:

        .. math::
            \\alpha = mass_2d / r / \\pi

        with alpha is the deflection angle

        :param r: radius (in angular units)
        :param kwargs: list of keyword arguments of lens model parameters matching the lens model classes
        :param bool_list: list of bools that are part of the output
        :return: projected mass (in angular units, modulo epsilon_crit)
        """

        print("Note: The computation of the 2d mass ignores the LOS corrections.")

        kwargs_main, kwargs_los = self.split_lens_los(kwargs)
        mass_2d = self._main_lens.mass_2d(r=r, kwargs=kwargs_main, bool_list=bool_list)

        return mass_2d

    def density(self, r, kwargs, bool_list=None):
        """
        3d mass density at radius r *for the main lens only*
        The integral in the LOS projection of this quantity results in the convergence quantity.

        :param r: radius (in angular units)
        :param kwargs: list of keyword arguments of lens model parameters matching the lens model classes
        :param bool_list: list of bools that are part of the output
        :return: mass density at radius r (in angular units, modulo epsilon_crit)
        """

        print("Note: The computation of the density ignores the LOS corrections.")

        kwargs_main, kwargs_los = self.split_lens_los(kwargs)
        density = self._main_lens.density(r=r, kwargs=kwargs_main, bool_list=bool_list)

        return density

    def potential(self, x, y, kwargs, k=None):
        """
        Lensing potential *of the main lens only*
        In the presence of LOS corrections, the system generally does not admit
        a potential, in the sense that the curl of alpha is generally non-zero

        :param x: x-position (preferentially arcsec)
        :type x: numpy array
        :param y: y-position (preferentially arcsec)
        :type y: numpy array
        :param kwargs: list of keyword arguments of lens model parameters matching the lens model classes
        :param k: only evaluate the k-th lens model
        :return: lensing potential in units of arcsec^2
        """

        print("Note: The computation of the potential ignores the LOS corrections.\
              In the presence of LOS corrections, a lensing system does not always\
              derive from a potential.")

        # kwargs_main, kwargs_los = self.split_lens_los(kwargs)
        potential = self._main_lens.potential(x, y, kwargs, k=k)

        return potential
