import numpy as np
import typing
import scipy


def _sign(val):
    return -1 * (val < 0) + 1 * (val >= 0)


class PathTracker:
    """
    Path trakcer for the homotopy continuation method. Uses a
    predictor-corrector scheme to trace a path defined by a homotopy.

    References:
      [1] T. Chen and T.-Y. Li, “Homotopy continuation method for solving
          systems of nonlinear and polynomial equations,” Communications in
          Information and Systems, vol. 15, no. 2, pp. 119-307, 2015
    """

    def __init__(self, homotopy, euler_step_size=0.05, euler_max_tries=10, newton_max_iter=5,
                 newton_convergence_threshold=0.001, sigma=1):
        self._homotopy = homotopy
        self._euler_step_size = euler_step_size
        self._euler_max_tries = euler_max_tries
        self._newton_max_iter = newton_max_iter
        self._newton_convergence_threshold = newton_convergence_threshold
        self._sigma = sigma

    def step(self, y):
        """Perform one predictor-corrector step."""
        return self.transparent_step(y)[3]

    def transparent_step(self, y) -> typing.Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray,]:
        """Perform one predictor-corrector step, returning intermediate results."""
        for i in range(self._euler_max_tries):
            step_size = self._euler_step_size / (1 << i)

            y_hat, y_prime = self._perform_euler_predictor_step(y, step_size)
            y_hat_n = self._perform_newtown_corrector_step(y_hat)

            return y, y_prime, y_hat, y_hat_n

        raise RuntimeError("Newton corrector did not converge")

    # TODO: Make sure the implementation with null_space instead of QR is correct
    def _perform_euler_predictor_step(self, y, step_size) -> typing.Tuple[np.ndarray, np.ndarray]:
        # Obtain y_prime

        DH = self._homotopy.DH(y)
        ns = scipy.linalg.null_space(DH)

        y_prime = ns[:, 0] * self._sigma

        # Q, R = np.linalg.qr(np.transpose(DH), mode="complete")
        # y_prime = Q[:, 2]
        #
        # if _sign(np.linalg.det(Q)*np.linalg.det(R[:2, :])) != _sign(self._sigma):
        #     y_prime = -y_prime

        # Perform prediction

        y_hat = y + step_size*y_prime

        return y_hat, y_prime

    def _perform_newtown_corrector_step(self, y) -> np.ndarray:
        prev_y = y

        for _ in range(self._newton_max_iter):
            # Perform correction

            DH = self._homotopy.DH(y)
            DH_pinv = np.linalg.pinv(DH)

            y = y - np.transpose(DH_pinv @ self._homotopy.H(y))[0]

            # Check stopping criterion

            if np.linalg.norm(y - prev_y) < self._newton_convergence_threshold:
                return y

            prev_y = y

        raise RuntimeError("Newton corrector did not converge")

    def set_sigma(self, sigma):
        self._sigma = sigma