import numpy as np
from lenstronomy.LensModel.Profiles.shear import ShearReduced
from lenstronomy.LensModel.Profiles.curved_arc_sis_mst import CurvedArcSISMST
from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase

__all__ = ['CurvedArcSPT']


class CurvedArcSPT(LensProfileBase):
    """
    Curved arc model based on SIS+MST with an additional non-linear shear distortions applied on the source coordinates
    around the center.
    This profile is effectively a Source Position Transform of a curved arc and a shear distortion.

    """
    param_names = ['tangential_stretch', 'radial_stretch', 'curvature', 'direction', 'gamma1', 'gamma2', 'center_x',
                   'center_y']
    lower_limit_default = {'tangential_stretch': -100, 'radial_stretch': -5, 'curvature': 0.000001, 'direction': -np.pi,
                           'gamma1': -0.5, 'gamma2': -0.5, 'center_x': -100, 'center_y': -100}
    upper_limit_default = {'tangential_stretch': 100, 'radial_stretch': 5, 'curvature': 100, 'direction': np.pi,
                           'gamma1': 0.5, 'gamma2': 0.5, 'center_x': 100, 'center_y': 100}

    def __init__(self):
        self._curve = CurvedArcSISMST()
        self._distort = ShearReduced()
        super(CurvedArcSPT, self).__init__()

    def function(self, x, y, tangential_stretch, radial_stretch, curvature, direction, gamma1, gamma2, center_x,
                 center_y):
        """
        ATTENTION: there may not be a global lensing potential!

        :param x:
        :param y:
        :param tangential_stretch: float, stretch of intrinsic source in tangential direction
        :param radial_stretch: float, stretch of intrinsic source in radial direction
        :param curvature: 1/curvature radius
        :param direction: float, angle in radian
        :param gamma1: non-linear reduced shear distortion in the source plane
        :param gamma2: non-linear reduced shear distortion in the source plane
        :param center_x: center of source in image plane
        :param center_y: center of source in image plane
        :return:
        """
        raise NotImplemented('lensing potential for regularly curved arc is not implemented')

    def derivatives(self, x, y, tangential_stretch, radial_stretch, curvature, direction, gamma1, gamma2, center_x,
                    center_y):
        """

        :param x:
        :param y:
        :param tangential_stretch: float, stretch of intrinsic source in tangential direction
        :param radial_stretch: float, stretch of intrinsic source in radial direction
        :param curvature: 1/curvature radius
        :param direction: float, angle in radian
        :param gamma1: non-linear reduced shear distortion in the source plane
        :param gamma2: non-linear reduced shear distortion in the source plane
        :param center_x: center of source in image plane
        :param center_y: center of source in image plane
        :return:
        """
        # computed regular curved arc deflection
        f_x_c, f_y_c = self._curve.derivatives(x, y, tangential_stretch, radial_stretch, curvature, direction,
                                               center_x, center_y)
        # map to source plane coordinate system
        beta_x, beta_y = x - f_x_c, y - f_y_c
        # distort source plane coordinate system around (center_x, center_y)
        f_x_b, f_y_b = self._distort.derivatives(beta_x, beta_y, gamma1, gamma2, ra_0=center_x, dec_0=center_y)
        beta_x_, beta_y_ = beta_x - f_x_b, beta_y - f_y_b
        # compute total deflection between initial coordinate and final source coordinate to match lens equation
        # beta = theta - alpha
        f_x, f_y = x - beta_x_, y - beta_y_
        return f_x, f_y

    def hessian(self, x, y, tangential_stretch, radial_stretch, curvature, direction, gamma1, gamma2, center_x,
                center_y):
        """

        :param x:
        :param y:
        :param tangential_stretch: float, stretch of intrinsic source in tangential direction
        :param radial_stretch: float, stretch of intrinsic source in radial direction
        :param curvature: 1/curvature radius
        :param direction: float, angle in radian
        :param gamma1: non-linear reduced shear distortion in the source plane
        :param gamma2: non-linear reduced shear distortion in the source plane
        :param center_x: center of source in image plane
        :param center_y: center of source in image plane
        :return:
        """
        alpha_ra, alpha_dec = self.derivatives(x, y, tangential_stretch, radial_stretch, curvature, direction, gamma1,
                                               gamma2, center_x, center_y)
        diff = 0.0000001
        alpha_ra_dx, alpha_dec_dx = self.derivatives(x + diff, y, tangential_stretch, radial_stretch, curvature,
                                                     direction, gamma1, gamma2, center_x, center_y)
        alpha_ra_dy, alpha_dec_dy = self.derivatives(x, y + diff, tangential_stretch, radial_stretch, curvature,
                                                     direction, gamma1, gamma2, center_x, center_y)

        f_xx = (alpha_ra_dx - alpha_ra) / diff
        f_xy = (alpha_ra_dy - alpha_ra) / diff
        f_yx = (alpha_dec_dx - alpha_dec) / diff
        f_yy = (alpha_dec_dy - alpha_dec) / diff
        return f_xx, f_xy, f_yx, f_yy
