import unittest
from unittest import mock
from csst_cpic_sim.io import obsid_parser, primary_hdu, frame_header, save_fits_simple
import csst_cpic_sim.io as io
from astropy.io import fits
import numpy as np
import yaml

cstar = {
    'magnitude': 5,
    'ra': '0d',
    'dec': '0d',
    'sptype': 'G0III',
    'distance': 10,
    'mag_input_band': 'f661'
}

params = {
    'target': {'cstar': cstar},
    'skybg': 21,
    'expt': 20,
    'nframe': 5,
    'band': 'f661',
    'emgain': 10,
    'obsid': '51012345678',
    'rotation': 20,
    'shift': [0, 1],
}


class TestIO(unittest.TestCase):
    def test_obsid_parser(self):
        self.assertRaises(ValueError, obsid_parser, '20190101')
        self.assertRaises(ValueError, obsid_parser, '123456789012')

        self.assertEqual(obsid_parser('50012345678'), 'BIAS')
        self.assertEqual(obsid_parser('50112345678'), 'DARK')
        self.assertEqual(obsid_parser('50212345678'), 'FLAT')
        self.assertEqual(obsid_parser('50312345678'), 'BKGD')
        self.assertEqual(obsid_parser('51012345678'), 'SCIE')

    def test_primary_hdu(self):
        hdu1 = primary_hdu(params, {}, filename_output=False)
        hdu2, folder, filename = primary_hdu(params, {}, filename_output=True)
        self.assertIsInstance(hdu1, fits.PrimaryHDU)
        self.assertIsInstance(hdu2, fits.PrimaryHDU)
        self.assertIsInstance(folder, str)
        self.assertIsInstance(filename, str)

    def test_frame_header(self):
        header = frame_header(params, 1, '2021-01-01T00:00:00')
        self.assertEqual(header['IMGINDEX'], 1)
        self.assertIsInstance(header, fits.Header)

    def test_write_fits(self):
        images = np.zeros((5, 10, 10))
        yaml_str = """
obsid: 51012345678
expt: 300
nframe: 10
band: "f661"
shift: [0, 0]
rotation: 0
emgain: 100
skybg: 21

target:
    cstar:
        ra: "10.684792d"
        dec: "41.26917d"
        sptype: "M0.5"
        magnitude: 3.4
    planets:
        - ra: 10.684792
          dec: 41.26917
          sptype: "M0.5"
          magnitude: 3.4
"""
        print(io.tmp_folder_path)
        parameters = yaml.load(yaml_str, Loader=yaml.FullLoader)
        mock_fits_writeto = mock.MagicMock()
        io.fits.writeto = mock_fits_writeto
        io.tmp_folder_path = 'test_folder_for_unit_test'
        io.save_fits_simple(images, parameters)
        output_name = mock_fits_writeto.call_args[0][0]
        self.assertEqual(
            output_name[:len(io.tmp_folder_path)], io.tmp_folder_path)


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