import unittest
import numpy as np

from utility import simulations
from utility import noise


class CountBitErrorsTestCase(unittest.TestCase):
    """Test case for bit error counting."""
    def test_count_bit_errors(self):
        d1 = np.array([0, 0, 0, 0])
        y_hat1 = np.array([0, 1, 0, 1])

        d2 = np.array([0, 0, 0, 0])
        y_hat2 = np.array([0, 0, 0, 0])

        d3 = np.array([0, 0, 0, 0])
        y_hat3 = np.array([1, 1, 1, 1])

        self.assertEqual(simulations.count_bit_errors(d1, y_hat1), 2)
        self.assertEqual(simulations.count_bit_errors(d2, y_hat2), 0)
        self.assertEqual(simulations.count_bit_errors(d3, y_hat3), 4)


# TODO: Is this correct?
class NoiseAmpFromSNRTestCase(unittest.TestCase):
    """Test case for noise amplitude calculation."""
    def test_get_noise_amp_from_SNR(self):
        SNR1 = 0
        SNR2 = 6
        SNR3 = 20
        SNR4 = -20
        SNR5 = 60

        self.assertEqual(noise.get_noise_amp_from_SNR(SNR1, signal_amp=1), 1)
        self.assertAlmostEqual(noise.get_noise_amp_from_SNR(SNR2, signal_amp=1), 0.5, places=2)
        self.assertEqual(noise.get_noise_amp_from_SNR(SNR3, signal_amp=1), 0.1)
        self.assertEqual(noise.get_noise_amp_from_SNR(SNR4, signal_amp=1), 10)
        self.assertEqual(noise.get_noise_amp_from_SNR(SNR5, signal_amp=2), 0.002)


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