import unittest
import numpy as np

from utility import simulation, 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(simulation.count_bit_errors(d1, y_hat1), 2)
        self.assertEqual(simulation.count_bit_errors(d2, y_hat2), 0)
        self.assertEqual(simulation.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 = 3
        SNR3 = 20
        SNR4 = -20

        var1 = noise.get_noise_variance_from_SNR(SNR1, n=8, k=8)
        var2 = noise.get_noise_variance_from_SNR(SNR2, n=8, k=8)
        var3 = noise.get_noise_variance_from_SNR(SNR3, n=8, k=8)
        var4 = noise.get_noise_variance_from_SNR(SNR4, n=8, k=8)

        self.assertEqual(var1, 1 * 0.5)
        self.assertAlmostEqual(var2, 0.5 * 0.5, places=2)
        self.assertEqual(var3, 0.01 * 0.5)
        self.assertEqual(var4, 100 * 0.5)


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()