"""
Advanced Predictor-Corrector Methods for Fractional Differential Equations

This module provides advanced predictor-corrector methods including
Adams-Bashforth-Moulton schemes, variable step size control, and error estimation.
"""

import numpy as np
from typing import Union, Optional, Tuple, Callable

from ..core.definitions import FractionalOrder

# Use adapter system for gamma function instead of direct imports
def _get_gamma_function():
    """Get gamma function through adapter system."""
    try:
        from ..special.gamma_beta import gamma_function as gamma
        return gamma
    except Exception:
        # Fallback to scipy
        from scipy.special import gamma
        return gamma

gamma = _get_gamma_function()


class PredictorCorrectorSolver:
    """
    Advanced predictor-corrector solver for fractional differential equations.

    Implements Adams-Bashforth-Moulton type schemes with variable step size
    control and error estimation.
    """

    def __init__(
        self,
        derivative_type: str = "caputo",
        order: int = 1,
        adaptive: bool = True,
        tol: float = 1e-6,
        max_iter: int = 10,
        min_h: float = 1e-8,
        max_h: float = 1e-2,
    ):
        """
        Initialize predictor-corrector solver.

        Args:
            derivative_type: Type of fractional derivative
            order: Order of the method
            adaptive: Use adaptive step size control
            tol: Tolerance for convergence
            max_iter: Maximum number of corrector iterations
            min_h: Minimum step size
            max_h: Maximum step size
        """
        # Accept both string types and objects (e.g., FractionalOrder) gracefully
        try:
            self.derivative_type = derivative_type.lower()  # type: ignore[attr-defined]
        except Exception:
            self.derivative_type = str(derivative_type).lower()
        self.order = order
        self.adaptive = adaptive
        self.tol = tol
        self.max_iter = max_iter
        self.min_h = min_h
        self.max_h = max_h

        # Validate derivative type
        valid_derivatives = [
            "caputo", "riemann_liouville", "grunwald_letnikov"]
        if self.derivative_type not in valid_derivatives:
            raise ValueError(
                f"Derivative type must be one of {valid_derivatives}")

    def solve(
        self,
        f: Callable,
        t_span: Tuple[float, float],
        y0: Union[float, np.ndarray],
        alpha: Union[float, FractionalOrder],
        h0: Optional[float] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Solve fractional differential equation using predictor-corrector method.

        Args:
            f: Right-hand side function f(t, y)
            t_span: Time interval (t0, tf)
            y0: Initial condition(s)
            alpha: Fractional order
            h: Step size (None for adaptive)
            **kwargs: Additional solver parameters

        Returns:
            Tuple of (t_values, y_values)
        """
        t0, tf = t_span

        if h0 is None:
            h0 = (tf - t0) / 100

        if self.adaptive:
            return self._solve_adaptive(f, t0, tf, y0, alpha, h0, **kwargs)
        else:
            return self._solve_fixed_step(f, t0, tf, y0, alpha, h0, **kwargs)

    def _solve_fixed_step(
        self,
        f: Callable,
        t0: float,
        tf: float,
        y0: Union[float, np.ndarray],
        alpha: Union[float, FractionalOrder],
        h0: float,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Solve with fixed step size.

        Args:
            f: Right-hand side function
            t0: Initial time
            tf: Final time
            y0: Initial condition
            alpha: Fractional order
            h: Step size
            **kwargs: Additional parameters

        Returns:
            Tuple of (t_values, y_values)
        """
        # Convert to arrays if needed
        if np.isscalar(y0):
            y0 = np.array([y0])

        # Time grid
        t_values = np.arange(t0, tf + h0, h0)
        N = len(t_values)

        # Solution array
        y_values = np.zeros((N, len(y0)))
        y_values[0] = y0

        # Compute fractional derivative coefficients
        coeffs = self._compute_fractional_coefficients(alpha, N)

        # Main iteration loop
        for n in range(1, N):
            t_values[n]

            # Predictor step
            y_pred = self._predictor_step(
                f, t_values, y_values, n, alpha, coeffs, h0)

            # Corrector step with iteration
            y_corr = self._corrector_step(
                f, t_values, y_values, y_pred, n, alpha, coeffs, h0
            )

            # Iterative correction
            for _ in range(self.max_iter):
                y_old = y_corr.copy()
                y_corr = self._corrector_step(
                    f, t_values, y_values, y_pred, n, alpha, coeffs, h0
                )

                if np.allclose(y_corr, y_old, rtol=self.tol):
                    break

            y_values[n] = y_corr

        return t_values, y_values

    def _solve_adaptive(
        self,
        f: Callable,
        t0: float,
        tf: float,
        y0: Union[float, np.ndarray],
        alpha: Union[float, FractionalOrder],
        h0: float,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Solve with adaptive step size control.

        Args:
            f: Right-hand side function
            t0: Initial time
            tf: Final time
            y0: Initial condition
            alpha: Fractional order
            h0: Initial step size
            **kwargs: Additional parameters

        Returns:
            Tuple of (t_values, y_values)
        """
        # Convert to arrays if needed
        if np.isscalar(y0):
            y0 = np.array([y0])

        # Initialize solution arrays
        t_values = [t0]
        y_values = [y0.copy()]

        t_current = t0
        y_current = y0.copy()
        h_current = h0

        iteration_count = 0
        max_iterations = int((tf - t0) / self.min_h) + 1000  # Safety limit

        while t_current < tf and iteration_count < max_iterations:
            iteration_count += 1

            # Compute solution with current step size
            t_next = min(t_current + h_current, tf)
            h_actual = t_next - t_current

            # Predictor step
            y_pred = self._adaptive_predictor_step(
                f, t_current, t_next, y_current, alpha, h_actual
            )

            # Corrector step
            y_corr = self._adaptive_corrector_step(
                f, t_current, t_next, y_current, y_pred, alpha, h_actual
            )

            # Estimate error
            error = self._estimate_error(y_pred, y_corr)

            # Check if error is acceptable
            if error <= self.tol:
                # Accept step
                t_values.append(t_next)
                y_values.append(y_corr)
                t_current = t_next
                y_current = y_corr

                # Adjust step size for next step
                h_current = min(self.max_h, h_current *
                                (self.tol / error) ** 0.5)
            else:
                # Reject step and reduce step size
                h_current = max(self.min_h, h_current *
                                (self.tol / error) ** 0.25)

                # Safety check: if step size is at minimum and still not
                # converging, force progress
                if h_current <= self.min_h and error > self.tol:
                    # Force acceptance with warning
                    import warnings

                    warnings.warn(
                        f"Step size at minimum ({self.min_h}), forcing acceptance of step with error {error:.2e}"
                    )
                    t_values.append(t_next)
                    y_values.append(y_corr)
                    t_current = t_next
                    y_current = y_corr
                    h_current = self.min_h * 2  # Try to increase step size

        return np.array(t_values), np.array(y_values)

    def _compute_fractional_coefficients(
        self, alpha: Union[float, FractionalOrder], N: int
    ) -> np.ndarray:
        """
        Compute fractional derivative coefficients.

        Args:
            alpha: Fractional order
            N: Number of time steps

        Returns:
            Array of coefficients
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        coeffs = np.zeros(N)
        coeffs[0] = 1.0

        for j in range(1, N):
            if self.derivative_type == "caputo":
                coeffs[j] = (j + 1) ** alpha_val - j**alpha_val
            elif self.derivative_type == "grunwald_letnikov":
                coeffs[j] = coeffs[j - 1] * (1 - (alpha_val + 1) / j)
            else:  # Riemann-Liouville
                coeffs[j] = (
                    (-1) ** j
                    * gamma(alpha_val + 1)
                    / (gamma(j + 1) * gamma(alpha_val - j + 1))
                )

        return coeffs

    def _predictor_step(
        self,
        f: Callable,
        t_values: np.ndarray,
        y_values: np.ndarray,
        n: int,
        alpha: Union[float, FractionalOrder],
        coeffs: np.ndarray,
        h: float,
    ) -> np.ndarray:
        """
        Predictor step (Adams-Bashforth type).

        Args:
            f: Right-hand side function
            t_values: Time points
            y_values: Solution values
            n: Current time step
            alpha: Fractional order
            coeffs: Fractional coefficients
            h: Step size

        Returns:
            Predicted solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Compute fractional derivative term
        frac_term = 0.0
        for j in range(n):
            frac_term += coeffs[j] * (y_values[n - j] - y_values[n - j - 1])

        # Predictor formula
        t_n = t_values[n]
        y_pred = y_values[n - 1] + (h**alpha_val / gamma(alpha_val + 1)) * (
            f(t_n, y_values[n - 1]) - frac_term
        )

        return y_pred

    def _corrector_step(
        self,
        f: Callable,
        t_values: np.ndarray,
        y_values: np.ndarray,
        y_pred: np.ndarray,
        n: int,
        alpha: Union[float, FractionalOrder],
        coeffs: np.ndarray,
        h: float,
    ) -> np.ndarray:
        """
        Corrector step (Adams-Moulton type).

        Args:
            f: Right-hand side function
            t_values: Time points
            y_values: Solution values
            y_pred: Predicted solution
            n: Current time step
            alpha: Fractional order
            coeffs: Fractional coefficients
            h: Step size

        Returns:
            Corrected solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Compute fractional derivative term
        frac_term = 0.0
        for j in range(n):
            frac_term += coeffs[j] * (y_values[n - j] - y_values[n - j - 1])

        # Corrector formula
        t_n = t_values[n]
        y_corr = y_values[n - 1] + (h**alpha_val / gamma(alpha_val + 1)) * (
            0.5 * (f(t_n, y_pred) +
                   f(t_values[n - 1], y_values[n - 1])) - frac_term
        )

        return y_corr

    def _adaptive_predictor_step(
        self,
        f: Callable,
        t_current: float,
        t_next: float,
        y_current: np.ndarray,
        alpha: Union[float, FractionalOrder],
        h: float,
    ) -> np.ndarray:
        """
        Adaptive predictor step.

        Args:
            f: Right-hand side function
            t_current: Current time
            t_next: Next time
            y_current: Current solution
            alpha: Fractional order
            h: Step size

        Returns:
            Predicted solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Simple predictor
        y_pred = y_current + (h**alpha_val / gamma(alpha_val + 1)) * f(
            t_current, y_current
        )

        return y_pred

    def _adaptive_corrector_step(
        self,
        f: Callable,
        t_current: float,
        t_next: float,
        y_current: np.ndarray,
        y_pred: np.ndarray,
        alpha: Union[float, FractionalOrder],
        h: float,
    ) -> np.ndarray:
        """
        Adaptive corrector step.

        Args:
            f: Right-hand side function
            t_current: Current time
            t_next: Next time
            y_current: Current solution
            y_pred: Predicted solution
            alpha: Fractional order
            h: Step size

        Returns:
            Corrected solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Corrector with iteration
        y_corr = y_pred
        for _ in range(self.max_iter):
            y_old = y_corr.copy()
            y_corr = y_current + (h**alpha_val / gamma(alpha_val + 1)) * (
                0.5 * (f(t_next, y_pred) + f(t_current, y_current))
            )

            if np.allclose(y_corr, y_old, rtol=self.tol):
                break

        return y_corr

    def _estimate_error(self, y_pred: np.ndarray, y_corr: np.ndarray) -> float:
        """
        Estimate local truncation error.

        Args:
            y_pred: Predicted solution
            y_corr: Corrected solution

        Returns:
            Estimated error
        """
        return np.linalg.norm(y_corr - y_pred)


class AdamsBashforthMoultonSolver(PredictorCorrectorSolver):
    """
    Adams-Bashforth-Moulton solver for fractional differential equations.

    Implements the classical Adams-Bashforth-Moulton scheme adapted for
    fractional differential equations.
    """

    def __init__(
        self,
        derivative_type: str = "caputo",
        order: int = 1,
        adaptive: bool = True,
        tol: float = 1e-6,
        max_iter: int = 10,
    ):
        """
        Initialize Adams-Bashforth-Moulton solver.

        Args:
            derivative_type: Type of fractional derivative
            order: Order of the method
            adaptive: Use adaptive step size control
            tol: Tolerance for convergence
            max_iter: Maximum number of corrector iterations
        """
        super().__init__(derivative_type, order, adaptive, tol, max_iter)

    def _predictor_step(
        self,
        f: Callable,
        t_values: np.ndarray,
        y_values: np.ndarray,
        n: int,
        alpha: Union[float, FractionalOrder],
        coeffs: np.ndarray,
        h: float,
    ) -> np.ndarray:
        """
        Adams-Bashforth predictor step.

        Args:
            f: Right-hand side function
            t_values: Time points
            y_values: Solution values
            n: Current time step
            alpha: Fractional order
            coeffs: Fractional coefficients
            h: Step size

        Returns:
            Predicted solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Adams-Bashforth weights
        weights = self._compute_adams_bashforth_weights(n, alpha_val)

        # Predictor formula
        t_values[n]
        y_pred = y_values[n - 1]

        for j in range(n):
            t_j = t_values[j]
            y_pred += (
                h**alpha_val / gamma(alpha_val + 1) *
                weights[j] * f(t_j, y_values[j])
            )

        return y_pred

    def _corrector_step(
        self,
        f: Callable,
        t_values: np.ndarray,
        y_values: np.ndarray,
        y_pred: np.ndarray,
        n: int,
        alpha: Union[float, FractionalOrder],
        coeffs: np.ndarray,
        h: float,
    ) -> np.ndarray:
        """
        Adams-Moulton corrector step.

        Args:
            f: Right-hand side function
            t_values: Time points
            y_values: Solution values
            y_pred: Predicted solution
            n: Current time step
            alpha: Fractional order
            coeffs: Fractional coefficients
            h: Step size

        Returns:
            Corrected solution
        """
        if isinstance(alpha, FractionalOrder):
            alpha_val = alpha.alpha
        else:
            alpha_val = alpha

        # Adams-Moulton weights
        weights = self._compute_adams_moulton_weights(n, alpha_val)

        # Corrector formula
        t_values[n]
        y_corr = y_values[n - 1]

        for j in range(n):
            t_j = t_values[j]
            if j == n - 1:
                y_corr += (
                    h**alpha_val / gamma(alpha_val + 1) *
                    weights[j] * f(t_j, y_pred)
                )
            else:
                y_corr += (
                    h**alpha_val
                    / gamma(alpha_val + 1)
                    * weights[j]
                    * f(t_j, y_values[j])
                )

        return y_corr

    def _compute_adams_bashforth_weights(
            self, n: int, alpha: float) -> np.ndarray:
        """
        Compute Adams-Bashforth weights.

        Args:
            n: Current time step
            alpha: Fractional order

        Returns:
            Array of weights
        """
        weights = np.zeros(n)

        # Simplified Adams-Bashforth weights
        for j in range(n):
            weights[j] = (n - j) ** alpha - (n - j - 1) ** alpha

        return weights

    def _compute_adams_moulton_weights(
            self, n: int, alpha: float) -> np.ndarray:
        """
        Compute Adams-Moulton weights.

        Args:
            n: Current time step
            alpha: Fractional order

        Returns:
            Array of weights
        """
        weights = np.zeros(n)

        # Simplified Adams-Moulton weights
        for j in range(n):
            if j == n - 1:
                weights[j] = 1.0
            else:
                weights[j] = (n - j) ** alpha - (n - j - 1) ** alpha

        return weights


class VariableStepPredictorCorrector(PredictorCorrectorSolver):
    """
    Variable step size predictor-corrector solver.

    Implements predictor-corrector methods with sophisticated step size
    control and error estimation.
    """

    def __init__(
        self,
        derivative_type: str = "caputo",
        order: int = 1,
        tol: float = 1e-6,
        max_iter: int = 10,
        min_h: float = 1e-8,
        max_h: float = 1e-2,
        safety_factor: float = 0.9,
    ):
        """
        Initialize variable step predictor-corrector solver.

        Args:
            derivative_type: Type of fractional derivative
            order: Order of the method
            tol: Tolerance for error control
            max_iter: Maximum number of corrector iterations
            min_h: Minimum step size
            max_h: Maximum step size
            safety_factor: Safety factor for step size control
        """
        super().__init__(derivative_type, order, True, tol, max_iter, min_h, max_h)
        self.safety_factor = safety_factor

    def _estimate_error(self, y_pred: np.ndarray, y_corr: np.ndarray) -> float:
        """
        Enhanced error estimation.

        Args:
            y_pred: Predicted solution
            y_corr: Corrected solution

        Returns:
            Estimated error
        """
        # Use relative error for better control
        error = np.linalg.norm(y_corr - y_pred) / \
            (np.linalg.norm(y_corr) + 1e-12)
        return error

    def _adaptive_step_size_control(
        self, error: float, h_current: float, order: int
    ) -> float:
        """
        Adaptive step size control.

        Args:
            error: Estimated error
            h_current: Current step size
            order: Method order
            alpha: Fractional order

        Returns:
            New step size
        """
        if error <= self.tol:
            # Accept step and increase step size
            factor = min(
                2.0, self.safety_factor *
                (self.tol / error) ** (1.0 / (order + 1))
            )
            h_new = min(self.max_h, h_current * factor)
        else:
            # Reject step and decrease step size
            factor = max(
                0.1, self.safety_factor *
                (self.tol / error) ** (1.0 / (order + 1))
            )
            h_new = max(self.min_h, h_current * factor)

        return h_new


# Convenience functions
def solve_predictor_corrector(
    f: Callable,
    t_span: Tuple[float, float],
    y0: Union[float, np.ndarray],
    alpha: Union[float, FractionalOrder],
    derivative_type: str = "caputo",
    method: str = "standard",
    adaptive: bool = True,
    h: Optional[float] = None,
    **kwargs,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Solve using predictor-corrector method.

    Args:
        f: Right-hand side function f(t, y)
        t_span: Time interval (t0, tf)
        y0: Initial condition(s)
        alpha: Fractional order
        derivative_type: Type of fractional derivative
        method: Method type ("standard", "adams_bashforth_moulton", "variable_step")
        adaptive: Use adaptive step size control
        h: Step size (None for adaptive)
        **kwargs: Additional solver parameters

    Returns:
        Tuple of (t_values, y_values)
    """
    if method == "standard":
        solver = PredictorCorrectorSolver(derivative_type, adaptive=adaptive)
    elif method == "adams_bashforth_moulton":
        solver = AdamsBashforthMoultonSolver(
            derivative_type, adaptive=adaptive)
    elif method == "variable_step":
        solver = VariableStepPredictorCorrector(derivative_type)
    else:
        raise ValueError(f"Unknown method: {method}")

    return solver.solve(f, t_span, y0, alpha, h, **kwargs)


def solve_adams_bashforth_moulton(
    f: Callable,
    t_span: Tuple[float, float],
    y0: Union[float, np.ndarray],
    alpha: Union[float, FractionalOrder],
    derivative_type: str = "caputo",
    adaptive: bool = True,
    h: Optional[float] = None,
    **kwargs,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Solve using Adams-Bashforth-Moulton method.

    Args:
        f: Right-hand side function f(t, y)
        t_span: Time interval (t0, tf)
        y0: Initial condition(s)
        alpha: Fractional order
        derivative_type: Type of fractional derivative
        adaptive: Use adaptive step size control
        h: Step size (None for adaptive)
        **kwargs: Additional solver parameters

    Returns:
        Tuple of (t_values, y_values)
    """
    solver = AdamsBashforthMoultonSolver(derivative_type, adaptive=adaptive)
    return solver.solve(f, t_span, y0, alpha, h, **kwargs)


def solve_variable_step_predictor_corrector(
    f: Callable,
    t_span: Tuple[float, float],
    y0: Union[float, np.ndarray],
    alpha: Union[float, FractionalOrder],
    derivative_type: str = "caputo",
    h0: Optional[float] = None,
    **kwargs,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Solve using variable step size predictor-corrector method.

    Args:
        f: Right-hand side function f(t, y)
        t_span: Time interval (t0, tf)
        y0: Initial condition(s)
        alpha: Fractional order
        derivative_type: Type of fractional derivative
        h0: Initial step size
        **kwargs: Additional solver parameters

    Returns:
        Tuple of (t_values, y_values)
    """
    solver = VariableStepPredictorCorrector(derivative_type)
    return solver.solve(f, t_span, y0, alpha, h0, **kwargs)
