import numpy as np
import pytest

import autofit as af
from autolens.lens.util import lens_fit_util as util
from autolens.data.instrument import abstract_data
from autolens.data.instrument 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.hyper import hyper_data as hi
from autolens.model.galaxy import galaxy as g
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 = abstract_data.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_7x7 = 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (
                fit.mask(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([4.0, 3.0])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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 = abstract_data.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_7x7 = 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (
                fit.mask(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([4.0, 0.0])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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 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 = abstract_data.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_7x7 = 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (
                fit.mask(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([2.0, 0.0])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == np.array(
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 2.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(return_in_2d=True)
                == 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))

        def test__hyper_image_changes_background_sky__reflected_in_likelihood(self):

            psf = abstract_data.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_7x7 = 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

            hyper_image_sky = hi.HyperImageSky(sky_scale=1.0)

            fit = lens_fit.LensProfileFit(
                lens_data=lens_data_7x7, tracer=tracer, hyper_image_sky=hyper_image_sky
            )

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

            assert (fit.image_1d == np.array([6.0, 5.0])).all()
            assert (
                fit.image(return_in_2d=True)
                == np.array(
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 6.0, 5.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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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([5.0, 4.0])).all()
            assert (
                fit.residual_map(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([5.0, 4.0])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == 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.chi_squared_map_1d == np.array([25.0, 16.0])).all()
            assert (
                fit.chi_squared_map(return_in_2d=True)
                == np.array(
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 25.0, 16.0, 0.0], [0.0, 0.0, 0.0, 0.0]]
                )
            ).all()

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

        def test__hyper_background_changes_background_noise_map__reflected_in_likelihood(
            self
        ):

            psf = abstract_data.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_7x7 = 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

            hyper_background_noise = hi.HyperBackgroundNoise(noise_scale=1.0)

            fit = lens_fit.LensProfileFit(
                lens_data=lens_data_7x7,
                tracer=tracer,
                hyper_background_noise=hyper_background_noise,
            )

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (
                fit.mask(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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, 1.0])).all()
            assert (
                fit.model_image(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([2.0, 1.5])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == np.array(
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 2.0, 1.5, 0.0], [0.0, 0.0, 0.0, 0.0]]
                )
            ).all()

            assert (fit.chi_squared_map_1d == np.array([4.0, 2.25])).all()
            assert (
                fit.chi_squared_map(return_in_2d=True)
                == np.array(
                    [[0.0, 0.0, 0.0, 0.0], [0.0, 4.0, 2.25, 0.0], [0.0, 0.0, 0.0, 0.0]]
                )
            ).all()

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

        def test__hyper_noise_map_max_changes_noise_map__reflected_in_likelihood(self):

            psf = abstract_data.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_7x7 = ld.LensData(
                ccd_data=ccd_data, mask=mask, sub_grid_size=1, hyper_noise_map_max=1.0
            )

            # 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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

            hyper_background_noise = hi.HyperBackgroundNoise(noise_scale=1.0)

            fit = lens_fit.LensProfileFit(
                lens_data=lens_data_7x7,
                tracer=tracer,
                hyper_background_noise=hyper_background_noise,
            )

            assert (fit.mask_1d == np.array([False, False])).all()
            assert (
                fit.mask(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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.normalized_residual_map_1d == np.array([4.0, 3.0])).all()
            assert (
                fit.normalized_residual_map(return_in_2d=True)
                == 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(return_in_2d=True)
                == 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)
            )

    class TestCompareToManual:
        def test___all_lens_fit_quantities__no_hyper_methods(self, lens_data_7x7):

            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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            assert lens_data_7x7.noise_map(return_in_2d=True) == pytest.approx(
                fit.noise_map(return_in_2d=True)
            )

            model_image_1d = tracer.blurred_profile_image_plane_image_1d_from_convolver_image(
                convolver_image=lens_data_7x7.convolver_image
            )
            model_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_7x7.image_1d,
                mask=lens_data_7x7.mask_1d,
                model_data=model_image_1d,
            )
            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )

            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=lens_data_7x7.noise_map(return_in_2d=True),
                mask=lens_data_7x7.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_7x7
        ):
            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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1, g2], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            g0_image_plane_image_1d = g0.intensities_from_grid(
                grid=lens_data_7x7.grid_stack.sub,
                return_in_2d=False,
                return_binned=True,
            )
            g0_image_plane_bluring_image_1d = g0.intensities_from_grid(
                grid=lens_data_7x7.grid_stack.blurring,
                return_in_2d=False,
                return_binned=False,
            )

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

            g1_image_plane_image_1d = g1.intensities_from_grid(
                grid=tracer.source_plane.grid_stack.sub,
                return_in_2d=False,
                return_binned=True,
            )
            g1_image_plane_bluring_image_1d = g1.intensities_from_grid(
                grid=tracer.source_plane.grid_stack.blurring,
                return_in_2d=False,
                return_binned=False,
            )

            g1_blurred_image_plane_image_1d = lens_data_7x7.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_7x7.scaled_array_2d_from_array_1d(
                array_1d=g0_blurred_image_plane_image_1d
            )
            g1_blurred_image_plane_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True) == 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_7x7
        ):

            hyper_noise_map_max = 0.2
            lens_data_7x7.hyper_noise_map_max = hyper_noise_map_max

            hyper_image_sky = hi.HyperImageSky(sky_scale=1.0)

            hyper_background_noise = hi.HyperBackgroundNoise(noise_scale=1.0)

            image_1d = hyper_image_sky.image_scaled_sky_from_image(
                image=lens_data_7x7.image_1d
            )

            hyper_noise_map_background_1d = hyper_background_noise.noise_map_scaled_noise_from_noise_map(
                noise_map=lens_data_7x7.noise_map_1d
            )

            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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

            fit = lens_fit.LensDataFit.for_data_and_tracer(
                lens_data=lens_data_7x7,
                tracer=tracer,
                hyper_image_sky=hyper_image_sky,
                hyper_background_noise=hyper_background_noise,
            )

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

            hyper_noise_map_1d = hyper_noise_map_background_1d + hyper_noise_1d

            hyper_noise_map_1d[
                hyper_noise_map_1d > hyper_noise_map_max
            ] = hyper_noise_map_max

            hyper_noise_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

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

            model_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=image_1d, mask=lens_data_7x7.mask_1d, model_data=model_image_1d
            )
            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )

            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=hyper_noise_map_2d, mask=lens_data_7x7.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_7x7
        ):

            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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1], image_plane_grid_stack=lens_data_7x7.grid_stack
            )

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

            blurred_profile_image_2d_of_planes = tracer.blurred_profile_image_plane_image_2d_of_planes_from_convolver_image(
                convolver_image=lens_data_7x7.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 = tracer.unmasked_blurred_profile_image_plane_image_from_psf(
                psf=lens_data_7x7.psf
            )

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

            unmasked_blurred_profile_image_of_planes = tracer.unmasked_blurred_profile_image_plane_image_of_planes_from_psf(
                psf=lens_data_7x7.psf
            )

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

            unmasked_blurred_profile_image_of_galaxies = tracer.unmasked_blurred_profile_image_plane_image_of_plane_and_galaxies_from_psf(
                psf=lens_data_7x7.psf
            )

            assert (
                unmasked_blurred_profile_image_of_galaxies[0][0]
                == fit.unmasked_blurred_profile_image_plane_image_of_planes_and_galaxies[
                    0
                ][
                    0
                ]
            ).all()
            assert (
                unmasked_blurred_profile_image_of_galaxies[1][0]
                == fit.unmasked_blurred_profile_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_7x7):

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficient=1.0)

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

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g.Galaxy(redshift=0.5), g0],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

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

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

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_7x7.image_1d,
                mask=lens_data_7x7.mask_1d,
                model_data=inversion.reconstructed_data_1d,
            )
            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )
            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_7x7.mask_2d,
                noise_map=lens_data_7x7.noise_map(return_in_2d=True),
            )

            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_7x7
        ):
            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficient=1.0)

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

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

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

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                mapper=mapper,
                regularization=reg,
                image_1d=lens_data_7x7.image_1d,
                noise_map_1d=lens_data_7x7.noise_map_1d,
                convolver=lens_data_7x7.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(return_in_2d=True) == pytest.approx(
                fit.galaxy_image_2d_dict[g1], 1.0e-4
            )

        def test___all_lens_fit_quantities__include_hyper_methods(self, lens_data_7x7):

            hyper_noise_map_max = 0.2
            lens_data_7x7.hyper_noise_map_max = hyper_noise_map_max

            hyper_image_sky = hi.HyperImageSky(sky_scale=1.0)

            hyper_background_noise = hi.HyperBackgroundNoise(noise_scale=1.0)

            image_1d = hyper_image_sky.image_scaled_sky_from_image(
                image=lens_data_7x7.image_1d
            )

            hyper_noise_map_background_1d = hyper_background_noise.noise_map_scaled_noise_from_noise_map(
                noise_map=lens_data_7x7.noise_map_1d
            )

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficient=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.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g.Galaxy(redshift=0.5), g0],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

            fit = lens_fit.LensDataFit.for_data_and_tracer(
                lens_data=lens_data_7x7,
                tracer=tracer,
                hyper_image_sky=hyper_image_sky,
                hyper_background_noise=hyper_background_noise,
            )

            hyper_noise_1d = tracer.hyper_noise_map_1d_from_noise_map_1d(
                noise_map_1d=lens_data_7x7.noise_map_1d
            )
            hyper_noise_map_1d = hyper_noise_map_background_1d + hyper_noise_1d
            hyper_noise_map_1d[
                hyper_noise_map_1d > hyper_noise_map_max
            ] = hyper_noise_map_max

            hyper_noise_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

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

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

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=image_1d,
                mask=lens_data_7x7.mask_1d,
                model_data=inversion.reconstructed_data_1d,
            )
            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )
            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_7x7.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_7x7
        ):

            pix = pixelizations.Rectangular(shape=(3, 3))
            reg = regularization.Constant(coefficient=1.0)

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

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g.Galaxy(redshift=0.5), g0],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

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

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

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


class TestLensProfileInversionFit:
    class TestCompareToManual:
        def test___all_lens_fit_quantities__no_hyper_methods(self, lens_data_7x7):
            galaxy_light = g.Galaxy(
                redshift=0.5, light_profile=lp.EllipticalSersic(intensity=1.0)
            )

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

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[galaxy_light, galaxy_pix],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

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

            blurred_profile_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

            profile_subtracted_image_1d = (
                lens_data_7x7.image_1d - blurred_profile_image_1d
            )
            profile_subtracted_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

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

            model_image_1d = blurred_profile_image_1d + inversion.reconstructed_data_1d
            model_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_7x7.image_1d,
                mask=lens_data_7x7.mask_1d,
                model_data=model_image_1d,
            )
            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=lens_data_7x7.noise_map_1d,
            )
            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_7x7.mask_2d,
                noise_map=lens_data_7x7.noise_map(return_in_2d=True),
            )

            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_7x7
        ):

            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(coefficient=1.0)
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[g0, g1, g2, galaxy_pix],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

            g0_image_plane_image_1d = g0.intensities_from_grid(
                grid=lens_data_7x7.grid_stack.sub,
                return_in_2d=False,
                return_binned=True,
            )
            g0_image_plane_bluring_image_1d = g0.intensities_from_grid(
                grid=lens_data_7x7.grid_stack.blurring,
                return_in_2d=False,
                return_binned=False,
            )

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

            g1_image_plane_image_1d = g1.intensities_from_grid(
                grid=tracer.source_plane.grid_stack.sub,
                return_in_2d=False,
                return_binned=True,
            )
            g1_image_plane_bluring_image_1d = g1.intensities_from_grid(
                grid=tracer.source_plane.grid_stack.blurring,
                return_in_2d=False,
                return_binned=False,
            )

            g1_blurred_image_plane_image_1d = lens_data_7x7.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_7x7.image_1d - blurred_profile_image_1d
            )
            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_7x7.grid_stack, border=None
            )

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d,
                noise_map_1d=lens_data_7x7.noise_map_1d,
                convolver=lens_data_7x7.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_7x7.scaled_array_2d_from_array_1d(
                array_1d=g0_blurred_image_plane_image_1d
            )
            g1_blurred_image_plane_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True) == 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_7x7):

            hyper_noise_map_max = 0.2
            lens_data_7x7.hyper_noise_map_max = hyper_noise_map_max

            hyper_image_sky = hi.HyperImageSky(sky_scale=1.0)

            hyper_background_noise = hi.HyperBackgroundNoise(noise_scale=1.0)

            image_1d = hyper_image_sky.image_scaled_sky_from_image(
                image=lens_data_7x7.image_1d
            )

            hyper_noise_map_background_1d = hyper_background_noise.noise_map_scaled_noise_from_noise_map(
                noise_map=lens_data_7x7.noise_map_1d
            )

            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(coefficient=1.0)
            galaxy_pix = g.Galaxy(redshift=1.0, pixelization=pix, regularization=reg)

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[galaxy_light, galaxy_pix],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

            fit = lens_fit.LensDataFit.for_data_and_tracer(
                lens_data=lens_data_7x7,
                tracer=tracer,
                hyper_image_sky=hyper_image_sky,
                hyper_background_noise=hyper_background_noise,
            )

            hyper_noise_1d = tracer.hyper_noise_map_1d_from_noise_map_1d(
                noise_map_1d=lens_data_7x7.noise_map_1d
            )
            hyper_noise_map_1d = hyper_noise_map_background_1d + hyper_noise_1d
            hyper_noise_map_1d[
                hyper_noise_map_1d > hyper_noise_map_max
            ] = hyper_noise_map_max

            hyper_noise_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

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

            blurred_profile_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

            profile_subtracted_image_1d = image_1d - blurred_profile_image_1d
            profile_subtracted_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

            mapper = pix.mapper_from_grid_stack_and_border(
                grid_stack=lens_data_7x7.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_7x7.convolver_mapping_matrix,
                mapper=mapper,
                regularization=reg,
            )

            model_image_1d = blurred_profile_image_1d + inversion.reconstructed_data_1d
            model_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            residual_map_1d = af.fit_util.residual_map_from_data_mask_and_model_data(
                data=image_1d, mask=lens_data_7x7.mask_1d, model_data=model_image_1d
            )

            residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True))

            normalized_residual_map_1d = af.fit_util.normalized_residual_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1d,
                mask=lens_data_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )

            normalized_residual_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=normalized_residual_map_1d
            )

            assert normalized_residual_map_1d == pytest.approx(
                fit.normalized_residual_map_1d, 1e-4
            )
            assert normalized_residual_map_2d == pytest.approx(
                fit.normalized_residual_map(return_in_2d=True)
            )

            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_7x7.mask_1d,
                noise_map=hyper_noise_map_1d,
            )

            chi_squared_map_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                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(return_in_2d=True)
            )

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

            noise_normalization = af.fit_util.noise_normalization_from_noise_map_and_mask(
                mask=lens_data_7x7.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_7x7
        ):
            galaxy_light = g.Galaxy(
                redshift=0.5, light_profile=lp.EllipticalSersic(intensity=1.0)
            )

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

            tracer = ray_tracing.Tracer.from_galaxies_and_image_plane_grid_stack(
                galaxies=[galaxy_light, galaxy_pix],
                image_plane_grid_stack=lens_data_7x7.grid_stack,
                border=None,
            )

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

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

            blurred_profile_image_2d = lens_data_7x7.scaled_array_2d_from_array_1d(
                array_1d=blurred_profile_image_1d
            )

            profile_subtracted_image_1d = (
                lens_data_7x7.image_1d - blurred_profile_image_1d
            )

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

            inversion = inversions.Inversion.from_data_1d_mapper_and_regularization(
                image_1d=profile_subtracted_image_1d,
                noise_map_1d=lens_data_7x7.noise_map_1d,
                convolver=lens_data_7x7.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)
