import numpy as np
import pytest

import autofit as af
from autolens.lens.util import lens_fit_util as util
from autolens.data import ccd
from autolens.data.array import mask as msk
from autolens.lens import lens_data as ld
from autolens.lens import ray_tracing, lens_fit
from autolens.model.galaxy import galaxy as g
from autolens.model.galaxy.util import galaxy_util
from autolens.model.inversion import inversions
from autolens.model.inversion import pixelizations
from autolens.model.inversion import regularization
from autolens.model.profiles import light_profiles as lp
from autolens.model.profiles import mass_profiles as mp
from test.unit.mock.model.mock_profiles import MockLightProfile


class TestLensProfileFit:
    class TestLikelihood:

        def test__1x2_image__no_psf_blurring__tracing_fits_data_with_chi_sq_5(self):
            # The image plane image generated by the galaxy is [1.0, 1.0]

            # Thus the chi squared is 4.0**2.0 + 3.0**2.0 = 25.0

            psf = ccd.PSF(array=(np.array([[0.0, 0.0, 0.0],
                                           [0.0, 1.0, 0.0],
                                           [0.0, 0.0, 0.0]])), pixel_scale=1.0)

            ccd_data = ccd.CCDData(5.0 * np.ones((3, 4)), pixel_scale=1.0, psf=psf,
                                   noise_map=np.ones((3, 4)))
            ccd_data.image[1, 2] = 4.0

            mask = msk.Mask(array=np.array([[True, True, True, True],
                                            [True, False, False, True],
                                            [True, True, True, True]]), pixel_scale=1.0)

            lens_data_5x5 = ld.LensData(ccd_data=ccd_data, mask=mask, sub_grid_size=1)

            # Setup as a ray trace instance, using a light profile for the lens

            g0 = g.Galaxy(redshift=0.5,
                          light_profile=MockLightProfile(value=1.0, size=2))
            tracer = ray_tracing.TracerImagePlane(lens_galaxies=[g0],
                                                  image_plane_grid_stack=lens_data_5x5.grid_stack)

            fit = lens_fit.LensProfileFit(lens_data=lens_data_5x5, tracer=tracer)

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (fit.mask_2d == np.array([[True, True, True, True],
                                             [True, False, False, True],
                                             [True, True, True, True]])).all()

            assert (fit.image_1d == np.array([5.0, 4.0])).all()
            assert (fit.image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                              [0.0, 5.0, 4.0, 0.0],
                                              [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.noise_map_1d == np.array([1.0, 1.0])).all()
            assert (fit.noise_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                  [0.0, 1.0, 1.0, 0.0],
                                                  [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.model_image_1d == np.array([1.0, 1.0])).all()
            assert (fit.model_image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                    [0.0, 1.0, 1.0, 0.0],
                                                    [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.residual_map_1d == np.array([4.0, 3.0])).all()
            assert (fit.residual_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                     [0.0, 4.0, 3.0, 0.0],
                                                     [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.chi_squared_map_1d == np.array([16.0, 9.0])).all()
            assert (fit.chi_squared_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                        [0.0, 16.0, 9.0, 0.0],
                                                        [0.0, 0.0, 0.0, 0.0]])).all()

            assert fit.chi_squared == 25.0
            assert fit.reduced_chi_squared == 25.0 / 2.0
            assert fit.noise_normalization == (2.0 * np.log(2 * np.pi * 1.0 ** 2.0))
            assert fit.likelihood == -0.5 * (
                        25.0 + 2.0 * np.log(2 * np.pi * 1.0 ** 2.0))

        def test__1x2_image__include_psf_blurring__tracing_fits_data_with_chi_sq_4(
                self):
            # This PSF changes the blurred image plane image from [1.0, 1.0] to [1.0, 5.0]

            # Thus, the chi squared is 4.0**2.0 + 0.0**2.0 = 16.0

            psf = ccd.PSF(array=(np.array([[0.0, 0.0, 0.0],
                                           [0.0, 1.0, 3.0],
                                           [0.0, 0.0, 0.0]])), pixel_scale=1.0)

            ccd_data = ccd.CCDData(5.0 * np.ones((3, 4)), pixel_scale=1.0, psf=psf,
                                   noise_map=np.ones((3, 4)))
            ccd_data.image[1, 2] = 4.0

            mask = msk.Mask(array=np.array([[True, True, True, True],
                                            [True, False, False, True],
                                            [True, True, True, True]]), pixel_scale=1.0)

            lens_data_5x5 = ld.LensData(ccd_data=ccd_data, mask=mask, sub_grid_size=1)

            # Setup as a ray trace instance, using a light profile for the lens

            g0 = g.Galaxy(redshift=0.5,
                          light_profile=MockLightProfile(value=1.0, size=2))
            tracer = ray_tracing.TracerImagePlane(lens_galaxies=[g0],
                                                  image_plane_grid_stack=lens_data_5x5.grid_stack)

            fit = lens_fit.LensProfileFit(lens_data=lens_data_5x5, tracer=tracer)

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (fit.mask_2d == np.array([[True, True, True, True],
                                             [True, False, False, True],
                                             [True, True, True, True]])).all()

            assert (fit.image_1d == np.array([5.0, 4.0])).all()
            assert (fit.image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                              [0.0, 5.0, 4.0, 0.0],
                                              [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.noise_map_1d == np.array([1.0, 1.0])).all()
            assert (fit.noise_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                  [0.0, 1.0, 1.0, 0.0],
                                                  [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.model_image_1d == np.array([1.0, 4.0])).all()
            assert (fit.model_image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                    [0.0, 1.0, 4.0, 0.0],
                                                    [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.residual_map_1d == np.array([4.0, 0.0])).all()
            assert (fit.residual_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                     [0.0, 4.0, 0.0, 0.0],
                                                     [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.chi_squared_map_1d == np.array([16.0, 0.0])).all()
            assert (fit.chi_squared_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                        [0.0, 16.0, 0.0, 0.0],
                                                        [0.0, 0.0, 0.0, 0.0]])).all()

            assert fit.chi_squared == 16.0
            assert fit.reduced_chi_squared == 16.0 / 2.0
            assert fit.noise_normalization == (2.0 * np.log(2 * np.pi * 1.0 ** 2.0))
            assert fit.likelihood == -0.5 * (
                        16.0 + 2.0 * np.log(2 * np.pi * 1.0 ** 2.0))

        def test_hyper_galaxy_changes_noise_above_from_1_to_2__reflected_in_likelihood(
                self):
            # This PSF changes the blurred image plane image from [1.0, 1.0] to [1.0, 5.0]

            # Thus, the chi squared is 4.0**2.0 + 0.0**2.0 = 16.0

            # The hyper galaxy increases the noise in both pixels by 1.0, to 2.0.

            # This reduces the chi squared to 2.0 instead of 4.0

            psf = ccd.PSF(array=(np.array([[0.0, 0.0, 0.0],
                                           [0.0, 1.0, 3.0],
                                           [0.0, 0.0, 0.0]])), pixel_scale=1.0)

            ccd_data = ccd.CCDData(5.0 * np.ones((3, 4)), pixel_scale=1.0, psf=psf,
                                   noise_map=np.ones((3, 4)))
            ccd_data.image[1, 2] = 4.0

            mask = msk.Mask(array=np.array([[True, True, True, True],
                                            [True, False, False, True],
                                            [True, True, True, True]]), pixel_scale=1.0)

            lens_data_5x5 = ld.LensData(ccd_data=ccd_data, mask=mask, sub_grid_size=1)

            # Setup as a ray trace instance, using a light profile for the lens

            g0 = g.Galaxy(redshift=0.5,
                          light_profile=MockLightProfile(value=1.0, size=2),
                          hyper_galaxy=g.HyperGalaxy(contribution_factor=1.0,
                                                     noise_factor=1.0, noise_power=1.0),
                          hyper_model_image_1d=np.ones(2),
                          hyper_galaxy_image_1d=np.ones(2), hyper_minimum_value=0.0)

            tracer = ray_tracing.TracerImagePlane(lens_galaxies=[g0],
                                                  image_plane_grid_stack=lens_data_5x5.grid_stack)

            fit = lens_fit.LensProfileFit(lens_data=lens_data_5x5, tracer=tracer)

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (fit.mask_2d == np.array([[True, True, True, True],
                                             [True, False, False, True],
                                             [True, True, True, True]])).all()

            assert (fit.image_1d == np.array([5.0, 4.0])).all()
            assert (fit.image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                              [0.0, 5.0, 4.0, 0.0],
                                              [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.noise_map_1d == np.array([2.0, 2.0])).all()
            assert (fit.noise_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                  [0.0, 2.0, 2.0, 0.0],
                                                  [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.model_image_1d == np.array([1.0, 4.0])).all()
            assert (fit.model_image_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                    [0.0, 1.0, 4.0, 0.0],
                                                    [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.residual_map_1d == np.array([4.0, 0.0])).all()
            assert (fit.residual_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                     [0.0, 4.0, 0.0, 0.0],
                                                     [0.0, 0.0, 0.0, 0.0]])).all()

            assert (fit.chi_squared_map_1d == np.array([4.0, 0.0])).all()
            assert (fit.chi_squared_map_2d == np.array([[0.0, 0.0, 0.0, 0.0],
                                                        [0.0, 4.0, 0.0, 0.0],
                                                        [0.0, 0.0, 0.0, 0.0]])).all()

            assert fit.chi_squared == 4.0
            assert fit.reduced_chi_squared == 4.0 / 2.0
            assert fit.noise_normalization == (2.0 * np.log(2 * np.pi * 2.0 ** 2.0))
            assert fit.likelihood == -0.5 * (4.0 + 2.0 * np.log(2 * np.pi * 2.0 ** 2.0))

    class TestCompareToManual:

        def test___all_lens_fit_quantities__no_hyper_methods(self, lens_data_5x5):
            g0 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=1.0),
                          mass_profile=mp.SphericalIsothermal(einstein_radius=1.0))
            g1 = g.Galaxy(redshift=1.0,
                          light_profile=lp.EllipticalSersic(intensity=1.0))

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0],
                                                         source_galaxies=[g1],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack)

            padded_tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g0], source_galaxies=[g1],
                image_plane_grid_stack=lens_data_5x5.padded_grid_stack)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer,
                                                           padded_tracer=padded_tracer)

            assert lens_data_5x5.noise_map_1d == pytest.approx(fit.noise_map_1d, 1e-4)
            assert lens_data_5x5.noise_map_2d == pytest.approx(fit.noise_map_2d, 1e-4)

            model_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_5x5.convolver_image)
            model_image_2d = lens_data_5x5.map_to_scaled_array(array_1d=model_image_1d)

            assert model_image_1d == pytest.approx(fit.model_image_1d, 1e-4)
            assert model_image_2d == pytest.approx(fit.model_image_2d, 1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=model_image_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=lens_data_5x5.noise_map_1d)

            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=lens_data_5x5.noise_map_2d, mask=lens_data_5x5.mask_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared,
                noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)
            assert likelihood == fit.figure_of_merit

        def test___lens_fit_galaxy_image_dict__corresponds_to_blurred_galaxy_images(
                self, lens_data_5x5):
            g0 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=1.0),
                          mass_profile=mp.SphericalIsothermal(einstein_radius=1.0))
            g1 = g.Galaxy(redshift=1.0,
                          light_profile=lp.EllipticalSersic(intensity=1.0))
            g2 = g.Galaxy(redshift=1.0)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0],
                                                         source_galaxies=[g1, g2],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            g0_image_plane_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=lens_data_5x5.grid_stack.sub, galaxies=[g0])
            g0_image_plane_bluring_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=lens_data_5x5.grid_stack.blurring, galaxies=[g0])

            g0_blurred_image_plane_image_1d = lens_data_5x5.convolver_image.convolve_image(
                image_array=g0_image_plane_image_1d,
                blurring_array=g0_image_plane_bluring_image_1d)

            g1_image_plane_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=tracer.source_plane.grid_stack.sub, galaxies=[g1])
            g1_image_plane_bluring_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=tracer.source_plane.grid_stack.blurring, galaxies=[g1])

            g1_blurred_image_plane_image_1d = lens_data_5x5.convolver_image.convolve_image(
                image_array=g1_image_plane_image_1d,
                blurring_array=g1_image_plane_bluring_image_1d)

            assert fit.galaxy_image_1d_dict[g0] == pytest.approx(
                g0_blurred_image_plane_image_1d, 1.0e-4)
            assert fit.galaxy_image_1d_dict[g1] == pytest.approx(
                g1_blurred_image_plane_image_1d, 1.0e-4)
            assert (fit.galaxy_image_1d_dict[g2] == np.zeros(9)).all()

            assert fit.model_image_1d == pytest.approx(
                fit.galaxy_image_1d_dict[g0] + fit.galaxy_image_1d_dict[g1], 1.0e-4)

            g0_blurred_image_plane_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=g0_blurred_image_plane_image_1d)
            g1_blurred_image_plane_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=g1_blurred_image_plane_image_1d)

            assert fit.galaxy_image_2d_dict[g0] == pytest.approx(
                g0_blurred_image_plane_image_2d, 1.0e-4)
            assert fit.galaxy_image_2d_dict[g1] == pytest.approx(
                g1_blurred_image_plane_image_2d, 1.0e-4)

            assert fit.model_image_2d == pytest.approx(
                fit.galaxy_image_2d_dict[g0] + fit.galaxy_image_2d_dict[g1], 1.0e-4)

        def test___all_lens_fit_quantities__including_hyper_methods(self,
                                                                    lens_data_5x5):
            g0 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=1.0),
                          mass_profile=mp.SphericalIsothermal(einstein_radius=1.0),
                          hyper_galaxy=g.HyperGalaxy(contribution_factor=1.0,
                                                     noise_factor=1.0, noise_power=1.0),
                          hyper_model_image_1d=np.ones(9),
                          hyper_galaxy_image_1d=np.ones(9), hyper_minimum_value=0.0)
            g1 = g.Galaxy(redshift=1.0,
                          light_profile=lp.EllipticalSersic(intensity=1.0))

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0],
                                                         source_galaxies=[g1],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack)

            padded_tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g0], source_galaxies=[g1],
                image_plane_grid_stack=lens_data_5x5.padded_grid_stack)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer,
                                                           padded_tracer=padded_tracer)

            hyper_noise_1d = tracer.hyper_noise_map_1d_from_noise_map_1d(
                noise_map_1d=lens_data_5x5.noise_map_1d)

            hyper_noise_map_1d = lens_data_5x5.noise_map_1d + hyper_noise_1d
            hyper_noise_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=hyper_noise_map_1d)
            assert hyper_noise_map_1d == pytest.approx(fit.noise_map_1d, 1e-4)
            assert hyper_noise_map_2d == pytest.approx(fit.noise_map_2d, 1e-4)

            model_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_5x5.convolver_image)

            model_image_2d = lens_data_5x5.map_to_scaled_array(array_1d=model_image_1d)

            assert model_image_1d == pytest.approx(fit.model_image_1d, 1e-4)
            assert model_image_2d == pytest.approx(fit.model_image_2d, 1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=model_image_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=hyper_noise_map_1d)

            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=hyper_noise_map_2d, mask=lens_data_5x5.mask_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared, noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)
            assert likelihood == fit.figure_of_merit

        def test___blurred_and_model_images_of_planes_and_unmasked_blurred_profile_image_properties(
                self, lens_data_5x5):
            g0 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=1.0),
                          mass_profile=mp.SphericalIsothermal(einstein_radius=1.0))
            g1 = g.Galaxy(redshift=1.0,
                          light_profile=lp.EllipticalSersic(intensity=1.0))

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0],
                                                         source_galaxies=[g1],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack)

            padded_tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g0], source_galaxies=[g1],
                image_plane_grid_stack=lens_data_5x5.padded_grid_stack)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer,
                                                           padded_tracer=padded_tracer)

            blurred_profile_image_2d_of_planes = \
                tracer.blurred_profile_image_plane_image_2d_of_planes_from_convolver_image(
                    convolver_image=lens_data_5x5.convolver_image)

            assert (blurred_profile_image_2d_of_planes[0] ==
                    fit.model_image_2d_of_planes[0]).all()
            assert (blurred_profile_image_2d_of_planes[1] ==
                    fit.model_image_2d_of_planes[1]).all()

            unmasked_blurred_profile_image = lens_data_5x5.padded_grid_stack.unmasked_blurred_image_from_psf_and_unmasked_image(
                psf=lens_data_5x5.psf,
                unmasked_image_1d=padded_tracer.profile_image_plane_image_1d)

            assert (
                        unmasked_blurred_profile_image == fit.unmasked_blurred_image_plane_image).all()

            unmasked_blurred_profile_image_of_planes = \
                padded_tracer.unmasked_blurred_profile_image_plane_image_of_planes_from_psf(
                    psf=lens_data_5x5.psf)

            assert (unmasked_blurred_profile_image_of_planes[0] ==
                    fit.unmasked_blurred_image_plane_image_of_planes[0]).all()
            assert (unmasked_blurred_profile_image_of_planes[1] ==
                    fit.unmasked_blurred_image_plane_image_of_planes[1]).all()

            unmasked_blurred_profile_image_of_galaxies = \
                padded_tracer.unmasked_blurred_profile_image_plane_image_of_plane_and_galaxies_from_psf(
                    psf=lens_data_5x5.psf)

            assert (unmasked_blurred_profile_image_of_galaxies[0][0] ==
                    fit.unmasked_blurred_image_plane_image_of_planes_and_galaxies[0][
                        0]).all()
            assert (unmasked_blurred_profile_image_of_galaxies[1][0] ==
                    fit.unmasked_blurred_image_plane_image_of_planes_and_galaxies[1][
                        0]).all()


class TestLensInversionFit:
    class TestCompareToManual:

        def test___all_lens_fit_quantities__no_hyper_methods(self, lens_data_5x5):
            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))

            g0 = g.Galaxy(redshift=0.5, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g.Galaxy(redshift=0.5)], source_galaxies=[g0],
                image_plane_grid_stack=lens_data_5x5.grid_stack,
                border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)
            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                mapper=mapper, regularization=reg, image_1d=lens_data_5x5.image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix)

            assert inversion.reconstructed_data_1d == pytest.approx(fit.model_image_1d,
                                                                    1e-4)
            assert inversion.reconstructed_data_2d == pytest.approx(fit.model_image_2d,
                                                                    1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=inversion.reconstructed_data_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=lens_data_5x5.noise_map_1d)
            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)

            assert chi_squared_map_1d == pytest.approx(fit.chi_squared_map_1d, 1e-4)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_5x5.mask_2d, noise_map=lens_data_5x5.noise_map_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared, noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)

            likelihood_with_regularization = \
                util.likelihood_with_regularization_from_chi_squared_regularization_term_and_noise_normalization(
                    chi_squared=chi_squared,
                    regularization_term=inversion.regularization_term,
                    noise_normalization=noise_normalization)

            assert likelihood_with_regularization == pytest.approx(
                fit.likelihood_with_regularization, 1e-4)

            evidence = util.evidence_from_inversion_terms(
                chi_squared=chi_squared,
                regularization_term=inversion.regularization_term,
                log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term,
                log_regularization_term=inversion.log_det_regularization_matrix_term,
                noise_normalization=noise_normalization)

            assert evidence == fit.evidence
            assert evidence == fit.figure_of_merit

        def test___lens_fit_galaxy_image_dict__has_inversion_reconstructed_data(self,
                                                                                lens_data_5x5):
            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))

            g0 = g.Galaxy(redshift=0.5)
            g1 = g.Galaxy(redshift=0.5, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0],
                                                         source_galaxies=[g1],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack,
                                                         border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)
            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                mapper=mapper, regularization=reg, image_1d=lens_data_5x5.image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix)

            assert (fit.galaxy_image_1d_dict[g0] == np.zeros(9)).all()

            assert fit.galaxy_image_1d_dict[g1] == pytest.approx(
                inversion.reconstructed_data_1d, 1.0e-4)

            assert fit.model_image_1d == pytest.approx(fit.galaxy_image_1d_dict[g1],
                                                       1.0e-4)

            assert fit.galaxy_image_2d_dict[g1] == pytest.approx(
                inversion.reconstructed_data_2d, 1.0e-4)

            assert fit.model_image_2d == pytest.approx(fit.galaxy_image_2d_dict[g1],
                                                       1.0e-4)

        def test___all_lens_fit_quantities__include_hyper_methods(self, lens_data_5x5):
            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))

            g0 = g.Galaxy(redshift=0.5, pixelization=pix, regularization=reg,
                          hyper_galaxy=g.HyperGalaxy(contribution_factor=1.0,
                                                     noise_factor=1.0, noise_power=1.0),
                          hyper_model_image_1d=np.ones(9),
                          hyper_galaxy_image_1d=np.ones(9), hyper_minimum_value=0.0)

            tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g.Galaxy(redshift=0.5)], source_galaxies=[g0],
                image_plane_grid_stack=lens_data_5x5.grid_stack,
                border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            hyper_noise_1d = tracer.hyper_noise_map_1d_from_noise_map_1d(
                noise_map_1d=lens_data_5x5.noise_map_1d)
            hyper_noise_map_1d = lens_data_5x5.noise_map_1d + hyper_noise_1d
            hyper_noise_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=hyper_noise_map_1d)
            assert hyper_noise_map_1d == pytest.approx(fit.noise_map_1d, 1e-4)
            assert hyper_noise_map_2d == pytest.approx(fit.noise_map_2d, 1e-4)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)
            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                mapper=mapper, regularization=reg, image_1d=lens_data_5x5.image_1d,
                noise_map_1d=hyper_noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix)

            assert inversion.reconstructed_data_1d == pytest.approx(fit.model_image_1d,
                                                                    1e-4)
            assert inversion.reconstructed_data_2d == pytest.approx(fit.model_image_2d,
                                                                    1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=inversion.reconstructed_data_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=hyper_noise_map_1d)
            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)

            assert chi_squared_map_1d == pytest.approx(fit.chi_squared_map_1d, 1e-4)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_5x5.mask_2d, noise_map=hyper_noise_map_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared, noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)

            likelihood_with_regularization = \
                util.likelihood_with_regularization_from_chi_squared_regularization_term_and_noise_normalization(
                    chi_squared=chi_squared,
                    regularization_term=inversion.regularization_term,
                    noise_normalization=noise_normalization)

            assert likelihood_with_regularization == pytest.approx(
                fit.likelihood_with_regularization, 1e-4)

            evidence = util.evidence_from_inversion_terms(
                chi_squared=chi_squared,
                regularization_term=inversion.regularization_term,
                log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term,
                log_regularization_term=inversion.log_det_regularization_matrix_term,
                noise_normalization=noise_normalization)

            assert evidence == fit.evidence
            assert evidence == fit.figure_of_merit

        def test___blurred_and_model_images_of_planes_and_unmasked_blurred_profile_image_properties(
                self, lens_data_5x5):
            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))

            g0 = g.Galaxy(redshift=0.5, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(
                lens_galaxies=[g.Galaxy(redshift=0.5)], source_galaxies=[g0],
                image_plane_grid_stack=lens_data_5x5.grid_stack,
                border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                mapper=mapper, regularization=reg, image_1d=lens_data_5x5.image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix)

            assert None == fit.model_image_2d_of_planes[0]
            assert inversion.reconstructed_data_2d == pytest.approx(
                fit.model_image_2d_of_planes[1], 1.0e-4)

            with pytest.raises(Exception):
                fit.unmasked_blurred_image_plane_image
                fit.unmasked_blurred_image_plane_image_of_planes
                fit.unmasked_blurred_image_plane_image_of_planes_and_galaxies


class TestLensProfileInversionFit:
    class TestCompareToManual:

        def test___all_lens_fit_quantities__no_hyper_methods(self, lens_data_5x5):
            galaxy_light = g.Galaxy(redshift=0.5,
                                    light_profile=lp.EllipticalSersic(intensity=1.0))

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[galaxy_light],
                                                         source_galaxies=[galaxy_pix],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack,
                                                         border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            blurred_profile_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_5x5.convolver_image)

            blurred_profile_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=blurred_profile_image_1d)

            assert blurred_profile_image_1d == pytest.approx(
                fit.blurred_profile_image_1d, 1e-4)
            assert blurred_profile_image_2d == pytest.approx(
                fit.blurred_profile_image_2d, 1e-4)

            profile_subtracted_image_1d = lens_data_5x5.image_1d - blurred_profile_image_1d
            profile_subtracted_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=profile_subtracted_image_1d)

            assert profile_subtracted_image_1d == pytest.approx(
                fit.profile_subtracted_image_1d)
            assert profile_subtracted_image_2d == pytest.approx(
                fit.profile_subtracted_image_2d)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix, mapper=mapper,
                regularization=reg)

            model_image_1d = blurred_profile_image_1d + inversion.reconstructed_data_1d
            model_image_2d = lens_data_5x5.map_to_scaled_array(array_1d=model_image_1d)

            assert model_image_1d == pytest.approx(fit.model_image_1d, 1e-4)
            assert model_image_2d == pytest.approx(fit.model_image_2d, 1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=model_image_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=lens_data_5x5.noise_map_1d)
            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)

            assert chi_squared_map_1d == pytest.approx(fit.chi_squared_map_1d, 1e-4)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_5x5.mask_2d, noise_map=lens_data_5x5.noise_map_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared, noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)

            likelihood_with_regularization = \
                util.likelihood_with_regularization_from_chi_squared_regularization_term_and_noise_normalization(
                    chi_squared=chi_squared,
                    regularization_term=inversion.regularization_term,
                    noise_normalization=noise_normalization)

            assert likelihood_with_regularization == pytest.approx(
                fit.likelihood_with_regularization, 1e-4)

            evidence = util.evidence_from_inversion_terms(
                chi_squared=chi_squared,
                regularization_term=inversion.regularization_term,
                log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term,
                log_regularization_term=inversion.log_det_regularization_matrix_term,
                noise_normalization=noise_normalization)

            assert evidence == fit.evidence
            assert evidence == fit.figure_of_merit

        def test___lens_fit_galaxy_image_dict__has_blurred_profile_images_and_inversion_reconstructed_data(
                self, lens_data_5x5):
            g0 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=1.0))
            g1 = g.Galaxy(redshift=0.5,
                          light_profile=lp.EllipticalSersic(intensity=2.0))
            g2 = g.Galaxy(redshift=0.5)

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[g0, g1, g2],
                                                         source_galaxies=[galaxy_pix],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack,
                                                         border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            g0_image_plane_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=lens_data_5x5.grid_stack.sub, galaxies=[g0])
            g0_image_plane_bluring_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=lens_data_5x5.grid_stack.blurring, galaxies=[g0])

            g0_blurred_image_plane_image_1d = lens_data_5x5.convolver_image.convolve_image(
                image_array=g0_image_plane_image_1d,
                blurring_array=g0_image_plane_bluring_image_1d)

            g1_image_plane_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=tracer.source_plane.grid_stack.sub, galaxies=[g1])
            g1_image_plane_bluring_image_1d = galaxy_util.intensities_of_galaxies_from_grid(
                grid=tracer.source_plane.grid_stack.blurring, galaxies=[g1])

            g1_blurred_image_plane_image_1d = lens_data_5x5.convolver_image.convolve_image(
                image_array=g1_image_plane_image_1d,
                blurring_array=g1_image_plane_bluring_image_1d)

            blurred_profile_image_1d = g0_blurred_image_plane_image_1d + g1_blurred_image_plane_image_1d

            profile_subtracted_image_1d = lens_data_5x5.image_1d - blurred_profile_image_1d
            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix, mapper=mapper,
                regularization=reg)

            assert fit.galaxy_image_1d_dict[g0] == pytest.approx(
                g0_blurred_image_plane_image_1d, 1.0e-4)
            assert fit.galaxy_image_1d_dict[g1] == pytest.approx(
                g1_blurred_image_plane_image_1d, 1.0e-4)
            assert (fit.galaxy_image_1d_dict[g2] == np.zeros(9)).all()
            assert fit.galaxy_image_1d_dict[galaxy_pix] == pytest.approx(
                inversion.reconstructed_data_1d, 1.0e-4)

            assert fit.model_image_1d == pytest.approx(
                fit.galaxy_image_1d_dict[g0] + fit.galaxy_image_1d_dict[g1] +
                fit.galaxy_image_1d_dict[galaxy_pix], 1.0e-4)

            g0_blurred_image_plane_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=g0_blurred_image_plane_image_1d)
            g1_blurred_image_plane_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=g1_blurred_image_plane_image_1d)

            assert fit.galaxy_image_2d_dict[g0] == pytest.approx(
                g0_blurred_image_plane_image_2d, 1.0e-4)
            assert fit.galaxy_image_2d_dict[g1] == pytest.approx(
                g1_blurred_image_plane_image_2d, 1.0e-4)
            assert fit.galaxy_image_2d_dict[galaxy_pix] == pytest.approx(
                inversion.reconstructed_data_2d, 1.0e-4)

            assert fit.model_image_2d == pytest.approx(fit.galaxy_image_2d_dict[g0] +
                                                       fit.galaxy_image_2d_dict[g1] +
                                                       inversion.reconstructed_data_2d,
                                                       1.0e-4)

        def test___all_lens_fit_quantities__include_hyper_methods(self, lens_data_5x5):
            galaxy_light = g.Galaxy(redshift=0.5,
                                    light_profile=lp.EllipticalSersic(intensity=1.0),
                                    hyper_galaxy=g.HyperGalaxy(contribution_factor=1.0,
                                                               noise_factor=1.0,
                                                               noise_power=1.0),
                                    hyper_model_image_1d=np.ones(9),
                                    hyper_galaxy_image_1d=np.ones(9),
                                    hyper_minimum_value=0.0)

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[galaxy_light],
                                                         source_galaxies=[galaxy_pix],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack,
                                                         border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            hyper_noise_1d = tracer.hyper_noise_map_1d_from_noise_map_1d(
                noise_map_1d=lens_data_5x5.noise_map_1d)
            hyper_noise_map_1d = lens_data_5x5.noise_map_1d + hyper_noise_1d
            hyper_noise_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=hyper_noise_map_1d)
            assert hyper_noise_map_1d == pytest.approx(fit.noise_map_1d, 1e-4)
            assert hyper_noise_map_2d == pytest.approx(fit.noise_map_2d, 1e-4)

            blurred_profile_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_5x5.convolver_image)

            blurred_profile_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=blurred_profile_image_1d)

            assert blurred_profile_image_1d == pytest.approx(
                fit.blurred_profile_image_1d, 1e-4)
            assert blurred_profile_image_2d == pytest.approx(
                fit.blurred_profile_image_2d, 1e-4)

            profile_subtracted_image_1d = lens_data_5x5.image_1d - blurred_profile_image_1d
            profile_subtracted_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=profile_subtracted_image_1d)

            assert profile_subtracted_image_1d == pytest.approx(
                fit.profile_subtracted_image_1d)
            assert profile_subtracted_image_2d == pytest.approx(
                fit.profile_subtracted_image_2d)

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d, noise_map_1d=hyper_noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix, mapper=mapper,
                regularization=reg)

            model_image_1d = blurred_profile_image_1d + inversion.reconstructed_data_1d
            model_image_2d = lens_data_5x5.map_to_scaled_array(array_1d=model_image_1d)

            assert model_image_1d == pytest.approx(fit.model_image_1d, 1e-4)
            assert model_image_2d == pytest.approx(fit.model_image_2d, 1e-4)

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_5x5.image_1d, mask=lens_data_5x5.mask_1d,
                model_data=model_image_1d)
            residual_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=residual_map_1d)

            assert residual_map_1d == pytest.approx(fit.residual_map_1d, 1e-4)
            assert residual_map_2d == pytest.approx(fit.residual_map_2d, 1e-4)

            chi_squared_map_1d = af.fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d, mask=lens_data_5x5.mask_1d,
                noise_map=hyper_noise_map_1d)
            chi_squared_map_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=chi_squared_map_1d)

            assert chi_squared_map_1d == pytest.approx(fit.chi_squared_map_1d, 1e-4)
            assert chi_squared_map_2d == pytest.approx(fit.chi_squared_map_2d, 1e-4)

            chi_squared = af.fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_2d, mask=lens_data_5x5.mask_2d)

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_5x5.mask_2d, noise_map=hyper_noise_map_2d)

            likelihood = af.fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared, noise_normalization=noise_normalization)

            assert likelihood == pytest.approx(fit.likelihood, 1e-4)

            likelihood_with_regularization = \
                util.likelihood_with_regularization_from_chi_squared_regularization_term_and_noise_normalization(
                    chi_squared=chi_squared,
                    regularization_term=inversion.regularization_term,
                    noise_normalization=noise_normalization)

            assert likelihood_with_regularization == pytest.approx(
                fit.likelihood_with_regularization, 1e-4)

            evidence = util.evidence_from_inversion_terms(
                chi_squared=chi_squared,
                regularization_term=inversion.regularization_term,
                log_curvature_regularization_term=inversion.log_det_curvature_reg_matrix_term,
                log_regularization_term=inversion.log_det_regularization_matrix_term,
                noise_normalization=noise_normalization)

            assert evidence == fit.evidence
            assert evidence == fit.figure_of_merit

        def test___blurred_and_model_images_of_planes_and_unmasked_blurred_profile_image_properties(
                self, lens_data_5x5):
            galaxy_light = g.Galaxy(redshift=0.5,
                                    light_profile=lp.EllipticalSersic(intensity=1.0))

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficients=(1.0,))
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.TracerImageSourcePlanes(lens_galaxies=[galaxy_light],
                                                         source_galaxies=[galaxy_pix],
                                                         image_plane_grid_stack=lens_data_5x5.grid_stack,
                                                         border=None)

            fit = lens_fit.LensDataFit.for_data_and_tracer(lens_data=lens_data_5x5,
                                                           tracer=tracer)

            blurred_profile_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_5x5.convolver_image)

            blurred_profile_image_2d = lens_data_5x5.map_to_scaled_array(
                array_1d=blurred_profile_image_1d)

            profile_subtracted_image_1d = lens_data_5x5.image_1d - blurred_profile_image_1d

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_5x5.grid_stack, border=None)

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d,
                noise_map_1d=lens_data_5x5.noise_map_1d,
                convolver=lens_data_5x5.convolver_mapping_matrix, mapper=mapper,
                regularization=reg)

            assert blurred_profile_image_2d == pytest.approx(
                fit.model_image_2d_of_planes[0], 1.0e-4)
            assert inversion.reconstructed_data_2d == pytest.approx(
                fit.model_image_2d_of_planes[1], 1.0e-4)


class MockTracerPositions:

    def __init__(self, positions, noise=None):
        self.positions = positions
        self.noise = noise


class TestPositionFit:

    def test__x1_positions__mock_position_tracer__maximum_separation_is_correct(self):
        tracer = MockTracerPositions(positions=[np.array([[0.0, 0.0], [0.0, 1.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == 1.0

        tracer = MockTracerPositions(positions=[np.array([[0.0, 0.0], [1.0, 1.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(2)

        tracer = MockTracerPositions(positions=[np.array([[0.0, 0.0], [1.0, 3.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(np.square(1.0) + np.square(3.0))

        tracer = MockTracerPositions(positions=[np.array([[-2.0, -4.0], [1.0, 3.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(np.square(3.0) + np.square(7.0))

        tracer = MockTracerPositions(positions=[np.array([[8.0, 4.0], [-9.0, -4.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(np.square(17.0) + np.square(8.0))

    def test_multiple_positions__mock_position_tracer__maximum_separation_is_correct(
            self):
        tracer = MockTracerPositions(
            positions=[np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 0.5]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == 1.0

        tracer = MockTracerPositions(
            positions=[np.array([[0.0, 0.0], [0.0, 0.0], [3.0, 3.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(18)

        tracer = MockTracerPositions(
            positions=[np.array([[0.0, 0.0], [1.0, 1.0], [3.0, 3.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(18)

        tracer = MockTracerPositions(
            positions=[np.array([[-2.0, -4.0], [1.0, 3.0], [0.1, 0.1], [-0.1, -0.1],
                                 [0.3, 0.4], [-0.6, 0.5]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(np.square(3.0) + np.square(7.0))

        tracer = MockTracerPositions(
            positions=[np.array([[8.0, 4.0], [8.0, 4.0], [-9.0, -4.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.maximum_separations[0] == np.sqrt(np.square(17.0) + np.square(8.0))

    def test_multiple_sets_of_positions__multiple_sets_of_max_distances(self):
        tracer = MockTracerPositions(
            positions=[np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 0.5]]),
                       np.array([[0.0, 0.0], [0.0, 0.0], [3.0, 3.0]]),
                       np.array([[0.0, 0.0], [1.0, 1.0], [3.0, 3.0]])])

        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)

        assert fit.maximum_separations[0] == 1.0
        assert fit.maximum_separations[1] == np.sqrt(18)
        assert fit.maximum_separations[2] == np.sqrt(18)

    def test__likelihood__is_sum_of_separations_divided_by_noise(self):
        tracer = MockTracerPositions(
            positions=[np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 0.5]]),
                       np.array([[0.0, 0.0], [0.0, 0.0], [3.0, 3.0]]),
                       np.array([[0.0, 0.0], [1.0, 1.0], [3.0, 3.0]])])

        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)
        assert fit.chi_squared_map[0] == 1.0
        assert fit.chi_squared_map[1] == pytest.approx(18.0, 1e-4)
        assert fit.chi_squared_map[2] == pytest.approx(18.0, 1e-4)
        assert fit.figure_of_merit == pytest.approx(-0.5 * (1.0 + 18 + 18), 1e-4)

        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=2.0)
        assert fit.chi_squared_map[0] == (1.0 / 2.0) ** 2.0
        assert fit.chi_squared_map[1] == pytest.approx((np.sqrt(18.0) / 2.0) ** 2.0,
                                                       1e-4)
        assert fit.chi_squared_map[2] == pytest.approx((np.sqrt(18.0) / 2.0) ** 2.0,
                                                       1e-4)
        assert fit.figure_of_merit == pytest.approx(
            -0.5 * ((1.0 / 2.0) ** 2.0 + (np.sqrt(18.0) / 2.0) ** 2.0 +
                    (np.sqrt(18.0) / 2.0) ** 2.0), 1e-4)

    def test__threshold__if_not_met_returns_ray_tracing_exception(self):
        tracer = MockTracerPositions(positions=[np.array([[0.0, 0.0], [0.0, 1.0]])])
        fit = lens_fit.LensPositionFit(positions=tracer.positions, noise_map=1.0)

        assert fit.maximum_separation_within_threshold(threshold=100.0)
        assert not fit.maximum_separation_within_threshold(threshold=0.1)
