from bilby.gw import calibration
import unittest
import numpy as np


class TestBaseClass(unittest.TestCase):

    def setUp(self):
        self.model = calibration.Recalibrate()

    def tearDown(self):
        del self.model

    def test_repr(self):
        expected = 'Recalibrate(prefix={})'.format('\'recalib_\'')
        actual = repr(self.model)
        self.assertEqual(expected, actual)

    def test_calibration_factor(self):
        frequency_array = np.linspace(20, 1024, 1000)
        cal_factor = self.model.get_calibration_factor(frequency_array)
        assert np.alltrue(cal_factor.real == np.ones_like(frequency_array))


class TestCubicSpline(unittest.TestCase):

    def setUp(self):
        self.prefix = 'recalib_'
        self.minimum_frequency = 20
        self.maximum_frequency = 1024
        self.n_points = 5
        self.model = calibration.CubicSpline(
            prefix=self.prefix, minimum_frequency=self.minimum_frequency,
            maximum_frequency=self.maximum_frequency, n_points=self.n_points)
        self.parameters = {'recalib_{}_{}'.format(param, ii): 0.0
                           for ii in range(5)
                           for param in ['amplitude', 'phase']}

    def tearDown(self):
        del self.prefix
        del self.minimum_frequency
        del self.maximum_frequency
        del self.n_points
        del self.model
        del self.parameters

    def test_calibration_factor(self):
        frequency_array = np.linspace(20, 1024, 1000)
        cal_factor = self.model.get_calibration_factor(frequency_array,
                                                       **self.parameters)
        assert np.alltrue(cal_factor.real == np.ones_like(frequency_array))

    def test_repr(self):
        expected = 'CubicSpline(prefix=\'{}\', minimum_frequency={}, maximum_frequency={}, n_points={})'\
            .format(self.prefix, self.minimum_frequency, self.maximum_frequency, self.n_points)
        actual = repr(self.model)
        self.assertEqual(expected, actual)


class TestCubicSplineRequiresFourNodes(unittest.TestCase):

    def test_cannot_instantiate_with_too_few_nodes(self):
        for ii in range(6):
            if ii < 4:
                with self.assertRaises(ValueError):
                    calibration.CubicSpline('test', 1, 10, ii)
            else:
                calibration.CubicSpline('test', 1, 10, ii)


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