from typing import Any, Callable, Optional, Sequence

from vicentin.utils import Dispatcher
from vicentin.utils.logging import error, warn


def transform_numpy(*args, **kwargs):
    if args:
        F = args[0]
        if not isinstance(F, (list, tuple)) or len(F) < 2:
            raise ValueError("NumPy backend requires F=(f, grad_f)")

        f, grad_f = F
        new_args = (f, grad_f) + args[1:]
        return new_args, kwargs

    F = kwargs.pop("F")
    if not isinstance(F, (list, tuple)) or len(F) < 2:
        raise ValueError("NumPy backend requires F=(f, grad_f)")

    return args, {"f": F[0], "grad_f": F[1], **kwargs}


def transform_autodiff(*args, **kwargs):
    if args:
        F = args[0]
        f = F[0] if isinstance(F, (list, tuple)) else F

        new_args = (f,) + args[1:]
        return new_args, kwargs

    F = kwargs.pop("F")
    f = F[0] if isinstance(F, (list, tuple)) else F

    return args, {"f": f, **kwargs}


dispatcher = Dispatcher()

try:
    from .ncg_np import ncg as ncg_np

    dispatcher.register("numpy", ncg_np, transform_numpy)
except ModuleNotFoundError:
    warn("NumPy backend not found.")
except Exception:
    error("NumPy backend failed to load.")

try:
    from .ncg_torch import ncg as ncg_torch

    dispatcher.register("torch", ncg_torch, transform_autodiff)
except ModuleNotFoundError:
    warn("Torch backend not found.")
except Exception:
    error("Torch backend failed to load.")


def nonlinear_conjugate_gradient(
    F: Sequence[Callable] | Callable,
    x0: Any,
    method: str = "polak-ribiere",
    max_iter: int = 100,
    tol: float = 1e-6,
    epsilon: float = 1e-8,
    return_loss: bool = True,
    backend: Optional[str] = None,
):
    """
    Minimizes a function using Nonlinear Conjugate Gradient (NCG).

    NCG is an iterative method that constructs search directions $d_n$ as a
    linear combination of the current negative gradient and the previous search
    direction, ensuring conjugacy for quadratic problems.

    Update Methods (Beta Formulas):
    ------------------------------
    * 'polak-ribiere' (Default):
        - Pros: Very robust for non-quadratic functions; automatically resets
          to gradient descent when steps are small.
        - Cons: Requires an accurate line search to maintain stability.
    * 'fletcher-reeves':
        - Pros: Simplest implementation; mathematically elegant for pure quadratics.
        - Cons: Very sensitive to line search errors; can stall on non-linear surfaces.
    * 'hestenes-stiefel':
        - Pros: Often exhibits faster convergence in high-dimensional RKHS
          problems due to its use of the search direction in the denominator.
        - Cons: Mathematically equivalent to Polak-Ribière only on quadratics.
    * 'dai-yuan':
        - Pros: Guaranteed descent property even with less precise line searches.
        - Cons: Can be slower to converge than PR in practice.

    Parameters:
    -----------
    F : Sequence[Callable] or Callable
        Problem definition. (f, grad_f) for NumPy; f for Torch.
    x0 : Any
        Initial guess. Type determines the backend.
    method : str, optional (default="polak-ribiere")
        The beta calculation formula used to enforce conjugacy.
    max_iter : int, optional (default=100)
        Maximum number of iterations.
    tol : float, optional (default=1e-6)
        Convergence tolerance for function value change.
    epsilon : float, optional (default=1e-8)
        Convergence tolerance for gradient norm.
    return_loss : bool, optional (default=True)

    Returns:
    --------
    x : Any
        Approximate local minimum.
    loss : List[float], optional
    """

    ctx = dispatcher.detect_backend(x0, backend)
    x0 = dispatcher.cast_values(ctx, x0)

    return dispatcher(ctx, F, x0, method, max_iter, tol, epsilon, return_loss)
