import numpy
import numpy as np
import itertools


class SoftDecisionDecoder:
    """This class naively implements a soft decision decoder. This decoder calculates
    the posterior probability for each codeword and then chooses the one with the largest
    probability.
    """

    def __init__(self, G: np.array, H: np.array):
        """Construct a new SotDecisionDecoder 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 = self._codewords * 2 - 1  # 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], ...]
        """
        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-like modulated signal ([-1, 1]^n instead of [0, 1]^n)
        and an AWGN channel.

        :param y: Vector of received values. (y = x + n, where 'x' is element of [-1, 1]^m
            and 'n' is noise)
        :return: Most probably sent symbol
        """
        # TODO: Is there a nice numpy way to implement this for loop?
        correlations = []
        for c in self._codewords_bpsk:
            correlations.append(np.dot(y, c))

        return self._datawords[numpy.argmax(correlations)]