import numpy as np
import sympy as sp
from typing import List, Callable


class HomotopyGenerator:
    """Generates homotopy functions from a binary parity check matrix."""

    def __init__(self, parity_check_matrix: np.ndarray):
        """
        Initialize with a parity check matrix.

        Args:
            parity_check_matrix: Binary matrix where rows represent parity checks
                                 and columns represent variables.
        """
        self.H_matrix = parity_check_matrix
        self.num_checks, self.num_vars = parity_check_matrix.shape

        self.x_vars = [sp.symbols(f'x{i+1}') for i in range(self.num_vars)]
        self.t = sp.symbols('t')

        self._F = self._create_F()
        self._F_lambda = self._create_F_lambda(self._F)
        # self._G = self._create_G(received)
        # self._G_lambda = self._create_G_lambda(self._G)
        # self._H = self._create_H()
        # self._H_lambda = self._create_H_lambda(self._H)
        # self._DH = self._create_DH()
        # self._DH_lambda = self._create_DH_lambda(self._DH)

    def update_received(self, received: np.ndarray):
        """Update the received vector and recompute G."""
        self._G = self._create_G(received)
        self._G_lambda = self._create_G_lambda(self._G)
        self._H = self._create_H()
        self._H_lambda = self._create_H_lambda(self._H)
        self._DH = self._create_DH()
        self._DH_lambda = self._create_DH_lambda(self._DH)

    def _create_F(self) -> sp.MutableMatrix:
        F = []

        for var in self.x_vars:
            F.append(1 - var**2)

        for row in self.H_matrix:
            term = 1
            for i, bit in enumerate(row):
                if bit == 1:
                    term *= self.x_vars[i]

            F.append(1 - term)

        groebner_basis = sp.groebner(F)
        return sp.MutableMatrix(groebner_basis)

    def _create_G(self, received) -> sp.MutableMatrix:
        G = []
        F_y = self._F_lambda(*received)

        for f, f_y_i in zip(self._F, F_y):
            G.append(f - f_y_i)

        return sp.MutableMatrix(G)

    def _create_H(self) -> sp.MutableMatrix:
        H = []

        for f, g in zip(self._F, self._G):
            H.append((1 - self.t) * g + self.t * f)

        return sp.MutableMatrix(H)

    def _create_DH(self) -> sp.MutableMatrix:
        all_vars = self.x_vars + [self.t]
        DH = sp.Matrix([[sp.diff(expr, var)
                         for var in all_vars] for expr in self._H])

        return DH

    def _create_F_lambda(self, expr) -> Callable:
        all_vars = self.x_vars
        return sp.lambdify(all_vars, expr, 'numpy')

    def _create_G_lambda(self, expr) -> Callable:
        all_vars = self.x_vars
        return sp.lambdify(all_vars, expr, 'numpy')

    def _create_H_lambda(self, expr) -> Callable:
        all_vars = self.x_vars + [self.t]
        return sp.lambdify(all_vars, expr, 'numpy')

    def _create_DH_lambda(self, expr) -> Callable:
        all_vars = self.x_vars + [self.t]
        return sp.lambdify(all_vars, expr, 'numpy')

    def H(self, y):
        return np.array(self._H_lambda(*y))

    def DH(self, y):
        return np.array(self._DH_lambda(*y))


def main():
    import utility
    H = utility.read_alist_file(
        "/home/andreas/workspace/work/hiwi/ba-sw/codes/BCH_7_4.alist")
    a = HomotopyGenerator(H)

    y = np.array([0.1, 0.2, 0.9, 0.1, -0.8, -0.5, -1.0, 0])

    print(a.DH(y))


if __name__ == "__main__":
    main()