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

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

        # Generate G, F, and H
        self.G = self._create_G()
        self.F = self._create_F()
        self.H = self._create_H()

        # Convert to callable functions
        self._H_lambda = self._create_H_lambda()
        self._DH_lambda = self._create_DH_lambda()

    def _create_G(self) -> List[sp.Expr]:
        """Create G polynomial system (the starting system)."""
        # For each variable xi, add the polynomial [xi]
        G = []
        for var in self.x_vars:
            G.append(var)

        return G

    def _create_F(self) -> List[sp.Expr]:
        """Create F polynomial system (the target system)."""
        F = []

        # Add 1 - xi^2 for each variable
        for var in self.x_vars:
            F.append(1 - var**2)

        # Add parity check polynomials: 1 - x1*x2*...*xk for each parity check
        for row in self.H_matrix:
            # Create product of variables that participate in this check
            term = 1
            for i, bit in enumerate(row):
                if bit == 1:
                    term *= self.x_vars[i]

            if term != 1:  # Only add if there are variables in this check
                F.append(1 - term)

        return F

    def _create_H(self) -> List[sp.Expr]:
        """Create the homotopy H = (1-t)*G + t*F."""
        H = []

        # Make sure G and F have the same length
        # Repeat variables from G if needed to match F's length
        G_extended = self.G.copy()
        while len(G_extended) < len(self.F):
            # Cycle through variables to repeat
            for i in range(min(self.num_vars, len(self.F) - len(G_extended))):
                G_extended.append(self.x_vars[i % self.num_vars])

        # Create the homotopy
        for g, f in zip(G_extended, self.F):
            H.append((1 - self.t) * g + self.t * f)

        return H

    def _create_H_lambda(self) -> Callable:
        """Create a lambda function to evaluate H."""
        all_vars = self.x_vars + [self.t]
        return sp.lambdify(all_vars, self.H, 'numpy')

    def _create_DH_lambda(self) -> Callable:
        """Create a lambda function to evaluate the Jacobian of H."""
        all_vars = self.x_vars + [self.t]
        jacobian = sp.Matrix([[sp.diff(expr, var)
                             for var in all_vars] for expr in self.H])
        return sp.lambdify(all_vars, jacobian, '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)