import numpy as np
import itertools


class MLDecoder:
    """This class naively implements a soft decision decoder. The decoder calculates
    the correlation between the received signal and each codeword and then chooses the
    one with the largest correlation.
    """

    def __init__(self, G: np.array, H: np.array):
        """Construct a new SoftDecisionDecoder object.

        :param G: Generator matrix
        :param H: Parity check matrix
        """
        self._G = G
        self._H = H
        self._datawords, self._codewords = self._gen_codewords()
        self._codewords_bpsk = 1 - 2 * self._codewords  # The codewords, but mapped to [-1, 1]^n

    def _gen_codewords(self) -> np.array:
        """Generate a list of all possible codewords.

        :return: Numpy array of the form [[codeword_1], [codeword_2], ...]
            (Each generated codeword is an element of [0, 1]^n)
        """
        k, n = self._G.shape

        # Generate a list of all possible data words
        u_lst = [list(i) for i in itertools.product([0, 1], repeat=k)]
        u_lst = np.array(u_lst)

        # Map each data word onto a codeword
        c_lst = np.dot(u_lst, self._G) % 2

        return u_lst, c_lst

    def decode(self, y: np.array) -> np.array:
        """Decode a received signal.

        This function assumes a BPSK modulated signal.

        :param y: Vector of received values. (y = x + w, where 'x' is element of [-1, 1]^n
            and 'w' is noise)
        :return: Most probably sent codeword (element of [0, 1]^k)
        """
        correlations = np.dot(self._codewords_bpsk, y)

        return self._codewords[np.argmax(correlations)]