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.G = self._create_G()
        self.F = self._create_F()
        self.H = self._create_H()
        self.DH = self._create_DH(self.H)

        self._H_lambda = self._create_H_lambda()
        self._DH_lambda = self._create_DH_lambda()

    def _create_G(self) -> List[sp.Expr]:
        G = []
        for var in self.x_vars:
            G.append(var)

        return G

    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_H(self) -> List[sp.Expr]:
        H = []

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

        return H

    def _create_DH(self, H: List[sp.Expr]) -> 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_H_lambda(self) -> Callable:
        all_vars = self.x_vars + [self.t]
        return sp.lambdify(all_vars, self.H, 'numpy')

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

    def evaluate_H(self, y: np.ndarray) -> np.ndarray:
        """
        Evaluate H at point y.

        Args:
            y: Array of form [x1, x2, ..., xn, t] where xi are the variables
               and t is the homotopy parameter.

        Returns:
            Array containing H evaluated at y.
        """
        return np.array(self._H_lambda(*y))

    def evaluate_DH(self, y: np.ndarray) -> np.ndarray:
        """
        Evaluate the Jacobian of H at point y.

        Args:
            y: Array of form [x1, x2, ..., xn, t] where xi are the variables
               and t is the homotopy parameter.

        Returns:
            Matrix containing the Jacobian of H evaluated at y.
        """
        return np.array(self._DH_lambda(*y), dtype=float)