import unittest
from csst_cpic_sim.camera import CpicVisEmccd
import numpy as np


class TestEMCCD(unittest.TestCase):
    def test_emccd_normal(self):
        emccd = CpicVisEmccd()
        image_focal = np.zeros(emccd.flat_shape) + 10
        iamge_cosmic_ray = np.zeros(emccd.dark_shape)
        em_set = 200
        expt = 0
        image = emccd.readout(image_focal, em_set, expt, iamge_cosmic_ray)

        self.assertEqual(image.shape[0], emccd.bias_shape[0])
        self.assertEqual(image.shape[1], emccd.bias_shape[1])

    def test_emccd_allclose(self):
        config = {
            'switch': {'shutter': False},
            'dark': 1e-2,
            'cic': 0.2,
            'flat': np.zeros((1024, 1024)),
            'ccd_temp': -200
        }

        emccd = CpicVisEmccd(config)
        for key in emccd.switch:
            emccd.switch[key] = False

        image_focal = np.zeros(emccd.flat_shape) + 10
        em_set = 1024
        expt = 10
        image = emccd.readout(image_focal, em_set, expt)

        self.assertEqual(image.shape[0], emccd.bias_shape[0])
        self.assertEqual(image.shape[1], emccd.bias_shape[1])

    def test_emccd_blooming(self):
        import numpy as np
        emccd = CpicVisEmccd()
        image_focal = np.zeros(emccd.flat_shape) + 1000
        image_focal[100:200, 100:200] = 100_000
        iamge_cosmic_ray = np.zeros(emccd.dark_shape)
        emgain = 800
        expt = 1
        emccd.ccd_temp = -100
        emccd.time_syn(10, readout=False)
        image = emccd.readout(image_focal, None, expt, iamge_cosmic_ray, emgain=emgain)
        self.assertEqual(image.shape[0], emccd.bias_shape[0])
        self.assertEqual(image.shape[1], emccd.bias_shape[1])
    def test_em_fix_fun(self):
        emccd = CpicVisEmccd()
        emgain = emccd.em_fix_fuc_fit(-5)
        self.assertEqual(emgain, 0.0)

    def test_emccd_update(self):
        emccd = CpicVisEmccd()
        emccd.ccd_temp = -100
        emgain = emccd.emgain_set(1024, ccd_temp=None, self_update=False )
        self.assertAlmostEqual(emgain, 1.23, places=2)


# if __name__ == '__main__':
#     unittest.main()
    # from CpicImgSim.camera import CPIC_VIS_EMCCD
    # emccd = CPIC_VIS_EMCCD()
    # emccd.emgain_fun(1023, -30)

    # bias_images = []
    # for _ in range(10):
    #     bias = emccd.bias_frame(show=True)
    #     bias_images.append(bias)
    # bias_images = np.array(bias_images)
    # from astropy.io import fits
    # fits.writeto("bias.fits", bias_images)