import numpy as np
import coolest
from lenstronomy.Util.param_util import phi_q2_ellipticity, shear_polar2cartesian


def shapelet_amp_coolest_to_lenstronomy(value):
    """
    Transforms shapelets coefficients from COOLEST conventions (x to the right)
    to lenstronomy conventions (x following ra, to the left)

    :param value: amplitude of the shapelet (float or np.array) in COOLEST conventions
    :return: amplitude of the shapelet (float or np.array) in lenstronomy conventions
    """
    if value is None:
        return None
    else:
        n = 0  # this is the shapelet order
        k = 0  # this is the index of a coefficient for a given order
        new_value = []
        for coeff in value:
            if n % 2 == 0:
                if k % 2 == 1:
                    coeff = -coeff
            elif n % 2 == 1:
                if k % 2 == 0:
                    coeff = -coeff
            new_value.append(coeff)
            k += 1
            if k == n + 1:
                n += 1
                k = 0
        return new_value


def degree_coolest_to_radian_lenstronomy(value):
    """
    Transform an angle in degree in COOLEST conventions (from y to negative x - aka East of North)
    into an angle in radian in lenstronomy conventions (from x to y with x pointing to the left - aka North of East)

    :param value: float, angle in COOLEST conventions
    :return: float, angle in lenstronomy conventions
    """
    if value is None:
        return None
    else:
        lenstro_oriented_degree = - value + 90.
        if lenstro_oriented_degree >= 180.:
            lenstro_oriented_degree -= 180.
        elif lenstro_oriented_degree < 0.:
            lenstro_oriented_degree += 180
        return lenstro_oriented_degree * np.pi / 180.


def qphi_coolest_to_e1e2_lenstronomy(q, phi):
    """
    Transform q and phi (axis ratio, position angle East-of-North) to e1,e2 in lenstronomy

    :param q: float, axis ratio
    :param phi: float, position angle in COOLEST conventions
    :return: e1, e2, lenstronomy usual ellipticity parameters
    """
    if None in [q, phi]:
        return None, None
    else:
        angle = degree_coolest_to_radian_lenstronomy(phi)
        e1, e2 = phi_q2_ellipticity(angle, q)
        return e1, e2


def gamma_phi_coolest_to_g1_g2_lenstronomy(gamma_ext, phi_ext):
    """
    Transform gamma_ext and phi_ext (shear strength, position angle East-of-North)
    to gamma1,gamma2 in lenstronomy

    :param gamma_ext: float, shear strenght
    :param phi_ext: float, shear angle in COOLEST conventions
    :return: gamma1, gamma2, lenstronomy usual shear parameters
    """
    if None in [gamma_ext, phi_ext]:
        return None, None
    else:
        angle = degree_coolest_to_radian_lenstronomy(phi_ext)
        gamma1, gamma2 = shear_polar2cartesian(angle,gamma_ext)
        return gamma1, gamma2


def ellibounds_coolest_to_lenstronomy(q_down, q_up, phi_down, phi_up):
    """
    Transforms upper and lower bounds on coolest ellipticity parameters (q, phi) towards lenstronomy bound on e1, e2
    The mapping can not be perfect but it's the best we can do

    :param q_down: float, lower bound of axis ratio
    :param q_up: float, upper bound of axis ratio
    :param phi_down: float, lower bound of position angle in COOLEST conventions
    :param phi_up: float, upper bound of position angle in COOLEST conventions
    :return: e1_down, e1_up, e2_down, e2_up, bounds for lenstronomy usual ellipticity parameters
    """
    if None in [q_down, q_up, phi_down, phi_up]:
        return None, None, None, None
    else:
        e1_down = - (1 - q_down) / (1 + q_down)
        e2_down = - (1 - q_down) / (1 + q_down)

        e1_up = (1 - q_down) / (1 + q_down)
        e2_up = (1 - q_down) / (1 + q_down)

        # WE LOOSE THE ANGLE INFORMATION AND THE OTHER INFO IS NOT WELL TRANSLATED (best we can do)

        return e1_down, e1_up, e2_down, e2_up


def shearbounds_coolest_to_lenstronomy(gamma_ext_down, gamma_ext_up, phi_ext_down, phi_ext_up):
    """
    Transforms upper and lower bounds on coolest shear parameters (gamma_ext, phi_ext) towards lenstronomy bounds
    on gamma_1, gamma_2
    The mapping can not be perfect but it's the best we can do

    :param gamma_ext_down: float, lower bound of shear strenght
    :param gamma_ext_up: float, upper bound of shear strenght
    :param phi_ext_down: float, lower bound of shear position angle in COOLEST conventions
    :param phi_ext_up: float, upper bound of shear position angle in COOLEST conventions
    :return: gamma1_down, gamma1_up, gamma2_down, gamma2_up ; bounds for lenstronomy usual shear parameters
    """
    if None in [gamma_ext_down, gamma_ext_up, phi_ext_down, phi_ext_up]:
        return None, None, None, None
    else:
        gamma1_down = - gamma_ext_up
        gamma2_down = - gamma_ext_up

        gamma1_up = gamma_ext_up
        gamma2_up = gamma_ext_up

        # WE LOOSE THE ANGLE INFORMATION AND THE OTHER INFO IS NOT WELL TRANSLATED (best we can do)

        return gamma1_down, gamma1_up, gamma2_down, gamma2_up


def update_kwargs_shear(shear_idx, lens_model_list, kwargs_lens, kwargs_lens_init, kwargs_lens_up, kwargs_lens_down,
          kwargs_lens_fixed, kwargs_lens_sigma, cleaning=False):
    """
    Update the lens model list and kwargs with SHEAR mass model (gamma_ext - phi_ext)

    :param shear_idx: coolest.template.classes.profiles.mass.ExternalShear object
    :param lens_model_list: the usual lenstronomy lens_model_list
    :param kwargs_lens: the usual lenstronomy kwargs
    :param kwargs_lens_init: the usual lenstronomy kwargs
    :param kwargs_lens_up: the usual lenstronomy kwargs
    :param kwargs_lens_down: the usual lenstronomy kwargs
    :param kwargs_lens_fixed: the usual lenstronomy kwargs
    :param kwargs_lens_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed
    :return: updated list and kwargs
    """
    lens_model_list.append('SHEAR')
    for shear_name, shear_param in shear_idx.parameters.items():
        if shear_name == 'gamma_ext':
            gammaext = getattr(shear_param.point_estimate, 'value')
            gammaext_up = getattr(shear_param.definition_range, 'max_value')
            gammaext_down = getattr(shear_param.definition_range, 'min_value')
            gammaext_fixed = gammaext if getattr(shear_param, 'fixed') else None
        elif shear_name == 'phi_ext':
            psiext = getattr(shear_param.point_estimate, 'value')
            psiext_up = getattr(shear_param.definition_range, 'max_value')
            psiext_down = getattr(shear_param.definition_range, 'min_value')
            psiext_fixed = psiext if getattr(shear_param, 'fixed') else None
        else:
            print(f"{shear_name} not known")
    gamma1, gamma2 = gamma_phi_coolest_to_g1_g2_lenstronomy(gammaext, psiext)
    gamma1_fixed, gamma2_fixed = gamma_phi_coolest_to_g1_g2_lenstronomy(gammaext_fixed, psiext_fixed)
    gamma1_down, gamma1_up, gamma2_down, gamma2_up = shearbounds_coolest_to_lenstronomy(gammaext_down, gammaext_up, psiext_down,
                                                                            psiext_up)

    kw_1 = {'gamma1': gamma1, 'gamma2': gamma2, 'ra_0': 0., 'dec_0': 0.}
    kw_up_1 = {'gamma1': gamma1_up, 'gamma2': gamma2_up, 'ra_0': 0., 'dec_0': 0.}
    kw_down_1 = {'gamma1': gamma1_down, 'gamma2': gamma2_down, 'ra_0': 0., 'dec_0': 0.}
    kw_fixed_1 = {'gamma1': gamma1_fixed, 'gamma2': gamma2_fixed, 'ra_0': 0., 'dec_0': 0.}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'gamma1': 0.0, 'gamma2': -0.0, 'ra_0': 0., 'dec_0': 0.}
        kw_up_default = {'gamma1': 0.5, 'gamma2': 0.5, 'ra_0': 100, 'dec_0': 100}
        kw_down_default = {'gamma1': -0.5, 'gamma2': -0.5, 'ra_0': -100, 'dec_0': -100}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_lens.append(kw_)
    kwargs_lens_init.append(kw_init)
    kwargs_lens_up.append(kw_up)
    kwargs_lens_down.append(kw_down)
    kwargs_lens_fixed.append(kw_fixed)
    kwargs_lens_sigma.append({'gamma1': 0.1, 'gamma2': 0.1, 'ra_0': 0., 'dec_0': 0.})
    print('\t Shear correctly added')

    return


def update_kwargs_pemd(mass, lens_model_list, kwargs_lens, kwargs_lens_init, kwargs_lens_up, kwargs_lens_down, kwargs_lens_fixed,
         kwargs_lens_sigma, cleaning=False):
    """
    Update the lens list and kwargs with PEMD mass model

    :param mass: coolest.template.classes.profiles.mass.PEMD object
    :param lens_model_list: the usual lenstronomy lens_model_list
    :param kwargs_lens: the usual lenstronomy kwargs
    :param kwargs_lens_init: the usual lenstronomy kwargs
    :param kwargs_lens_up: the usual lenstronomy kwargs
    :param kwargs_lens_down: the usual lenstronomy kwargs
    :param kwargs_lens_fixed: the usual lenstronomy kwargs
    :param kwargs_lens_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed

    :return: updated list and kwargs
    """
    lens_model_list.append('PEMD')
    for mass_name, mass_param in mass.parameters.items():
        if mass_name == 'theta_E':
            te = getattr(mass_param.point_estimate, 'value')
            te_up = getattr(mass_param.definition_range, 'max_value')
            te_down = getattr(mass_param.definition_range, 'min_value')
            te_fixed = te if getattr(mass_param, 'fixed') else None
        elif mass_name == 'gamma':
            gamma = getattr(mass_param.point_estimate, 'value')
            gamma_up = getattr(mass_param.definition_range, 'max_value')
            gamma_down = getattr(mass_param.definition_range, 'min_value')
            gamma_fixed = gamma if getattr(mass_param, 'fixed') else None
        elif mass_name == 'q':
            q = getattr(mass_param.point_estimate, 'value')
            q_up = getattr(mass_param.definition_range, 'max_value')
            q_down = getattr(mass_param.definition_range, 'min_value')
            q_fixed = q if getattr(mass_param, 'fixed') else None
        elif mass_name == 'phi':
            phi = getattr(mass_param.point_estimate, 'value')
            phi_up = getattr(mass_param.definition_range, 'max_value')
            phi_down = getattr(mass_param.definition_range, 'min_value')
            phi_fixed = phi if getattr(mass_param, 'fixed') else None
        elif mass_name == 'center_x':
            center_x = - getattr(mass_param.point_estimate, 'value') if getattr(mass_param.point_estimate,
                                                                                'value') is not None else None
            center_x_up = getattr(mass_param.definition_range, 'max_value')
            center_x_down = getattr(mass_param.definition_range, 'min_value')
            center_x_fixed = center_x if getattr(mass_param, 'fixed') else None
        elif mass_name == 'center_y':
            center_y = getattr(mass_param.point_estimate, 'value')
            center_y_up = getattr(mass_param.definition_range, 'max_value')
            center_y_down = getattr(mass_param.definition_range, 'min_value')
            center_y_fixed = center_y if getattr(mass_param, 'fixed') else None
        else:
            print(f"{mass_name} not known")

    e1, e2 = qphi_coolest_to_e1e2_lenstronomy(q, phi)
    e1_fixed, e2_fixed = qphi_coolest_to_e1e2_lenstronomy(q_fixed, phi_fixed)
    e1_down, e1_up, e2_down, e2_up = ellibounds_coolest_to_lenstronomy(q_down, q_up, phi_down, phi_up)

    kw_1 = {'theta_E': te, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': center_x, 'center_y': center_y}
    kw_up_1 = {'theta_E': te_up, 'gamma': gamma_up, 'e1': e1_up, 'e2': e2_up, 'center_x': center_x_up,
               'center_y': center_y_up}
    kw_down_1 = {'theta_E': te_down, 'gamma': gamma_down, 'e1': e1_down, 'e2': e2_down,
                 'center_x': center_x_down, 'center_y': center_y_down}
    kw_fixed_1 = {'theta_E': te_fixed, 'gamma': gamma_fixed, 'e1': e1_fixed, 'e2': e2_fixed,
                  'center_x': center_x_fixed, 'center_y': center_y_fixed}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'theta_E': 1., 'gamma': 2., 'e1': 0.0, 'e2': -0.0, 'center_x': 0., 'center_y': 0.}
        kw_up_default = {'theta_E': 100, 'gamma': 2.5, 'e1': 0.5, 'e2': 0.5, 'center_x': 100, 'center_y': 100}
        kw_down_default = {'theta_E': 0, 'gamma': 1.5, 'e1': -0.5, 'e2': -0.5, 'center_x': -100, 'center_y': -100}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_lens.append(kw_)
    kwargs_lens_init.append(kw_init)
    kwargs_lens_up.append(kw_up)
    kwargs_lens_down.append(kw_down)
    kwargs_lens_fixed.append(kw_fixed)
    kwargs_lens_sigma.append({'theta_E': 1.5, 'gamma': 0.2, 'e1': 0.3, 'e2': 0.3, 'center_x': 0.5, 'center_y': 0.5})

    print('\t PEMD correctly added')

    return


def update_kwargs_sie(mass, lens_model_list, kwargs_lens, kwargs_lens_init, kwargs_lens_up, kwargs_lens_down, kwargs_lens_fixed,
        kwargs_lens_sigma, cleaning=False):
    """
    Update the lens list and kwargs with SIE mass model

    :param mass: coolest.template.classes.profiles.mass.SIE object
    :param lens_model_list: the usual lenstronomy lens_model_list
    :param kwargs_lens: the usual lenstronomy kwargs
    :param kwargs_lens_init: the usual lenstronomy kwargs
    :param kwargs_lens_up: the usual lenstronomy kwargs
    :param kwargs_lens_down: the usual lenstronomy kwargs
    :param kwargs_lens_fixed: the usual lenstronomy kwargs
    :param kwargs_lens_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed

    :return: updated list and kwargs
    """
    lens_model_list.append('SIE')
    for mass_name, mass_param in mass.parameters.items():
        if mass_name == 'theta_E':
            te = getattr(mass_param.point_estimate, 'value')
            te_up = getattr(mass_param.definition_range, 'max_value')
            te_down = getattr(mass_param.definition_range, 'min_value')
            te_fixed = te if getattr(mass_param, 'fixed') else None
        elif mass_name == 'q':
            q = getattr(mass_param.point_estimate, 'value')
            q_up = getattr(mass_param.definition_range, 'max_value')
            q_down = getattr(mass_param.definition_range, 'min_value')
            q_fixed = q if getattr(mass_param, 'fixed') else None
        elif mass_name == 'phi':
            phi = getattr(mass_param.point_estimate, 'value')
            phi_up = getattr(mass_param.definition_range, 'max_value')
            phi_down = getattr(mass_param.definition_range, 'min_value')
            phi_fixed = phi if getattr(mass_param, 'fixed') else None
        elif mass_name == 'center_x':
            center_x = - getattr(mass_param.point_estimate, 'value') if getattr(mass_param.point_estimate,
                                                                                'value') is not None else None
            center_x_up = getattr(mass_param.definition_range, 'max_value')
            center_x_down = getattr(mass_param.definition_range, 'min_value')
            center_x_fixed = center_x if getattr(mass_param, 'fixed') else None
        elif mass_name == 'center_y':
            center_y = getattr(mass_param.point_estimate, 'value')
            center_y_up = getattr(mass_param.definition_range, 'max_value')
            center_y_down = getattr(mass_param.definition_range, 'min_value')
            center_y_fixed = center_y if getattr(mass_param, 'fixed') else None
        else:
            print(f"{mass_name} not known")

    e1, e2 = qphi_coolest_to_e1e2_lenstronomy(q, phi)
    e1_fixed, e2_fixed = qphi_coolest_to_e1e2_lenstronomy(q_fixed, phi_fixed)
    e1_down, e1_up, e2_down, e2_up = ellibounds_coolest_to_lenstronomy(q_down, q_up, phi_down, phi_up)

    kw_1 = {'theta_E': te, 'e1': e1, 'e2': e2, 'center_x': center_x, 'center_y': center_y}
    kw_up_1 = {'theta_E': te_up, 'e1': e1_up, 'e2': e2_up, 'center_x': center_x_up,
               'center_y': center_y_up}
    kw_down_1 = {'theta_E': te_down, 'e1': e1_down, 'e2': e2_down,
                 'center_x': center_x_down, 'center_y': center_y_down}
    kw_fixed_1 = {'theta_E': te_fixed, 'e1': e1_fixed, 'e2': e2_fixed,
                  'center_x': center_x_fixed, 'center_y': center_y_fixed}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'theta_E': 1., 'e1': 0.0, 'e2': -0.0, 'center_x': 0., 'center_y': 0.}
        kw_up_default = {'theta_E': 100, 'e1': 0.5, 'e2': 0.5, 'center_x': 100, 'center_y': 100}
        kw_down_default = {'theta_E': 0, 'e1': -0.5, 'e2': -0.5, 'center_x': -100, 'center_y': -100}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_lens.append(kw_)
    kwargs_lens_init.append(kw_init)
    kwargs_lens_up.append(kw_up)
    kwargs_lens_down.append(kw_down)
    kwargs_lens_fixed.append(kw_fixed)
    kwargs_lens_sigma.append({'theta_E': 1.5, 'e1': 0.3, 'e2': 0.3, 'center_x': 0.5, 'center_y': 0.5})

    print('\t SIE correctly added')

    return


def update_kwargs_sersic(light, light_model_list, kwargs_light, kwargs_light_init, kwargs_light_up, kwargs_light_down,
           kwargs_light_fixed, kwargs_light_sigma, cleaning=False):
    """
    Update the source list and kwargs with SERSIC_ELLISPE light model

    :param light: coolest.template.classes.profiles.light.Sersic object
    :param light_model_list: the usual lenstronomy lens_light_model_list or source_light_model_list
    :param kwargs_light: the usual lenstronomy kwargs
    :param kwargs_light_init: the usual lenstronomy kwargs
    :param kwargs_light_up: the usual lenstronomy kwargs
    :param kwargs_light_down: the usual lenstronomy kwargs
    :param kwargs_light_fixed: the usual lenstronomy kwargs
    :param kwargs_light_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed

    :return: updated list and kwargs
    """
    light_model_list.append('SERSIC_ELLIPSE')
    for light_name, light_param in light.parameters.items():
        if light_name == 'I_eff':
            amp = getattr(light_param.point_estimate, 'value')
            amp_up = getattr(light_param.definition_range, 'max_value')
            amp_down = getattr(light_param.definition_range, 'min_value')
            amp_fixed = amp if getattr(light_param, 'fixed') else None
        elif light_name == 'theta_eff':
            R = getattr(light_param.point_estimate, 'value')
            R_up = getattr(light_param.definition_range, 'max_value')
            R_down = getattr(light_param.definition_range, 'min_value')
            R_fixed = R if getattr(light_param, 'fixed') else None
        elif light_name == 'n':
            n = getattr(light_param.point_estimate, 'value')
            n_up = getattr(light_param.definition_range, 'max_value')
            n_down = getattr(light_param.definition_range, 'min_value')
            n_fixed = n if getattr(light_param, 'fixed') else None
        elif light_name == 'q':
            q = getattr(light_param.point_estimate, 'value')
            q_up = getattr(light_param.definition_range, 'max_value')
            q_down = getattr(light_param.definition_range, 'min_value')
            q_fixed = q if getattr(light_param, 'fixed') else None
        elif light_name == 'phi':
            phi = getattr(light_param.point_estimate, 'value')
            phi_up = getattr(light_param.definition_range, 'max_value')
            phi_down = getattr(light_param.definition_range, 'min_value')
            phi_fixed = phi if getattr(light_param, 'fixed') else None
        elif light_name == 'center_x':
            cx = - getattr(light_param.point_estimate, 'value') if getattr(light_param.point_estimate,
                                                                           'value') is not None else None
            cx_up = getattr(light_param.definition_range, 'max_value')
            cx_down = getattr(light_param.definition_range, 'min_value')
            cx_fixed = cx if getattr(light_param, 'fixed') else None
        elif light_name == 'center_y':
            cy = getattr(light_param.point_estimate, 'value')
            cy_up = getattr(light_param.definition_range, 'max_value')
            cy_down = getattr(light_param.definition_range, 'min_value')
            cy_fixed = cy if getattr(light_param, 'fixed') else None
        else:
            print(f'Parameter {light_name} unknown in SersicEllipse profile.')

    e1, e2 = qphi_coolest_to_e1e2_lenstronomy(q, phi)
    e1_fixed, e2_fixed = qphi_coolest_to_e1e2_lenstronomy(q_fixed, phi_fixed)
    e1_down, e1_up, e2_down, e2_up = ellibounds_coolest_to_lenstronomy(q_down, q_up, phi_down, phi_up)

    kw_1 = {'amp': amp, 'R_sersic': R, 'n_sersic': n, 'e1': e1, 'e2': e2, 'center_x': cx, 'center_y': cy}
    kw_up_1 = {'R_sersic': R_up, 'n_sersic': n_up, 'e1': e1_up, 'e2': e2_up,
               'center_x': cx_up, 'center_y': cy_up}
    kw_down_1 = {'R_sersic': R_down, 'n_sersic': n_down, 'e1': e1_down, 'e2': e2_down,
                 'center_x': cx_down, 'center_y': cy_down}
    kw_fixed_1 = {'R_sersic': R_fixed, 'n_sersic': n_fixed, 'e1': e1_fixed,
                  'e2': e2_fixed, 'center_x': cx_fixed, 'center_y': cy_fixed}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'amp': 1., 'R_sersic': 1., 'n_sersic': 3.5, 'e1': 0.0, 'e2': -0.0, 'center_x': 0.,
                           'center_y': 0.}
        kw_up_default = {'amp': 100, 'R_sersic': 100, 'n_sersic': 8, 'e1': 0.5, 'e2': 0.5, 'center_x': 100,
                         'center_y': 100}
        kw_down_default = {'amp': 0, 'R_sersic': 0, 'n_sersic': 0.5, 'e1': -0.5, 'e2': -0.5, 'center_x': -100,
                           'center_y': -100}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_light.append(kw_)
    kwargs_light_init.append(kw_init)
    kwargs_light_up.append(kw_up)
    kwargs_light_down.append(kw_down)
    kwargs_light_fixed.append(kw_fixed)
    kwargs_light_sigma.append(
        {'R_sersic': 1.5, 'n_sersic': 1.5, 'e1': 0.2, 'e2': 0.2, 'center_x': 0.5, 'center_y': 0.5})
    print('\t Sersic (Ellipse) correctly added')
    return


def update_kwargs_shapelets(light, light_model_list, kwargs_light, kwargs_light_init, kwargs_light_up, kwargs_light_down,
              kwargs_light_fixed, kwargs_light_sigma, cleaning=False):
    """
    Update the source list and kwargs with SHAPELETS light model

    :param light: coolest.template.classes.profiles.light.Shapelets object
    :param light_model_list: the usual lenstronomy lens_light_model_list or source_light_model_list
    :param kwargs_light: the usual lenstronomy kwargs
    :param kwargs_light_init: the usual lenstronomy kwargs
    :param kwargs_light_up: the usual lenstronomy kwargs
    :param kwargs_light_down: the usual lenstronomy kwargs
    :param kwargs_light_fixed: the usual lenstronomy kwargs
    :param kwargs_light_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed

    :return: updated list and kwargs
    """
    light_model_list.append('SHAPELETS')
    for light_name, light_param in light.parameters.items():
        if light_name == 'beta':
            b = getattr(light_param.point_estimate, 'value')
            b_up = getattr(light_param.definition_range, 'max_value')
            b_down = getattr(light_param.definition_range, 'min_value')
            b_fixed = b if getattr(light_param, 'fixed') else None
        elif light_name == 'n_max':
            nmax = getattr(light_param.point_estimate, 'value')
            nmax_fixed = nmax if getattr(light_param, 'fixed') else None
        elif light_name == 'center_x':
            cx = - getattr(light_param.point_estimate, 'value') if getattr(light_param.point_estimate,
                                                                           'value') is not None else None
            cx_up = getattr(light_param.definition_range, 'max_value')
            cx_down = getattr(light_param.definition_range, 'min_value')
            cx_fixed = cx if getattr(light_param, 'fixed') else None
        elif light_name == 'center_y':
            cy = getattr(light_param.point_estimate, 'value')
            cy_up = getattr(light_param.definition_range, 'max_value')
            cy_down = getattr(light_param.definition_range, 'min_value')
            cy_fixed = cy if getattr(light_param, 'fixed') else None
        elif light_name == 'amps':
            amp = shapelet_amp_coolest_to_lenstronomy(getattr(light_param.point_estimate, 'value'))
            amp_up = shapelet_amp_coolest_to_lenstronomy(getattr(light_param.definition_range, 'max_value'))
            amp_down = shapelet_amp_coolest_to_lenstronomy(getattr(light_param.definition_range, 'min_value'))
            amp_fixed = amp if getattr(light_param, 'fixed') else None
        else:
            print(f'Parameter {light_name} unknown in Shapelets profile.')

    kw_1 = {'amp': amp, 'beta': b, 'center_x': cx, 'center_y': cy, 'n_max': nmax}
    kw_up_1 = {'beta': b_up, 'center_x': cx_up, 'center_y': cy_up}
    kw_down_1 = {'beta': b_down, 'center_x': cx_down, 'center_y': cy_down}
    kw_fixed_1 = {'beta': b_fixed, 'center_x': cx_fixed, 'center_y': cy_fixed, 'n_max': nmax}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'amp': 1., 'beta': 0.1, 'center_x': 0., 'center_y': 0., 'n_max': -1}
        kw_up_default = {'amp': 100, 'beta': 100, 'center_x': 100, 'center_y': 100, 'n_max': 1000}
        kw_down_default = {'amp': 0, 'beta': 0, 'center_x': -100, 'center_y': -100, 'n_max': -1}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_light.append(kw_)
    kwargs_light_init.append(kw_init)
    kwargs_light_up.append(kw_up)
    kwargs_light_down.append(kw_down)
    kwargs_light_fixed.append(kw_fixed)
    kwargs_light_sigma.append({'beta': 0.5, 'center_x': 0.5, 'center_y': 0.5})
    print('\t Shapelets correctly added')
    return


def update_kwargs_lensed_ps(light, ps_model_list, kwargs_ps, kwargs_ps_init, kwargs_ps_up, kwargs_ps_down, kwargs_ps_fixed,
                        kwargs_ps_sigma, cleaning=False):
    """
    Update the source list and kwargs with lensed point source "LENSED_POSITION" light model

    :param light: coolest.template.classes.profiles.lightLensedPS object
    :param ps_model_list: the usual lenstronomy point_source_model_list
    :param kwargs_ps: the usual lenstronomy kwargs
    :param kwargs_ps_init: the usual lenstronomy kwargs
    :param kwargs_ps_up: the usual lenstronomy kwargs
    :param kwargs_ps_down: the usual lenstronomy kwargs
    :param kwargs_ps_fixed: the usual lenstronomy kwargs
    :param kwargs_ps_sigma: the usual lenstronomy kwargs
    :param cleaning: bool, if True, will update the empty fields with default values + cleans the kwargs_fixed

    :return: updated list and kwargs
    """
    ps_model_list.append('LENSED_POSITION')

    try:
        num_ps = len(getattr(light.parameters['ra_list'].point_estimate, 'value'))
    except:
        num_ps = 4

    for light_name, light_param in light.parameters.items():
        if light_name == 'ra_list':
            ra = - np.array(getattr(light_param.point_estimate, 'value')) if getattr(light_param.point_estimate,
                                                                           'value') is not None else None
            ra_up = getattr(light_param.definition_range, 'max_value')
            ra_down = getattr(light_param.definition_range, 'min_value')
            ra_fixed = ra if getattr(light_param, 'fixed') else None
        elif light_name == 'dec_list':
            dec = getattr(light_param.point_estimate, 'value')
            dec_up = getattr(light_param.definition_range, 'max_value')
            dec_down = getattr(light_param.definition_range, 'min_value')
            dec_fixed = dec if getattr(light_param, 'fixed') else None
        elif light_name == 'amps':
            amp = getattr(light_param.point_estimate, 'value')
            amp_up = getattr(light_param.definition_range, 'max_value')
            amp_down = getattr(light_param.definition_range, 'min_value')
            amp_fixed = amp if getattr(light_param, 'fixed') else None
        else:
            print(f'Parameter {light_name} unknown in LensedPS profile.')

    kw_1 = {'point_amp': np.array(amp) if amp is not None else None,
            'ra_image': np.array(ra) if ra is not None else None,
            'dec_image': np.array(dec) if dec is not None else None}
    kw_up_1 = {'ra_image': np.array(ra_up) if ra_up is not None else None,
               'dec_image': np.array(dec_up) if dec_up is not None else None}
    kw_down_1 = {'ra_image': np.array(ra_down) if ra_down is not None else None,
                 'dec_image': np.array(dec_down) if dec_down is not None else None}
    kw_fixed_1 = {'ra_image': ra_fixed,
                  'dec_image': dec_fixed}

    kw_ = kw_1.copy()
    kw_init = kw_1.copy()
    kw_up = kw_up_1.copy()
    kw_down = kw_down_1.copy()
    kw_fixed = kw_fixed_1.copy()

    if cleaning is True:
        kw_init_default = {'point_amp': np.ones(num_ps), 'ra_image': np.ones(num_ps), 'dec_image': np.ones(num_ps)}
        kw_up_default = {'ra_image': np.ones(num_ps) * 10, 'dec_image': np.ones(num_ps) * 10}
        kw_down_default = {'ra_image': np.ones(num_ps) * (-10), 'dec_image': np.ones(num_ps) * (-10)}
        for key, val in kw_1.items():
            if val is None:
                kw_init[key] = kw_init_default[key]
        for key, val in kw_up_1.items():
            if val is None:
                kw_up[key] = kw_up_default[key]
        for key, val in kw_down_1.items():
            if val is None:
                kw_down[key] = kw_down_default[key]
        for key, val in kw_fixed_1.items():
            if val is None:
                del kw_fixed[key]

    kwargs_ps.append(kw_)
    kwargs_ps_init.append(kw_init)
    kwargs_ps_up.append(kw_up)
    kwargs_ps_down.append(kw_down)
    kwargs_ps_fixed.append(kw_fixed)
    kwargs_ps_sigma.append({'ra_image': np.ones(num_ps), 'dec_image': np.ones(num_ps)})
    print('\t Lensed point sources correctly added')
    return

