import numpy as np
import pytest

from autofit.tools import fit_util
from autolens.data import ccd
from autolens.data.array import mask as msk
from autolens.model.galaxy import galaxy as g
from autolens.lens.util import lens_fit_util
from workspace_jam.scraps.stack import (
    ray_tracing_stack,
    lens_data_stack as lds,
    lens_fit_stack,
)
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


@pytest.fixture(name="lens_data_stack_blur")
def make_li_blur_stack():

    image_0 = np.array(
        [
            [0.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 1.0, 0.0],
            [0.0, 1.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 0.0],
        ]
    )
    psf_0 = ccd.PSF(
        array=(np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])),
        pixel_scale=1.0,
        renormalize=False,
    )
    ccd_0 = ccd.CCDData(image_0, pixel_scale=1.0, psf=psf_0, noise_map=np.ones((4, 4)))

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

    image_1 = np.array(
        [
            [0.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 1.0, 0.0],
            [0.0, 1.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 0.0],
        ]
    )
    psf_1 = ccd.PSF(
        array=(np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])),
        pixel_scale=1.0,
        renormalize=False,
    )
    ccd_1 = ccd.CCDData(image_1, pixel_scale=1.0, psf=psf_1, noise_map=np.ones((4, 4)))

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

    return lds.LensDataStack(
        ccd_datas=[ccd_0, ccd_1], masks=[mask_0, mask_1], sub_grid_size=1
    )


@pytest.fixture(name="lens_data_stack_manual")
def make_li_manual_stack():

    image_0 = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 2.0, 3.0, 0.0],
            [0.0, 4.0, 5.0, 6.0, 0.0],
            [0.0, 7.0, 8.0, 9.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0],
        ]
    )
    psf_0 = ccd.PSF(
        array=(np.array([[1.0, 5.0, 9.0], [2.0, 5.0, 1.0], [3.0, 4.0, 0.0]])),
        pixel_scale=1.0,
    )
    ccd_0 = ccd.CCDData(image_0, pixel_scale=1.0, psf=psf_0, noise_map=np.ones((5, 5)))
    mask_0 = msk.Mask(
        array=np.array(
            [
                [True, True, True, True, True],
                [True, False, False, False, True],
                [True, False, False, False, True],
                [True, False, False, False, True],
                [True, True, True, True, True],
            ]
        ),
        pixel_scale=1.0,
    )

    image_1 = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 2.0, 3.0, 0.0],
            [0.0, 4.0, 6.0, 6.0, 0.0],
            [0.0, 7.0, 8.0, 9.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0],
        ]
    )
    psf_1 = ccd.PSF(
        array=(np.array([[1.0, 1.0, 1.0], [2.0, 1.0, 1.0], [3.0, 1.0, 0.0]])),
        pixel_scale=1.0,
    )
    ccd_1 = ccd.CCDData(image_1, pixel_scale=1.0, psf=psf_1, noise_map=np.ones((5, 5)))
    mask_1 = msk.Mask(
        array=np.array(
            [
                [True, True, True, True, True],
                [True, False, False, False, True],
                [True, False, False, False, True],
                [True, False, False, True, True],
                [True, True, True, True, True],
            ]
        ),
        pixel_scale=1.0,
    )

    return lds.LensDataStack(
        ccd_datas=[ccd_0, ccd_1], masks=[mask_0, mask_1], sub_grid_size=1
    )


class TestLensProfileFit:
    class TestLikelihood:
        def test__image__tracing_fits_data_perfectly__no_psf_blurring__lh_is_noise_normalization(
            self
        ):

            psf_0 = 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_0 = ccd.CCDData(
                np.ones((3, 3)), pixel_scale=1.0, psf=psf_0, noise_map=np.ones((3, 3))
            )
            mask_0 = msk.Mask(
                array=np.array(
                    [[True, True, True], [True, False, True], [True, True, True]]
                ),
                pixel_scale=1.0,
            )

            psf_1 = 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_1 = ccd.CCDData(
                np.ones((3, 3)), pixel_scale=1.0, psf=psf_1, noise_map=np.ones((3, 3))
            )
            mask_1 = msk.Mask(
                array=np.array(
                    [[True, True, True], [True, False, True], [True, True, True]]
                ),
                pixel_scale=1.0,
            )

            li_stack = lds.LensDataStack(
                ccd_datas=[ccd_0, ccd_1], masks=[mask_0, mask_1], sub_grid_size=1
            )

            g0 = g.Galaxy(light_profile=MockLightProfile(value=1.0))
            tracer = ray_tracing_stack.TracerImagePlaneStack(
                lens_galaxies=[g0], image_plane_grid_stacks=li_stack.grid_stacks
            )

            fit = lens_fit_stack.LensProfileFitStack(
                lens_data_stack=li_stack, tracer=tracer
            )
            assert fit.likelihood == 2.0 * -0.5 * np.log(2 * np.pi * 1.0)

        def test__1x2_image__tracing_fits_data_with_chi_sq_5_and_chi_sq_4(self):

            psf_0 = 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_0 = ccd.CCDData(
                5.0 * np.ones((3, 4)),
                pixel_scale=1.0,
                psf=psf_0,
                noise_map=np.ones((3, 4)),
            )
            ccd_0.image[1, 2] = 4.0
            mask_0 = msk.Mask(
                array=np.array(
                    [
                        [True, True, True, True],
                        [True, False, False, True],
                        [True, True, True, True],
                    ]
                ),
                pixel_scale=1.0,
            )

            psf_1 = 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_1 = ccd.CCDData(
                5.0 * np.ones((3, 4)),
                pixel_scale=1.0,
                psf=psf_1,
                noise_map=np.ones((3, 4)),
            )
            mask_1 = msk.Mask(
                array=np.array(
                    [
                        [True, True, True, True],
                        [True, False, False, True],
                        [True, True, True, True],
                    ]
                ),
                pixel_scale=1.0,
            )

            li_stack = lds.LensDataStack(
                ccd_datas=[ccd_0, ccd_1], masks=[mask_0, mask_1], sub_grid_size=1
            )

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

            g0 = g.Galaxy(light_profile=MockLightProfile(value=1.0, size=2))
            tracer = ray_tracing_stack.TracerImagePlaneStack(
                lens_galaxies=[g0], image_plane_grid_stacks=li_stack.grid_stacks
            )

            fit = lens_fit_stack.LensProfileFitStack(
                lens_data_stack=li_stack, tracer=tracer
            )

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

    class TestCompareToManual:
        def test___manual_image_and_psf(self, lens_data_stack_manual):

            g0 = g.Galaxy(light_profile=lp.EllipticalSersic(intensity=1.0))
            g1 = g.Galaxy(mass_profile=mp.SphericalIsothermal(einstein_radius=1.0))

            tracer = ray_tracing_stack.TracerImageSourcePlanesStack(
                lens_galaxies=[g0, g1],
                source_galaxies=[g0],
                image_plane_grid_stacks=lens_data_stack_manual.grid_stacks,
            )

            padded_tracer = ray_tracing_stack.TracerImageSourcePlanesStack(
                lens_galaxies=[g0, g1],
                source_galaxies=[g0],
                image_plane_grid_stacks=lens_data_stack_manual.padded_grid_stacks,
            )

            fit = lens_fit_stack.fit_lens_image_stack_with_tracer(
                lens_data_stack=lens_data_stack_manual,
                tracer=tracer,
                padded_tracer=padded_tracer,
            )

            assert lens_data_stack_manual.noise_maps[0] == pytest.approx(
                fit.noise_maps[0], 1e-4
            )
            assert lens_data_stack_manual.noise_maps[1] == pytest.approx(
                fit.noise_maps[1], 1e-4
            )

            model_image_1d_0 = lens_fit_util.blurred_image_1d_from_1d_unblurred_and_blurring_images(
                unblurred_image_1d=tracer.image_plane_images_1d[0],
                blurring_image_1d=tracer.image_plane_blurring_images_1d[0],
                convolver=lens_data_stack_manual.convolvers_image[0],
            )

            model_image_1d_1 = lens_fit_util.blurred_image_1d_from_1d_unblurred_and_blurring_images(
                unblurred_image_1d=tracer.image_plane_images_1d[1],
                blurring_image_1d=tracer.image_plane_blurring_images_1d[1],
                convolver=lens_data_stack_manual.convolvers_image[1],
            )

            model_image_0 = lens_data_stack_manual.map_to_scaled_arrays[0](
                array_1d=model_image_1d_0
            )
            model_image_1 = lens_data_stack_manual.map_to_scaled_arrays[1](
                array_1d=model_image_1d_1
            )

            assert model_image_0 == pytest.approx(fit.model_images[0], 1e-4)
            assert model_image_1 == pytest.approx(fit.model_images[1], 1e-4)

            residual_map_0 = fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_stack_manual.images[0],
                mask=lens_data_stack_manual.masks[0],
                model_data=model_image_0,
            )
            residual_map_1 = fit_util.residual_map_from_data_mask_and_model_data(
                data=lens_data_stack_manual.images[1],
                mask=lens_data_stack_manual.masks[1],
                model_data=model_image_1,
            )

            assert residual_map_0 == pytest.approx(fit.residual_maps[0], 1e-4)
            assert residual_map_1 == pytest.approx(fit.residual_maps[1], 1e-4)

            chi_squared_map_0 = fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_0,
                mask=lens_data_stack_manual.masks[0],
                noise_map=lens_data_stack_manual.noise_maps[0],
            )
            chi_squared_map_1 = fit_util.chi_squared_map_from_residual_map_noise_map_and_mask(
                residual_map=residual_map_1,
                mask=lens_data_stack_manual.masks[1],
                noise_map=lens_data_stack_manual.noise_maps[1],
            )

            assert chi_squared_map_0 == pytest.approx(fit.chi_squared_maps[0], 1e-4)
            assert chi_squared_map_1 == pytest.approx(fit.chi_squared_maps[1], 1e-4)

            chi_squared_0 = fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_0, mask=lens_data_stack_manual.masks[0]
            )
            noise_normalization_0 = fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=lens_data_stack_manual.noise_maps[0],
                mask=lens_data_stack_manual.masks[0],
            )
            likelihood_0 = fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared_0, noise_normalization=noise_normalization_0
            )

            chi_squared_1 = fit_util.chi_squared_from_chi_squared_map_and_mask(
                chi_squared_map=chi_squared_map_1, mask=lens_data_stack_manual.masks[1]
            )
            noise_normalization_1 = fit_util.noise_normalization_from_noise_map_and_mask(
                noise_map=lens_data_stack_manual.noise_maps[1],
                mask=lens_data_stack_manual.masks[1],
            )
            likelihood_1 = fit_util.likelihood_from_chi_squared_and_noise_normalization(
                chi_squared=chi_squared_1, noise_normalization=noise_normalization_1
            )

            assert fit.chi_squareds == [chi_squared_0, chi_squared_1]
            assert fit.noise_normalizations == [
                noise_normalization_0,
                noise_normalization_1,
            ]
            assert fit.likelihoods == [likelihood_0, likelihood_1]

            assert likelihood_0 + likelihood_1 == pytest.approx(fit.likelihood, 1e-4)
            assert likelihood_0 + likelihood_1 == fit.figure_of_merit
            blurred_image_of_planes_0 = lens_fit_util.blurred_image_of_planes_from_1d_images_and_convolver(
                total_planes=tracer.total_planes,
                image_plane_image_1d_of_planes=tracer.image_plane_images_1d_of_planes[
                    0
                ],
                image_plane_blurring_image_1d_of_planes=tracer.image_plane_blurring_images_1d_of_planes[
                    0
                ],
                convolver=lens_data_stack_manual.convolvers_image[0],
                map_to_scaled_array=lens_data_stack_manual.map_to_scaled_arrays[0],
            )

            blurred_image_of_planes_1 = lens_fit_util.blurred_image_of_planes_from_1d_images_and_convolver(
                total_planes=tracer.total_planes,
                image_plane_image_1d_of_planes=tracer.image_plane_images_1d_of_planes[
                    1
                ],
                image_plane_blurring_image_1d_of_planes=tracer.image_plane_blurring_images_1d_of_planes[
                    1
                ],
                convolver=lens_data_stack_manual.convolvers_image[1],
                map_to_scaled_array=lens_data_stack_manual.map_to_scaled_arrays[1],
            )

            assert (
                blurred_image_of_planes_0[0] == fit.model_images_of_planes[0][0]
            ).all()
            assert (
                blurred_image_of_planes_0[1] == fit.model_images_of_planes[0][1]
            ).all()
            assert (
                blurred_image_of_planes_1[0] == fit.model_images_of_planes[1][0]
            ).all()
            assert (
                blurred_image_of_planes_1[1] == fit.model_images_of_planes[1][1]
            ).all()

            unmasked_blurred_image_0 = lens_fit_util.unmasked_blurred_image_from_padded_grid_stack_psf_and_unmasked_image(
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[0],
                psf=lens_data_stack_manual.psfs[0],
                unmasked_image_1d=padded_tracer.image_plane_images_1d[0],
            )

            unmasked_blurred_image_1 = lens_fit_util.unmasked_blurred_image_from_padded_grid_stack_psf_and_unmasked_image(
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[1],
                psf=lens_data_stack_manual.psfs[1],
                unmasked_image_1d=padded_tracer.image_plane_images_1d[1],
            )

            assert (unmasked_blurred_image_0 == fit.unmasked_model_images[0]).all()
            assert (unmasked_blurred_image_1 == fit.unmasked_model_images[1]).all()

            unmasked_blurred_image_of_galaxies_i0 = lens_fit_util.unmasked_blurred_image_of_galaxies_from_psf_and_unmasked_1d_galaxy_images(
                galaxies=padded_tracer.image_plane.galaxies,
                image_plane_image_1d_of_galaxies=padded_tracer.image_plane.image_plane_images_1d_of_galaxies[
                    0
                ],
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[0],
                psf=lens_data_stack_manual.psfs[0],
            )

            unmasked_blurred_image_of_galaxies_s0 = lens_fit_util.unmasked_blurred_image_of_galaxies_from_psf_and_unmasked_1d_galaxy_images(
                galaxies=padded_tracer.source_plane.galaxies,
                image_plane_image_1d_of_galaxies=padded_tracer.source_plane.image_plane_images_1d_of_galaxies[
                    0
                ],
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[0],
                psf=lens_data_stack_manual.psfs[0],
            )

            unmasked_blurred_image_of_galaxies_i1 = lens_fit_util.unmasked_blurred_image_of_galaxies_from_psf_and_unmasked_1d_galaxy_images(
                galaxies=padded_tracer.image_plane.galaxies,
                image_plane_image_1d_of_galaxies=padded_tracer.image_plane.image_plane_images_1d_of_galaxies[
                    0
                ],
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[1],
                psf=lens_data_stack_manual.psfs[1],
            )

            unmasked_blurred_image_of_galaxies_s1 = lens_fit_util.unmasked_blurred_image_of_galaxies_from_psf_and_unmasked_1d_galaxy_images(
                galaxies=padded_tracer.source_plane.galaxies,
                image_plane_image_1d_of_galaxies=padded_tracer.source_plane.image_plane_images_1d_of_galaxies[
                    0
                ],
                padded_grid_stack=lens_data_stack_manual.padded_grid_stacks[1],
                psf=lens_data_stack_manual.psfs[1],
            )

            assert (
                unmasked_blurred_image_of_galaxies_i0[0]
                == fit.unmasked_model_images_of_planes_and_galaxies[0][0][0]
            ).all()
            assert (
                unmasked_blurred_image_of_galaxies_s0[0]
                == fit.unmasked_model_images_of_planes_and_galaxies[0][1][0]
            ).all()
            assert (
                unmasked_blurred_image_of_galaxies_i1[0]
                == fit.unmasked_model_images_of_planes_and_galaxies[1][0][0]
            ).all()
            assert (
                unmasked_blurred_image_of_galaxies_s1[0]
                == fit.unmasked_model_images_of_planes_and_galaxies[1][1][0]
            ).all()
