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

from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.LightModel.Profiles.gaussian import Gaussian
from lenstronomy.LightModel.Profiles.starlets import SLIT_Starlets
from lenstronomy.Util import util


_force_no_pysap = (
    True  # if issues on Travis-CI to install pysap, force use python-only functions
)


class TestSLITStarlets(object):
    """Class to test SLIT_Starlets light profile."""

    def setup_method(self):
        # different versions of Starlet transforms
        self.starlets = SLIT_Starlets(
            fast_inverse=False, second_gen=False, force_no_pysap=_force_no_pysap
        )
        self.starlets_fast = SLIT_Starlets(
            fast_inverse=True, second_gen=False, force_no_pysap=_force_no_pysap
        )
        self.starlets_2nd = SLIT_Starlets(
            second_gen=True, force_no_pysap=_force_no_pysap
        )

        # define a test image with gaussian components
        self.num_pix = 50
        self.n_scales = 3
        self.n_pixels = self.num_pix**2
        self.x, self.y = util.make_grid(self.num_pix, 1)

        # build a non-trivial positive image from sum of gaussians
        gaussian = Gaussian()
        gaussian1 = gaussian.function(
            self.x, self.y, amp=100, sigma=1, center_x=-7, center_y=-7
        )
        gaussian2 = gaussian.function(
            self.x, self.y, amp=500, sigma=3, center_x=-3, center_y=-3
        )
        gaussian3 = gaussian.function(
            self.x, self.y, amp=2000, sigma=5, center_x=+5, center_y=+5
        )
        self.test_image = util.array2image(gaussian1 + gaussian2 + gaussian3)

        self.test_coeffs = np.zeros((self.n_scales, self.num_pix, self.num_pix))
        self.test_coeffs[0, :, :] = util.array2image(gaussian1)
        self.test_coeffs[1, :, :] = util.array2image(gaussian2)
        self.test_coeffs[2, :, :] = util.array2image(gaussian3)

    def test_reconstructions_2d(self):
        """

        :return:
        """
        # PySAP requires call to decomposition once before anything else
        self.starlets.decomposition_2d(self.test_image, self.n_scales)
        self.starlets_fast.decomposition_2d(self.test_image, self.n_scales)
        self.starlets_2nd.decomposition_2d(self.test_image, self.n_scales)

        image = self.starlets.function_2d(
            coeffs=self.test_coeffs, n_scales=self.n_scales, n_pixels=self.n_pixels
        )
        image_fast = self.starlets_fast.function_2d(
            coeffs=self.test_coeffs, n_scales=self.n_scales, n_pixels=self.n_pixels
        )
        assert image.shape == (self.num_pix, self.num_pix)
        assert image_fast.shape == (self.num_pix, self.num_pix)

        image_2nd = self.starlets_2nd.function_2d(
            coeffs=self.test_coeffs, n_scales=self.n_scales, n_pixels=self.n_pixels
        )
        assert image_2nd.shape == (self.num_pix, self.num_pix)
        assert np.all(image_2nd >= 0)

    def test_decompositions_2d(self):
        """

        :return:
        """
        # test equality between fast and std transform (which are identical)
        coeffs = self.starlets.decomposition_2d(self.test_image, self.n_scales)
        coeffs_fast = self.starlets_fast.decomposition_2d(
            self.test_image, self.n_scales
        )
        assert coeffs.shape == (self.n_scales, self.num_pix, self.num_pix)
        assert coeffs_fast.shape == (self.n_scales, self.num_pix, self.num_pix)
        npt.assert_almost_equal(coeffs, coeffs_fast, decimal=3)

        # test non-negativity of second generation starlet transform
        coeffs_2nd = self.starlets_2nd.decomposition_2d(self.test_image, self.n_scales)
        assert coeffs_2nd.shape == (self.n_scales, self.num_pix, self.num_pix)

    def test_function(self):
        """

        :return:
        """
        # PySAP requires call to decomposition once before anything else
        self.starlets.decomposition(self.test_image, self.n_scales)
        self.starlets_fast.decomposition(self.test_image, self.n_scales)
        self.starlets_2nd.decomposition(self.test_image, self.n_scales)

        # test with a 1D input
        self.starlets.decomposition(util.image2array(self.test_image), self.n_scales)

        coeffs_1d = self.test_coeffs.reshape(self.n_scales * self.num_pix**2)

        image_1d = self.starlets.function(
            self.x,
            self.y,
            amp=coeffs_1d,
            n_scales=self.n_scales,
            n_pixels=self.n_pixels,
        )
        assert image_1d.shape == (self.num_pix**2,)
        image_1d_fast = self.starlets_fast.function(
            self.x,
            self.y,
            amp=coeffs_1d,
            n_scales=self.n_scales,
            n_pixels=self.n_pixels,
        )
        assert image_1d_fast.shape == (self.num_pix**2,)
        image_1d_2nd = self.starlets_2nd.function(
            self.x,
            self.y,
            amp=coeffs_1d,
            n_scales=self.n_scales,
            n_pixels=self.n_pixels,
        )
        assert image_1d_2nd.shape == (self.num_pix**2,)

    def test_identity_operations_fast(self):
        """Test the decomposition/reconstruction.

        :return:
        """
        coeffs = self.starlets_fast.decomposition_2d(self.test_image, self.n_scales)
        test_image_recon = self.starlets_fast.function_2d(
            coeffs=coeffs, n_scales=self.n_scales, n_pixels=self.n_pixels
        )
        npt.assert_almost_equal(self.test_image, test_image_recon, decimal=5)

    def test_identity_operations_2nd(self):
        """Test the decomposition/reconstruction.

        :return:
        """
        coeffs = self.starlets_2nd.decomposition_2d(self.test_image, self.n_scales)
        test_image_recon = self.starlets_2nd.function_2d(
            coeffs=coeffs, n_scales=self.n_scales, n_pixels=self.n_pixels
        )
        npt.assert_almost_equal(self.test_image, test_image_recon, decimal=5)

    def test_delete_cache(self):
        amp = self.test_coeffs.reshape(self.n_scales * self.num_pix**2)
        kwargs_starlets = dict(
            amp=amp,
            n_scales=self.n_scales,
            n_pixels=self.n_pixels,
            center_x=0,
            center_y=0,
            scale=1,
        )
        output = self.starlets_fast.function(self.x, self.y, **kwargs_starlets)
        assert hasattr(self.starlets_fast.interpol, "_image_interp")
        self.starlets_fast.delete_cache()
        assert not hasattr(self.starlets_fast.interpol, "_image_interp")

    def test_coeffs2pysap(self):
        n_scales = 3
        num_pix = 20
        coeffs = np.ones((n_scales, num_pix, num_pix))
        pysap_list = self.starlets._coeffs2pysap(coeffs)
        assert len(pysap_list) == n_scales
        for i in range(n_scales):
            assert pysap_list[i].shape == coeffs[i].shape

    def test_pysap2coeffs(self):
        n_scales = 3
        num_pix = 20
        pysap_list = n_scales * [np.ones((num_pix, num_pix))]
        coeffs = self.starlets._pysap2coeffs(pysap_list)
        assert coeffs.shape == (n_scales, num_pix, num_pix)
        for i in range(n_scales):
            assert pysap_list[i].shape == coeffs[i].shape


class TestRaise(unittest.TestCase):
    def test_raise(self):
        with self.assertRaises(ValueError):
            # try to set decomposition scale to higher than maximal value
            starlets = SLIT_Starlets(force_no_pysap=True)
            # define a test image with gaussian components
            num_pix = 50
            x, y = util.make_grid(num_pix, 1)
            # build a non-trivial positive image from sum of gaussians
            gaussian = Gaussian()
            gaussian1 = gaussian.function(
                x, y, amp=100, sigma=1, center_x=-7, center_y=-7
            )
            gaussian2 = gaussian.function(
                x, y, amp=500, sigma=3, center_x=-3, center_y=-3
            )
            gaussian3 = gaussian.function(
                x, y, amp=2000, sigma=5, center_x=+5, center_y=+5
            )
            test_image = util.array2image(gaussian1 + gaussian2 + gaussian3)
            n_scales = 100
            _ = starlets.decomposition_2d(test_image, n_scales)
        with self.assertRaises(ValueError):
            # try to set decomposition scale to negative value
            starlets = SLIT_Starlets(force_no_pysap=True)
            # define a test image with gaussian components
            num_pix = 50
            x, y = util.make_grid(num_pix, 1)
            # build a non-trivial positive image from sum of gaussians
            gaussian = Gaussian()
            gaussian1 = gaussian.function(
                x, y, amp=100, sigma=1, center_x=-7, center_y=-7
            )
            gaussian2 = gaussian.function(
                x, y, amp=500, sigma=3, center_x=-3, center_y=-3
            )
            gaussian3 = gaussian.function(
                x, y, amp=2000, sigma=5, center_x=+5, center_y=+5
            )
            test_image = util.array2image(gaussian1 + gaussian2 + gaussian3)
            n_scales = -1
            _ = starlets.decomposition_2d(test_image, n_scales)
        with self.assertRaises(ValueError):
            # function_split is not supported/defined for pixel-based profiles
            light_model = LightModel(["SLIT_STARLETS"])
            num_pix = 20
            x, y = util.make_grid(num_pix, 1)
            kwargs_list = [
                {
                    "amp": np.ones((3, num_pix, num_pix)),
                    "n_scales": 3,
                    "n_pixels": 20**2,
                    "center_x": 0,
                    "center_y": 0,
                    "scale": 1,
                }
            ]
            _ = light_model.functions_split(x, y, kwargs_list)
        with self.assertRaises(ValueError):
            # provided a wrong shape for starlet coefficients
            starlet_class = SLIT_Starlets()
            num_pix = 20
            x, y = util.make_grid(num_pix, 1)
            coeffs_wrong = np.ones((3, num_pix**2))
            kwargs_list = {
                "amp": coeffs_wrong,
                "n_scales": 3,
                "n_pixels": 20**2,
                "center_x": 0,
                "center_y": 0,
                "scale": 1,
            }
            _ = starlet_class.function(x, y, **kwargs_list)
            image_wrong = np.ones((1, num_pix, num_pix))
            _ = starlet_class.decomposition(image_wrong, 3)
        with self.assertRaises(ValueError):
            # provided a wrong shape for image to be decomposed
            starlet_class = SLIT_Starlets()
            num_pix = 20
            image_wrong = np.ones((2, num_pix, num_pix))
            _ = starlet_class.decomposition(image_wrong, 3)


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