import numpy as np
from tqdm import tqdm


class ProximalDecoder:
    """Class implementing the Proximal Decoding algorithm. See "Proximal Decoding for LDPC Codes" by Tadashi
    Wadayama, and Satoshi Takabe.
    """
    def __init__(self, H: np.array, K: int = 10, step_size: float = 0.01, gamma: float = 0.05):
        """Construct a new ProximalDecoder Object.

        :param H: Parity Check Matrix
        :param K: Max number of iterations to perform when decoding
        :param step_size: Step size for the gradient descent process
        :param gamma: Positive constant. Arises in the approximation of the prior PDF
        """
        self._H = H
        self._K = K
        self._step_size = step_size
        self._gamma = gamma

    @staticmethod
    def _L_awgn(s: np.array, y: np.array) -> np.array:
        """Variation of the negative log-likelihood for the special case of AWGN noise. See 4.1, p. 4."""
        return s - y

    def _grad_h(self, x: np.array) -> np.array:
        """Gradient of the code-constraint polynomial. See 2.3, p. 2."""
        # Calculate first term

        result = 4 * (x**2 - 1) * x

        # Calculate second term

        for k, x_k in enumerate(x):
            # TODO: Perform this operation for each row simultaneously
            B_k = np.argwhere(self._H[:, k] == 1)
            B_k = B_k[:, 0]  # Get rid of one layer of arrays

            # TODO: Perform the summation with np.sum()
            sum_result = 0
            for i in B_k:
                # TODO: Perform this operation for each column simultaneously
                A_i = np.argwhere(self._H[i] == 1)
                A_i = A_i[:, 0]  # Get rid of one layer of arrays

                prod = 1
                for j in A_i:
                    prod *= x[j]

                sum_result += prod**2 - prod

            term_2 = 2 / x_k * sum_result

            result[k] += term_2

        return np.array(result)

    def _check_parity(self, y_hat: np.array) -> bool:
        """Perform a parity check for a given codeword.

        :param y_hat: codeword to be checked
        :return: True if the parity check passes, i.e. the codeword is valid. False otherwise
        """
        syndrome = np.dot(self._H, y_hat) % 2
        return not np.any(syndrome)

    def decode(self, y: np.array) -> np.array:
        """Decode a received signal. The algorithm is detailed in 3.2, p.3.

        This function assumes an AWGN channel.

        :param y: Vector of received values
        :return: Most probably sent symbol
        """
        s = 0
        x_hat = 0
        for k in range(self._K):
            r = s - self._step_size * self._L_awgn(s, y)
            s = r - self._gamma * self._grad_h(r)
            x_hat = (np.sign(s) == 1) * 1

            if self._check_parity(x_hat):
                break

        return x_hat