from quadmodel.data.quad_base import Quad
import numpy as np


class WFI2026(Quad):

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

        _zvalues = [0.2, 0.2342857142857143, 0.26857142857142857,
                          0.3028571428571429, 0.3371428571428572, 0.37142857142857144,
                          0.40571428571428575, 0.44000000000000006, 0.4742857142857143,
                          0.5085714285714287, 0.5428571428571429, 0.5771428571428572,
                          0.6114285714285714, 0.6457142857142857, 0.68,
                          0.7142857142857144, 0.7485714285714287, 0.7828571428571429,
                          0.8171428571428572, 0.8514285714285714, 0.8857142857142857,
                          0.9200000000000002, 0.9542857142857144, 0.9885714285714287,
                          1.022857142857143, 1.0571428571428572, 1.0914285714285714,
                          1.1257142857142857, 1.1600000000000001, 1.1942857142857144,
                          1.2285714285714286, 1.262857142857143, 1.2971428571428572,
                          1.3314285714285714, 1.3657142857142857]

        _zpdf = [0.025705329153604992, 0.025391849529780547,
                          0.04075235109717862, 0.04576802507836983, 0.058620689655172406,
                          0.07774294670846382, 0.07868338557993716, 0.10909090909090909,
                           0.1300940438871469, 0.15956112852664578, 0.19247648902821318,
                           0.21786833855799373, 0.27868338557993727, 0.30689655172413693,
                           0.3874608150470207, 0.4520376175548589, 0.5100313479623825,
                           0.600626959247649, 0.6793103448275862, 0.7611285266457679,
                           0.8087774294670793, 0.8705329153605016, 0.9677115987460815,
                           0.9479623824451411, 1.0, 0.9846394984326019, 0.9031347962382444,
                           0.870532915360496, 0.7758620689655172, 0.6482758620689655,
                           0.545141065830721, 0.3846394984326019, 0.2811912225705329,
                           0.2122257053291536, 0.121943573667]

        self.log10_host_halo_mass = 13.3
        self.log10_host_halo_mass_sigma = 0.3
        zlens = [np.array(_zvalues), np.array(_zpdf)]
        zsource = 2.23
        x = [0.187, 0.44, 0.023, -0.548]
        y = [-0.563, -0.348, 0.865, -0.179]
        m = [1.0, 0.75, 0.31, 0.28]
        delta_m = [0.02, 0.02/0.75, 0.02/0.31, 0.01/0.28]
        delta_xy = [0.005] * 4
        keep_flux_ratio_index = [0, 1, 2]

        kwargs_macromodel = {'shear_amplitude_min': 0.015, 'shear_amplitude_max': 0.16}

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