__author__ = 'sibirrer'

import numpy.testing as npt
import numpy as np
import pytest
import unittest

from lenstronomy.Analysis.kinematics_api import KinematicsAPI
import lenstronomy.Util.param_util as param_util


class TestKinematicsAPI(object):

    def setup(self):
        pass

    def test_velocity_dispersion(self):
        z_lens = 0.5
        z_source = 1.5
        kwargs_model = {'lens_model_list': ['SPEP', 'SHEAR', 'SIS', 'SIS', 'SIS'],
                          'lens_light_model_list': ['SERSIC_ELLIPSE', 'SERSIC']}

        theta_E = 1.5
        gamma = 1.8
        kwargs_lens = [{'theta_E': theta_E, 'e1': 0, 'center_x': -0.044798916793300093, 'center_y': 0.0054408937891703788, 'e2': 0, 'gamma': gamma},
                       {'e1': -0.050871696555354479, 'e2': -0.0061601733920590464}, {'center_y': 2.79985456, 'center_x': -2.32019894,
                        'theta_E': 0.28165274714097904}, {'center_y': 3.83985426,
                        'center_x': -2.32019933, 'theta_E': 0.0038110812674654873},
                       {'center_y': 4.31985428, 'center_x': -1.68019931, 'theta_E': 0.45552039839735037}]

        phi, q = -0.52624727893702705, 0.79703498156919605
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_lens_light = [{'n_sersic': 1.1212528655709217,
                              'center_x': -0.019674496231393473,
                              'e1': e1, 'e2': e2, 'amp': 1.1091367792010356, 'center_y': 0.076914975081560991,
                               'R_sersic': 0.42691611878867058},
                             {'R_sersic': 0.03025682660635394, 'amp': 139.96763298885992, 'n_sersic': 1.90000008624093865,
                              'center_x': -0.019674496231393473, 'center_y': 0.076914975081560991}]
        r_ani = 0.62
        kwargs_anisotropy = {'r_ani': r_ani}
        R_slit = 3.8
        dR_slit = 1.
        aperture_type = 'slit'
        kwargs_aperture = {'aperture_type': aperture_type, 'center_ra': 0, 'width': dR_slit, 'length': R_slit, 'angle': 0, 'center_dec': 0}

        psf_fwhm = 0.7
        kwargs_psf = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}
        anisotropy_model = 'OM'
        kwargs_mge = {'n_comp': 20}
        r_eff = 0.211919902322
        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture=kwargs_aperture, kwargs_seeing=kwargs_psf,
                                     lens_model_kinematics_bool=[True, False, False, False, False], anisotropy_model=anisotropy_model,
                                     kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge, sampling_number=1000,
                                     MGE_light=True)

        v_sigma = kinematicAPI.velocity_dispersion(kwargs_lens, kwargs_lens_light, kwargs_anisotropy, r_eff=r_eff)

        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture=kwargs_aperture,
                                     kwargs_seeing=kwargs_psf, lens_model_kinematics_bool=[True, False, False, False, False],
                                     anisotropy_model=anisotropy_model,
                                     kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge, sampling_number=1000,
                                     MGE_light=True, MGE_mass=True)
        v_sigma_mge_lens = kinematicAPI.velocity_dispersion(kwargs_lens, kwargs_lens_light, kwargs_anisotropy, r_eff=r_eff, theta_E=theta_E)
        #v_sigma_mge_lens = kinematicAPI.velocity_dispersion_numerical(kwargs_lens, kwargs_lens_light, kwargs_anisotropy, kwargs_aperture,
        #                                                          kwargs_psf, anisotropy_model, MGE_light=True, MGE_mass=True, theta_E=theta_E,
        #                                                          kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge,
        #                                                          r_eff=r_eff)
        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture=kwargs_aperture,
                                     kwargs_seeing=kwargs_psf,
                                     lens_model_kinematics_bool=[True, False, False, False, False],
                                     anisotropy_model=anisotropy_model,
                                     kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge, sampling_number=1000,
                                     MGE_light=False, MGE_mass=False, Hernquist_approx=True)
        v_sigma_hernquist = kinematicAPI.velocity_dispersion(kwargs_lens, kwargs_lens_light, kwargs_anisotropy,
                                                             r_eff=r_eff, theta_E=theta_E)
        #v_sigma_hernquist = kinematicAPI.velocity_dispersion_numerical(kwargs_lens, kwargs_lens_light, kwargs_anisotropy,
        #                                                          kwargs_aperture, kwargs_psf, anisotropy_model,
        #                                                          MGE_light=False, MGE_mass=False,
        #                                                          r_eff=r_eff, Hernquist_approx=True)

        vel_disp_temp = kinematicAPI.velocity_dispersion_analytical(theta_E, gamma, r_ani=r_ani, r_eff=r_eff)
        print(v_sigma, vel_disp_temp)
        #assert 1 == 0
        npt.assert_almost_equal(v_sigma / vel_disp_temp, 1, decimal=1)
        npt.assert_almost_equal(v_sigma_mge_lens / v_sigma, 1, decimal=1)
        npt.assert_almost_equal(v_sigma / v_sigma_hernquist, 1, decimal=1)

    def test_galkin_settings(self):
        z_lens = 0.5
        z_source = 1.5
        kwargs_model = {'lens_model_list': ['SIS'],
                        'lens_light_model_list': ['HERNQUIST']}

        kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}]
        kwargs_lens_light = [{'amp': 1, 'Rs': 1, 'center_x': 0, 'center_y': 0}]
        r_ani = 0.62
        kwargs_anisotropy = {'r_ani': r_ani}
        R_slit = 3.8
        dR_slit = 1.
        aperture_type = 'slit'
        kwargs_aperture = {'aperture_type': aperture_type, 'center_ra': 0, 'width': dR_slit, 'length': R_slit,
                           'angle': 0, 'center_dec': 0}

        psf_fwhm = 0.7
        kwargs_psf = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}
        anisotropy_model = 'OM'
        kwargs_mge = {'n_comp': 20}
        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture=kwargs_aperture,
                                     kwargs_seeing=kwargs_psf, analytic_kinematics=True,
                                     anisotropy_model=anisotropy_model,
                                     kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge, sampling_number=1000)
        galkin, kwargs_profile, kwargs_light = kinematicAPI.galkin_settings(kwargs_lens, kwargs_lens_light, r_eff=None,
                                                                            theta_E=None, gamma=None)
        npt.assert_almost_equal(kwargs_profile['gamma'], 2, decimal=2)

        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture=[kwargs_aperture],
                                     kwargs_seeing=[kwargs_psf], analytic_kinematics=True,
                                     anisotropy_model=anisotropy_model, multi_observations=True,
                                     kwargs_mge_light=kwargs_mge, kwargs_mge_mass=kwargs_mge, sampling_number=1000)
        galkin, kwargs_profile, kwargs_light = kinematicAPI.galkin_settings(kwargs_lens, kwargs_lens_light, r_eff=None,
                                                                            theta_E=None, gamma=None)
        npt.assert_almost_equal(kwargs_profile['gamma'], 2, decimal=2)

    def test_kinematic_light_profile(self):
        z_lens = 0.5
        z_source = 1.5
        kwargs_options = {'lens_light_model_list': ['HERNQUIST_ELLIPSE', 'SERSIC']}
        kwargs_mge = {'n_comp': 20}
        kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_options, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
        r_eff = 0.2
        kwargs_lens_light = [{'amp': 1, 'Rs': r_eff * 0.551, 'e1': 0., 'e2': 0, 'center_x': 0, 'center_y': 0},
                             {'amp': 1, 'R_sersic': 1, 'n_sersic': 2, 'center_x': -10, 'center_y': -10}]
        light_profile_list, kwargs_light = kinematicAPI.kinematic_light_profile(kwargs_lens_light, MGE_fit=True,
                                                                                r_eff=r_eff,
                                                                                model_kinematics_bool=[True, False],
                                                                                kwargs_mge=kwargs_mge)
        assert light_profile_list[0] == 'MULTI_GAUSSIAN'

        light_profile_list, kwargs_light = kinematicAPI.kinematic_light_profile(kwargs_lens_light, MGE_fit=False,
                                                                                r_eff=r_eff, model_kinematics_bool=[True, False])
        assert light_profile_list[0] == 'HERNQUIST_ELLIPSE'

        light_profile_list, kwargs_light = kinematicAPI.kinematic_light_profile(kwargs_lens_light, MGE_fit=False,
                                                                                Hernquist_approx=True, r_eff=r_eff,
                                                                                model_kinematics_bool=[True, False])
        assert light_profile_list[0] == 'HERNQUIST'
        npt.assert_almost_equal(kwargs_light[0]['Rs'] / kwargs_lens_light[0]['Rs'], 1, decimal=2)

    def test_kinematic_lens_profiles(self):
        z_lens = 0.5
        z_source = 1.5
        kwargs_options = {'lens_model_list': ['SPEP', 'SHEAR']}
        kin_api = KinematicsAPI(z_lens, z_source, kwargs_options, kwargs_aperture={}, kwargs_seeing={}, anisotropy_model='OM')
        kwargs_lens = [{'theta_E': 1.4272358196260446, 'e1': 0, 'center_x': -0.044798916793300093,
                        'center_y': 0.0054408937891703788, 'e2': 0, 'gamma': 1.8},
                       {'e1': -0.050871696555354479, 'e2': -0.0061601733920590464}
                       ]

        kwargs_mge = {'n_comp': 20}
        mass_profile_list, kwargs_profile = kin_api.kinematic_lens_profiles(kwargs_lens, MGE_fit=True,
                                                                             kwargs_mge=kwargs_mge, theta_E=1.4,
                                                                             model_kinematics_bool=[True, False])
        assert mass_profile_list[0] == 'MULTI_GAUSSIAN_KAPPA'

        mass_profile_list, kwargs_profile = kin_api.kinematic_lens_profiles(kwargs_lens, MGE_fit=False,
                                                                             model_kinematics_bool=[True, False])
        assert mass_profile_list[0] == 'SPEP'

    def test_model_dispersion(self):
        np.random.seed(42)
        z_lens = 0.5
        z_source = 1.5
        r_eff = 1.
        theta_E = 1.
        kwargs_model = {'lens_model_list': ['SIS'], 'lens_light_model_list': ['HERNQUIST']}
        kwargs_lens = [{'theta_E': theta_E, 'center_x': 0, 'center_y': 0}]
        kwargs_lens_light = [{'amp': 1, 'Rs': r_eff * 0.551, 'center_x': 0, 'center_y': 0}]
        kwargs_anisotropy = {'r_ani': 1}
        # settings

        R_slit = 3.8
        dR_slit = 1.
        aperture_type = 'slit'
        kwargs_aperture = {'aperture_type': aperture_type, 'center_ra': 0, 'width': dR_slit, 'length': R_slit,
                           'angle': 0, 'center_dec': 0}
        psf_fwhm = 0.7
        kwargs_seeing = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}
        anisotropy_model = 'OM'
        kin_api = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture, kwargs_seeing,
                                anisotropy_model=anisotropy_model)

        kwargs_numerics_galkin = {'interpol_grid_num': 2000, 'log_integration': True,
                                  'max_integrate': 1000, 'min_integrate': 0.0001}
        kin_api.kinematics_modeling_settings(anisotropy_model, kwargs_numerics_galkin, analytic_kinematics=True,
                                     Hernquist_approx=False, MGE_light=False, MGE_mass=False)
        vel_disp_analytic = kin_api.velocity_dispersion(kwargs_lens, kwargs_lens_light, kwargs_anisotropy, r_eff=r_eff,
                                                        theta_E=theta_E, gamma=2)

        kin_api.kinematics_modeling_settings(anisotropy_model, kwargs_numerics_galkin, analytic_kinematics=False,
                                             Hernquist_approx=False, MGE_light=False, MGE_mass=False)
        vel_disp_numerical = kin_api.velocity_dispersion(kwargs_lens, kwargs_lens_light, kwargs_anisotropy) #,
                                                         # r_eff=r_eff, theta_E=theta_E, gamma=2)
        npt.assert_almost_equal(vel_disp_numerical / vel_disp_analytic, 1, decimal=2)

        kin_api.kinematics_modeling_settings(anisotropy_model, kwargs_numerics_galkin, analytic_kinematics=False,
                                             Hernquist_approx=False, MGE_light=False, MGE_mass=False,
                                             kwargs_mge_light={'n_comp': 10}, kwargs_mge_mass={'n_comp': 5})
        assert kin_api._kwargs_mge_mass['n_comp'] == 5
        assert kin_api._kwargs_mge_light['n_comp'] == 10

    def test_velocity_dispersion_map(self):
        np.random.seed(42)
        z_lens = 0.5
        z_source = 1.5
        kwargs_options = {'lens_model_list': ['SIS'], 'lens_light_model_list': ['HERNQUIST']}
        r_eff = 1.
        theta_E = 1
        kwargs_lens = [{'theta_E': theta_E, 'center_x': 0, 'center_y': 0}]
        kwargs_lens_light = [{'amp': 1, 'Rs': r_eff * 0.551, 'center_x': 0, 'center_y': 0}]
        kwargs_anisotropy = {'r_ani': 1}

        r_bins = np.array([0, 0.5, 1])
        aperture_type = 'IFU_shells'
        kwargs_aperture = {'aperture_type': aperture_type, 'center_ra': 0, 'r_bins': r_bins, 'center_dec': 0}
        psf_fwhm = 0.7
        kwargs_seeing = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}
        anisotropy_model = 'OM'
        kin_api = KinematicsAPI(z_lens, z_source, kwargs_options, kwargs_aperture=kwargs_aperture,
                                kwargs_seeing=kwargs_seeing, anisotropy_model=anisotropy_model)

        kwargs_numerics_galkin = {'interpol_grid_num': 500, 'log_integration': True,
                                  'max_integrate': 10, 'min_integrate': 0.001}
        kin_api.kinematics_modeling_settings(anisotropy_model, kwargs_numerics_galkin, analytic_kinematics=True,
                                             Hernquist_approx=False, MGE_light=False, MGE_mass=False,
                                             num_kin_sampling=1000, num_psf_sampling=100)
        vel_disp_analytic = kin_api.velocity_dispersion_map(kwargs_lens, kwargs_lens_light, kwargs_anisotropy,
                                                            r_eff=r_eff, theta_E=theta_E, gamma=2)

        kin_api.kinematics_modeling_settings(anisotropy_model, kwargs_numerics_galkin, analytic_kinematics=False,
                                             Hernquist_approx=False, MGE_light=False, MGE_mass=False,
                                             num_kin_sampling=1000, num_psf_sampling=100)
        vel_disp_numerical = kin_api.velocity_dispersion_map(kwargs_lens, kwargs_lens_light, kwargs_anisotropy,
                                                             r_eff=r_eff, theta_E=theta_E, gamma=2)
        print(vel_disp_numerical, vel_disp_analytic)
        npt.assert_almost_equal(vel_disp_numerical, vel_disp_analytic, decimal=-1)

    def test_interpolated_sersic(self):
        from lenstronomy.Analysis.light2mass import light2mass_interpol
        kwargs_light = [{'n_sersic': 2, 'R_sersic': 0.5, 'amp': 1, 'center_x': 0.01, 'center_y': 0.01}]
        kwargs_lens = [{'n_sersic': 2, 'R_sersic': 0.5, 'k_eff': 1, 'center_x': 0.01, 'center_y': 0.01}]
        deltaPix = 0.1
        numPix = 100

        kwargs_interp = light2mass_interpol(['SERSIC'], kwargs_lens_light=kwargs_light, numPix=numPix,
                                                            deltaPix=deltaPix, subgrid_res=5)
        kwargs_lens_interp = [kwargs_interp]
        from lenstronomy.Analysis.kinematics_api import KinematicsAPI
        z_lens = 0.5
        z_source = 1.5
        r_ani = 0.62
        kwargs_anisotropy = {'r_ani': r_ani}
        R_slit = 3.8
        dR_slit = 1.
        aperture_type = 'slit'
        kwargs_aperture = {'center_ra': 0, 'width': dR_slit, 'length': R_slit, 'angle': 0, 'center_dec': 0, 'aperture_type': aperture_type}
        psf_fwhm = 0.7
        kwargs_psf = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}
        anisotropy_model = 'OM'
        r_eff = 0.5
        kwargs_model = {'lens_model_list': ['SERSIC'],
                          'lens_light_model_list': ['SERSIC']}
        kwargs_mge = {'n_comp': 20}
        kinematic_api = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_aperture, kwargs_seeing=kwargs_psf,
                                      anisotropy_model=anisotropy_model, MGE_light=True, MGE_mass=True,
                                      kwargs_mge_mass=kwargs_mge, kwargs_mge_light=kwargs_mge)

        v_sigma = kinematic_api.velocity_dispersion(kwargs_lens, kwargs_light, kwargs_anisotropy, r_eff=r_eff, theta_E=1)
        kwargs_model_interp = {'lens_model_list': ['INTERPOL'],
                                 'lens_light_model_list': ['SERSIC']}
        kinematic_api_interp = KinematicsAPI(z_lens, z_source, kwargs_model_interp, kwargs_aperture, kwargs_seeing=kwargs_psf,
                                      anisotropy_model=anisotropy_model, MGE_light=True, MGE_mass=True,
                                                                            kwargs_mge_mass=kwargs_mge,
                                                                            kwargs_mge_light=kwargs_mge)
        v_sigma_interp = kinematic_api_interp.velocity_dispersion(kwargs_lens_interp, kwargs_light, kwargs_anisotropy,
                                                        theta_E=1., r_eff=r_eff)
        npt.assert_almost_equal(v_sigma / v_sigma_interp, 1, 1)
        # use as kinematic constraints
        # compare with MGE Sersic kinematic estimate


class TestRaise(unittest.TestCase):

    def test_raise(self):
        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST']}
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kwargs_light = [{'Rs': 1, 'amp': 1, 'center_x': 0, 'center_y': 0}]
            kinematicAPI.kinematic_light_profile(kwargs_light, MGE_fit=False,
                                                 Hernquist_approx=True, r_eff=None, model_kinematics_bool=[True])
        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST']}
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kwargs_light = [{'Rs': 1, 'amp': 1, 'center_x': 0, 'center_y': 0}]
            kinematicAPI.kinematic_light_profile(kwargs_light, MGE_fit=False,
                                                 Hernquist_approx=False, r_eff=None, analytic_kinematics=True)
        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST'], 'lens_model_list': []}
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kwargs_light = [{'Rs': 1, 'amp': 1, 'center_x': 0, 'center_y': 0}]
            kinematicAPI.kinematic_lens_profiles(kwargs_light, MGE_fit=True, model_kinematics_bool=[True])
        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST'], 'lens_model_list': []}
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kinematicAPI.kinematic_lens_profiles(kwargs_lens=None, analytic_kinematics=True)

        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST'], 'lens_model_list': []}
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kwargs_lens_light = [{'Rs': 1, 'center_x': 0, 'center_y': 0}]
            kinematicAPI.kinematic_light_profile(kwargs_lens_light, r_eff=None, MGE_fit=True, model_kinematics_bool=None,
                                    Hernquist_approx=False, kwargs_mge=None)
        with self.assertRaises(ValueError):
            z_lens = 0.5
            z_source = 1.5
            kwargs_model = {'lens_light_model_list': ['HERNQUIST'], 'lens_model_list': ['SIS']}
            kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}]
            kinematicAPI = KinematicsAPI(z_lens, z_source, kwargs_model, kwargs_seeing={}, kwargs_aperture={}, anisotropy_model='OM')
            kinematicAPI.kinematic_lens_profiles(kwargs_lens, MGE_fit=True, model_kinematics_bool=None, theta_E=None,
                                kwargs_mge={})


if __name__ == '__main__':
    pytest.main()
