import unittest
from csst_cpic_sim.target import star_photlam
from csst_cpic_sim.optics import make_focus_image, focal_mask, filter_throughput
from csst_cpic_sim.config import which_focalplane, S
import numpy as np


class TestOptics(unittest.TestCase):
    def test_filter_throughtput(self):
        bandpass = filter_throughput('f661')
        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)

        bandpass = filter_throughput('F1265')
        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)

        bandpass = filter_throughput('deFault')
        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)

        bandpass = filter_throughput('none')
        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)

    def test_which_focalpalne(self):
        self.assertEqual(which_focalplane('f565'), 'vis')
        self.assertEqual(which_focalplane('F661'), 'vis')
        self.assertEqual(which_focalplane('f743'), 'vis')
        self.assertEqual(which_focalplane('f883'), 'vis')
        self.assertEqual(which_focalplane('F940'), 'nir')
        self.assertEqual(which_focalplane('f1265'), 'nir')
        self.assertEqual(which_focalplane('F1425'), 'nir')
        self.assertEqual(which_focalplane('f1542'), 'nir')
        self.assertEqual(which_focalplane('wfs'), 'wfs')
        self.assertRaises(ValueError, which_focalplane, 'what')

    def test_make_focus_image(self):
        # test fuction to generate psf
        def gaussian_psf(band, spectrum, shape, error=0.1):
            psf_shape = [shape, shape]

            xx, yy = np.mgrid[0:psf_shape[0], 0:psf_shape[1]]
            center = np.array([(psf_shape[0]-1)/2, (psf_shape[1]-1)/2])
            sigma = 10
            psf = np.exp(-((xx-center[0])**2 +
                         (yy-center[1])**2) / (2*sigma**2))
            psf = psf / psf.sum()

            filter = filter_throughput(band)
            return psf * (spectrum * filter).integrate()

        # test targets
        cstar = star_photlam(0, 'F2V', is_blackbody=True)
        targets = [
            [0, 0, cstar],
            [1, 1, star_photlam(10-5, 'B2V', is_blackbody=True)],
            [-2, 2, star_photlam(11-5, 'A2V', is_blackbody=True)],
            [3, -3, star_photlam(12-5, 'G2V', is_blackbody=True)],
            [100, 100, star_photlam(12, 'K2V', is_blackbody=True)],
            [100, 100, star_photlam(12, 'K2V', is_blackbody=True)],
        ]

        focus_image = make_focus_image(
            'f661',
            targets,
            gaussian_psf,
            init_shifts=[1, 1],
            rotation=45,
            platesize=[1024, 1024]
        )
        self.assertIsNotNone(focus_image)

        focus_image = make_focus_image(
            'f661',
            [],
            gaussian_psf,
            init_shifts=[1, 1],
            rotation=45,
            platesize=[1024, 1024]
        )

        self.assertIsNotNone(focus_image)

    def test_focal_mask(self):
        image = np.zeros((100, 100)) + 1
        image_out = focal_mask(image, 1, 0.1, throughtput=0)
        self.assertEqual((image - image_out).sum(), 2000+2000-400)


if __name__ == '__main__':
    unittest.main()