from lenstronomy.SimulationAPI.data_api import DataAPI
import lenstronomy.Util.util as util
import numpy as np
import unittest


class TestDataAPI(object):
    def setup_method(self):
        numpix = 10
        self.ccd_gain = 4.0
        self.pixel_scale = 0.13
        self.read_noise = 10.0
        kwargs_instrument = {
            "read_noise": self.read_noise,
            "pixel_scale": self.pixel_scale,
            "ccd_gain": self.ccd_gain,
        }

        exposure_time = 100
        sky_brightness = 20.0
        self.magnitude_zero_point = 21.0
        num_exposures = 2
        seeing = 0.9
        kwargs_observations = {
            "exposure_time": exposure_time,
            "sky_brightness": sky_brightness,
            "magnitude_zero_point": self.magnitude_zero_point,
            "num_exposures": num_exposures,
            "seeing": seeing,
            "psf_type": "GAUSSIAN",
            "kernel_point_source": None,
        }
        self.kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations)
        self.api = DataAPI(numpix=numpix, data_count_unit="ADU", **self.kwargs_data)

        kwargs_observations = {
            "exposure_time": exposure_time,
            "sky_brightness": sky_brightness,
            "magnitude_zero_point": self.magnitude_zero_point,
            "num_exposures": num_exposures,
            "seeing": seeing,
            "psf_type": "PIXEL",
            "kernel_point_source": np.ones((3, 3)),
        }
        kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations)
        self.api_pixel = DataAPI(numpix=numpix, data_count_unit="ADU", **kwargs_data)

        self.ra_at_xy_0 = 0.02
        self.dec_at_xy_0 = 0.02
        self.transform_pix2angle = [[-self.pixel_scale, 0], [0, self.pixel_scale]]
        kwargs_pixel_grid = {
            "ra_at_xy_0": self.ra_at_xy_0,
            "dec_at_xy_0": self.dec_at_xy_0,
            "transform_pix2angle": self.transform_pix2angle,
        }
        self.api_pixel_grid = DataAPI(
            numpix=numpix,
            kwargs_pixel_grid=kwargs_pixel_grid,
            data_count_unit="ADU",
            **self.kwargs_data
        )

    def test_data_class(self):
        data_class = self.api.data_class
        assert data_class.pixel_width == self.pixel_scale

    def test_psf_class(self):
        psf_class = self.api.psf_class
        assert psf_class.psf_type == "GAUSSIAN"
        psf_class = self.api_pixel.psf_class
        assert psf_class.psf_type == "PIXEL"

    def test_kwargs_data(self):
        kwargs_data = self.api.kwargs_data
        assert kwargs_data["ra_at_xy_0"] != self.ra_at_xy_0
        kwargs_data = self.api_pixel_grid.kwargs_data
        assert kwargs_data["ra_at_xy_0"] == self.ra_at_xy_0


class TestRaise(unittest.TestCase):
    def test_raise(self):
        numpix = 10
        self.ccd_gain = 4.0
        self.pixel_scale = 0.13
        self.read_noise = 10.0
        kwargs_instrument = {
            "read_noise": self.read_noise,
            "pixel_scale": self.pixel_scale,
            "ccd_gain": self.ccd_gain,
        }

        exposure_time = 100
        sky_brightness = 20.0
        magnitude_zero_point = 21.0
        num_exposures = 2
        seeing = 0.9
        kwargs_observations = {
            "exposure_time": exposure_time,
            "sky_brightness": sky_brightness,
            "magnitude_zero_point": magnitude_zero_point,
            "num_exposures": num_exposures,
            "seeing": seeing,
            "psf_type": "wrong",
            "kernel_point_source": None,
        }
        kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations)
        data_api = DataAPI(numpix=numpix, data_count_unit="ADU", **kwargs_data)
        print(data_api._psf_type)
        with self.assertRaises(ValueError):
            data_api = DataAPI(numpix=numpix, data_count_unit="ADU", **kwargs_data)
            psf_class = data_api.psf_class

        kwargs_observations = {
            "exposure_time": exposure_time,
            "sky_brightness": sky_brightness,
            "magnitude_zero_point": magnitude_zero_point,
            "num_exposures": num_exposures,
            "seeing": seeing,
            "psf_type": "PIXEL",
            "kernel_point_source": None,
        }
        kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations)
        with self.assertRaises(ValueError):
            data_api = DataAPI(numpix=numpix, data_count_unit="ADU", **kwargs_data)
            psf_class = data_api.psf_class

        kwargs_data["kernel_point_source"] = np.ones((3, 3))
        kwargs_pixel_grid = {"ra_at_xy_0": 0.02, "dec_at_xy_0": 0.02}
        with self.assertRaises(ValueError):
            data_api = DataAPI(
                numpix=numpix, kwargs_pixel_grid=kwargs_pixel_grid, **kwargs_data
            )
