import unittest
import numpy as np

from utility import simulations, noise, codes


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: Rewrite tests for new SNR calculation
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)


class CodesTestCase(unittest.TestCase):
    """Tests relating to the 'codes' utilities."""

    def test_get_systematic_H(self):
        # Hamming(7,4) code
        G = np.array([[1, 0, 0, 0, 0, 1, 1],
                      [0, 1, 0, 0, 1, 0, 1],
                      [0, 0, 1, 0, 1, 1, 0],
                      [0, 0, 0, 1, 1, 1, 1]])

        expected_H = np.array([[0, 1, 1, 1, 1, 0, 0],
                               [1, 0, 1, 1, 0, 1, 0],
                               [1, 1, 0, 1, 0, 0, 1]])

        H = codes.get_systematic_H(G)

        self.assertEqual(np.array_equal(expected_H, H), True)


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