from quadmodel.data.quad_base import Quad
from quadmodel.deflector_models.sis import SIS
import numpy as np


class WFI2033(Quad):

    def __init__(self, sourcemodel_type='NARROW_LINE_Gaussian',
                 macromodel_type ='EPL_FIXED_SHEAR_MULTIPOLE'):

        zlens = 0.66
        zsource = 1.66
        x = [-0.751, -0.039, 1.445, -0.668]
        y = [0.953, 1.068, -0.307, -0.585]
        m = [1.0, 0.65, 0.5, 0.53]
        delta_m = [0.03, 0.03/0.64, 0.02/0.5, 0.02/0.53]
        delta_xy = [0.005] * 4
        keep_flux_ratio_index = [0, 1, 2]
        self.log10_host_halo_mass = 13.3
        self.log10_host_halo_mass_sigma = 0.3

        kwargs_macromodel = {'shear_amplitude_min': 0.07, 'shear_amplitude_max': 0.26}

        super(WFI2033, self).__init__(zlens, zsource, x, y, m, delta_m, delta_xy, sourcemodel_type, {}, macromodel_type,
                                    kwargs_macromodel, keep_flux_ratio_index)

    def satellite_galaxy(self, sample=True):
        """
        If the deflector system has no satellites, return an empty list of lens components (see macromodel class)
        """

        theta_E = 0.03
        center_x = 0.245
        center_y = 2.037
        if sample:
            theta_E = abs(np.random.normal(theta_E, 0.03))
            center_x = np.random.normal(center_x, 0.05)
            center_y = np.random.normal(center_y, 0.05)

        kwargs_init_1 = [{'theta_E': theta_E, 'center_x': center_x, 'center_y': center_y}]
        satellite_1 = SIS(0.66, kwargs_init_1)
        params_1 = np.array([theta_E, center_x, center_y])
        param_names_1 = ['theta_E_1', 'center_x_1', 'center_y_1']

        # from Gilman et al. (2020)
        # theta_E = 0.93
        # center_x = -3.36
        # center_y = -0.08
        # the following were computed from a full fit to the imaging data
        theta_E = 0.9
        center_x = -3.667
        center_y = -0.102
        if sample:
            theta_E = abs(np.random.normal(theta_E, 0.05))
            center_x = np.random.normal(center_x, 0.05)
            center_y = np.random.normal(center_y, 0.05)

        kwargs_init_2 = [{'theta_E': theta_E, 'center_x': center_x, 'center_y': center_y}]
        satellite_2 = SIS(0.745, kwargs_init_2)
        params_2 = np.array([theta_E, center_x, center_y])
        param_names_2 = ['theta_E_2', 'center_x_2', 'center_y_2']
        return [satellite_1, satellite_2], np.append(params_1, params_2), param_names_1 + param_names_2
