import unittest
from csst_cpic_sim.target import _sptype2num, hybrid_albedo_spectrum, star_photlam, spectrum_generator, planet_contrast, extract_target_x_y
from csst_cpic_sim.config import S
from math import pi


class TestTarget(unittest.TestCase):
    def test_hybrid_albedo_spectrum(self):
        planet = hybrid_albedo_spectrum(0.5, 1)
        self.assertIsInstance(planet, S.spectrum.SpectralElement)

    def test_sptype(self):
        self.assertEqual(_sptype2num('M1V'), (6.1, 5))
        self.assertEqual(_sptype2num('O5IV'), (0.5, 4))
        self.assertEqual(_sptype2num('F3V'), (3.3, 5))
        self.assertEqual(_sptype2num('K4.5II'), (5.45, 2))
        self.assertEqual(_sptype2num('M1'), (6.1, 5))
        self.assertEqual(_sptype2num('M'), (6, 5))
        self.assertEqual(_sptype2num('K3Vvar'), (5.3, 5))
        self.assertEqual(_sptype2num('F6Vbwvar'), (3.6, 5))
        self.assertEqual(_sptype2num('K0IV SB'), (5.0, 4))
        self.assertEqual(_sptype2num('F5V+'), (3.5, 5))
        self.assertEqual(_sptype2num('G5IV/V +K1IV'), (4.5, 4))
        self.assertEqual(_sptype2num('F7IV-V'), (3.7, 4))
        self.assertEqual(_sptype2num('O4/O5IV'), (0.4, 5))
        self.assertEqual(_sptype2num('G3/G5V'), (4.3, 5))
        self.assertRaises(ValueError, _sptype2num, 'errorinput')

    def test_star_photlam(self):
        photlam = star_photlam(5, 'F0III', is_blackbody=True)
        self.assertIsInstance(photlam, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F2II', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        photlam2 = star_photlam(5, 'F2', is_blackbody=False)
        self.assertIsInstance(photlam2, S.spectrum.SourceSpectrum)
        self.assertRaises(ValueError, star_photlam, 5, 'T0III', is_blackbody=1)

    def test_planet_contrast(self):
        contrast = planet_contrast(1, 0, 90, 1)
        contrast_cal = (69_911 / 149_597_870.7)**2 / pi
        self.assertAlmostEqual(contrast, contrast_cal)
        self.assertRaises(ValueError, planet_contrast, 1, 0, 0, 1)

    def test_extract_target_x_y(self):
        target = {
            'ra': '120d',
            'dec': '40d',
            'distance': 10,
        }
        x, y = extract_target_x_y(target, '120d', '40d')
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 0)

        x, y = extract_target_x_y(target, '120d', '41d')
        self.assertAlmostEqual(y, -3600)

        x, y = extract_target_x_y(target, '8h', '40d')
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 0)

        x, y = extract_target_x_y(dict(pangle=0, separation=1))
        self.assertAlmostEqual(x, 0)
        self.assertAlmostEqual(y, 1)

        x, y = extract_target_x_y(dict(pangle=90, separation=1))
        self.assertAlmostEqual(x, 1)
        self.assertAlmostEqual(y, 0)

        self.assertRaises(ValueError, extract_target_x_y, dict(pangle=0))
        self.assertRaises(ValueError, extract_target_x_y, dict(ra='120d'))
        self.assertRaises(ValueError, extract_target_x_y,
                          dict(ra='120d', dec='40d'))

    def test_spectrum_generator(self):
        target_example = {
            'cstar': {
                'magnitude': 5,
                'ra': '120d',
                'dec': '40d',
                'distance': 10,
                'sptype': 'F0III',
            },
            'stars': [
                {
                    'magnitude': 5,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0III',
                    'isblackbody': True
                },
                {
                    'magnitude': 2,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0II',
                    'isblackbody': True
                },
                {
                    'magnitude': 2,
                    'ra': '120.0001d',
                    'dec': '40.0001d',
                    'sptype': 'G0',
                    'isblackbody': True
                },
            ],
            'planets': [
                {
                    'radius': 2,
                    'pangle': 60,
                    'separation': 0.5,
                    'phase_angle': 90,
                }
            ]
        }

        sps = spectrum_generator(target_example)
        num_ouput = len([sp for sp in sps if sp is not None])
        num_input = len(target_example['stars']) + \
            len(target_example['planets']) + 1
        self.assertEqual(num_ouput, num_input)


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