import unittest
from CpicImgSim.camera import CosmicRayFrameMaker, sky_frame_maker
from CpicImgSim.target import star_photlam
from CpicImgSim.optics import filter_throughput
import numpy as np


class TestCRandBKG(unittest.TestCase):
    def test_cr(self):
        image = np.zeros((1000, 1000))
        cr = CosmicRayFrameMaker()
        crimage = cr.make_cr_frame(image.shape, 30)
        self.assertEqual(crimage.shape[0], image.shape[0])
        self.assertEqual(crimage.shape[1], image.shape[1])

    def test_no_cr(self):
        cr = CosmicRayFrameMaker()
        crimage = cr.make_cr_frame((1000, 1000), 0)
        self.assertEqual(crimage.sum(), 0)

    def test_cr_rate(self):
        cr = CosmicRayFrameMaker()
        width_pix = 1000
        expt = 30
        cr_count = cr._number_rand(expt, (width_pix, width_pix))
        width_um = width_pix * cr.pitch
        width_cm = width_um / 1e4
        area_cm2 = width_cm**2
        cr_count_cal = int(cr.cr_rate * expt * area_cm2)
        self.assertEqual(cr_count, cr_count_cal)

        cr_count_random_1 = cr._number_rand(
            expt, (width_pix, width_pix), random=True, seed=10)
        cr_count_random_2 = cr._number_rand(
            expt, (width_pix, width_pix), random=True, seed=10)
        self.assertEqual(cr_count_random_1, cr_count_random_2)

    def test_back_ground(self):
        for band in ['f661', 'f743', 'f1542']:
            skframe = sky_frame_maker(band, 20, 1, (100, 100))
            self.assertEqual(skframe.shape[0], 100)
            self.assertEqual(skframe.shape[1], 100)
            self.assertAlmostEqual(skframe.std(), 0)

            star = star_photlam(20, 'F2V', mag_input_band=band)
            filter = filter_throughput(band)

            count = (star * filter).integrate()
            self.assertAlmostEqual(skframe[1, 1], count)


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