import unittest

import time

from csst_cpic_sim.target import star_photlam
from csst_cpic_sim.optics import ideal_focus_image, focal_convolve, focal_mask, filter_throughput, ideal_focus_image
from csst_cpic_sim.config import which_focalplane, S
import numpy as np
from astropy.io import fits
from csst_cpic_sim.optics import FILTERS

def make_test_sub_image(size, shape):
    shape = np.array([shape, shape])
    sub_image = np.zeros(shape)
    center = (shape-1)/2

    xx, yy = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]))
    xx = xx - center[1]
    yy = yy - center[0]
    sub_image[np.abs(xx) < 0.6] = 1
    sub_image[np.abs(yy) < 0.6] = 1

    sub_image[(np.abs(xx) < 0.6) & (np.abs(yy) < 0.6)] = 0
    sub_image[yy > 7] = 0
    sub_image[yy < -3] = 0
    sub_image[np.abs(xx) > 3] = 0

    sub_image2 = (np.sqrt(xx**2*2 + yy**2) < size).astype(int)
    sub_image = sub_image2 * (1 - sub_image)
    return sub_image
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()



class TestOptics(unittest.TestCase):
    def test_filter_throughtput(self):

        bands = list(FILTERS.keys())
        bandpass = filter_throughput(bands[0])
        self.assertIsInstance(bandpass, S.spectrum.SpectralElement)

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


    def test_which_focalpalne(self):
        self.assertEqual(which_focalplane('f565'), 'vis')

    def test_ideal_focus_image(self):
        targets = [
            [-20, 0, star_photlam(2, 'G2'), None],
            [5, 3, star_photlam(0, 'G2'), make_test_sub_image(4, 20)],
            [8, 0, star_photlam(-5, 'G2'), make_test_sub_image(10, 100)],
        ]
        bandpass = S.Box(6000, 500)
        
        foc = ideal_focus_image(
            bandpass,
            targets,
            0.0165,
            [1024, 1024],
        )
        # fits.writeto('foc.fits', foc, overwrite=True)

        foc = ideal_focus_image(
            bandpass,
            targets,
            0.0165,
            [1024, 1024],
            rotation=30,
        )

        foc = ideal_focus_image(
            bandpass,
            {},
            0.0165,
            [1024, 1024],
            rotation=30,
        )


        # fits.writeto('foc_rot30.fits', foc, overwrite=True)
    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)

    def test_convolve_psf(self):
        targets = [
            [0, 0, star_photlam(2, 'G2'), None],
            [5, 3, star_photlam(0, 'G2'), make_test_sub_image(4, 20)],
            [8, 0, star_photlam(-5, 'G2'), make_test_sub_image(10, 100)],
        ]

        img_final = focal_convolve('f661', {})

        img_final = focal_convolve('f661', targets)

        img_final = focal_convolve('f661', targets, init_shifts=[10, 10])

        # fits.writeto('cov.fits', img_final, overwrite=True)


# if __name__ == '__main__':
#     # # unittest.main()
#     # import time
#     # from CpicImgSim.target import star_photlam




    




    # test_convolve_psf()

    # # import matplotlib.pyplot as plt
    # # plt.imshow(make_test_sub_image(5, 6))
    # # plt.show()