import unittest
import numpy as np
from decoders import maximum_likelihood


class CodewordGenerationTestCase(unittest.TestCase):
    def test_codeword_generation(self):
        """Test case for data word and code word generation."""
        # Hamming(7,4) code
        G = np.array([[1, 1, 1, 0, 0, 0, 0],
                      [1, 0, 0, 1, 1, 0, 0],
                      [0, 1, 0, 1, 0, 1, 0],
                      [1, 1, 0, 1, 0, 0, 1]])
        H = np.array([[1, 0, 1, 0, 1, 0, 1],
                      [0, 1, 1, 0, 0, 1, 1],
                      [0, 0, 0, 1, 1, 1, 1]])

        decoder = maximum_likelihood.MLDecoder(G, H)

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

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

        self.assertEqual(np.array_equal(decoder._datawords, expected_datawords), True)
        self.assertEqual(np.array_equal(decoder._codewords, expected_codewords), True)