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

from lenstronomy.Data.imaging_data import ImageData
import lenstronomy.Util.util as util


class TestData(object):
    def setup_method(self):
        self.numPix = 10
        kwargs_data = {
            "image_data": np.zeros((self.numPix, self.numPix)),
            "noise_map": np.ones((self.numPix, self.numPix)),
        }
        self.Data = ImageData(**kwargs_data)

    def test_numData(self):
        assert self.Data.num_pixel == self.numPix**2

    def test_shift_coords(self):
        numPix = 10
        deltaPix = 0.05
        (
            x_grid,
            y_grid,
            ra_at_xy_0,
            dec_at_xy_0,
            x_at_radec_0,
            y_at_radec_0,
            Mpix2coord,
            Mcoord2pix,
        ) = util.make_grid_with_coordtransform(
            numPix=numPix, deltapix=deltaPix, subgrid_res=1, inverse=True
        )
        # mask (1= model this pixel, 0= leave blanck)

        kwargs_data = {
            "ra_at_xy_0": ra_at_xy_0,
            "dec_at_xy_0": dec_at_xy_0,
            "transform_pix2angle": Mpix2coord,
            "image_data": np.ones((numPix, numPix)),
        }
        data = ImageData(**kwargs_data)

        ra_shift = 0.05
        dec_shift = 0.0
        kwargs_data["ra_shift"] = ra_shift
        kwargs_data["dec_shift"] = dec_shift
        data_shift = ImageData(**kwargs_data)

        ra, dec = data.map_pix2coord(1, 1)
        ra_new, dec_new = data_shift.map_pix2coord(1, 1)
        npt.assert_almost_equal(ra_new - ra, ra_shift, decimal=10)
        npt.assert_almost_equal(dec_new - dec, dec_shift, decimal=10)

        ra_2, dec_2 = data_shift.map_pix2coord(2, 1)
        npt.assert_almost_equal(ra, ra_2, decimal=10)
        npt.assert_almost_equal(dec, dec_2, decimal=10)

        x, y = data.map_coord2pix(0, 0)
        x_new, y_new = data_shift.map_coord2pix(ra_shift, dec_shift)
        npt.assert_almost_equal(x, x_new, decimal=10)
        npt.assert_almost_equal(y, y_new, decimal=10)

    def test_shift_coordinate_system(self):
        x_shift = 0.05
        y_shift = 0

        numPix = 10
        deltaPix = 0.05
        (
            x_grid,
            y_grid,
            ra_at_xy_0,
            dec_at_xy_0,
            x_at_radec_0,
            y_at_radec_0,
            Mpix2coord,
            Mcoord2pix,
        ) = util.make_grid_with_coordtransform(
            numPix=numPix, deltapix=deltaPix, subgrid_res=1, inverse=True
        )

        kwargs_data = {
            "ra_at_xy_0": ra_at_xy_0,
            "dec_at_xy_0": dec_at_xy_0,
            "transform_pix2angle": Mpix2coord,
            "image_data": np.ones((numPix, numPix)),
        }

        data = ImageData(**kwargs_data)
        data_new = copy.deepcopy(data)
        data_new.shift_coordinate_system(x_shift, y_shift, pixel_unit=False)
        ra, dec = 0, 0
        x, y = data.map_coord2pix(ra, dec)
        x_new, y_new = data_new.map_coord2pix(ra + x_shift, dec + y_shift)
        npt.assert_almost_equal(x, x_new, decimal=10)
        npt.assert_almost_equal(y, y_new, decimal=10)

        ra, dec = data.map_pix2coord(x, y)
        ra_new, dec_new = data_new.map_pix2coord(x, y)
        npt.assert_almost_equal(ra, ra_new - x_shift, decimal=10)
        npt.assert_almost_equal(dec, dec_new - y_shift, decimal=10)

        x_coords, y_coords = data.pixel_coordinates
        x_coords_new, y_coords_new = data_new.pixel_coordinates
        npt.assert_almost_equal(x_coords[0], x_coords_new[0] - x_shift, decimal=10)
        npt.assert_almost_equal(y_coords[0], y_coords_new[0] - y_shift, decimal=10)

    def test_update_data(self):
        kwargs_data = {
            "image_data": np.zeros((self.numPix, self.numPix)),
            "noise_map": None,
            "exposure_time": 1,
            "background_rms": 1,
        }
        data = ImageData(**kwargs_data)
        C_D = data.C_D
        data.update_data(image_data=np.ones((self.numPix, self.numPix)))
        C_D_new = data.C_D
        assert C_D_new[0, 0] > C_D[0, 0]
        data_new = data.data
        npt.assert_almost_equal(data_new, np.ones((self.numPix, self.numPix)))


class TestRaise(unittest.TestCase):
    def test_raise(self):
        kwargs_data = {"image_data": np.zeros((10, 10))}
        Data = ImageData(**kwargs_data)
        image_data_new = np.zeros((5, 5))
        with self.assertRaises(ValueError):
            out = Data.update_data(image_data_new)
        with self.assertRaises(ValueError):
            ImageData(**kwargs_data, likelihood_method="WRONG")


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