from typing import Any, Callable, Optional

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

dispatcher = Dispatcher()

try:
    from .conjugate_gradient_np import conjugate_gradient as cg_np

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

try:
    from .conjugate_gradient_torch import conjugate_gradient as cg_torch

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


def conjugate_gradient(
    A: Callable[[Any], Any],
    b: Any,
    x0: Optional[Any] = None,
    max_iter: int = 100,
    tol: float = 1e-8,
    return_loss: bool = True,
    backend: Optional[str] = None,
):
    """
    Solves the linear system $Ax = b$ using the Linear Conjugate Gradient method.

    This algorithm is specifically designed for symmetric positive-definite
    matrices, such as the Gram matrices found in RKHS problems.
    Unlike Gradient Descent, it finds the exact solution in at most $N$
    steps for an $N \times N$ system in exact arithmetic.

    Advantages:
    -----------
    - Memory Efficient: Uses your backend's `@` operator and `chunk_size`
      to solve without materializing $A$.
    - Zero Search: Uses analytical formulas for step size, eliminating the
      need for line searches.

    Parameters:
    -----------
    A : Callable
        A function that computes the matrix-vector product $A @ v$.
        In your study, this is usually `lambda v: kernel @ v + lambda_reg * v`.
    b : Any
        The target vector (e.g., labels $y$).
    x0 : Any, optional
        Initial guess for the solution.
    max_iter : int, optional (default=100)
    tol : float, optional (default=1e-8)
    return_loss : bool, optional (default=True)
    """

    ctx = dispatcher.detect_backend(b, backend)

    x0, b = dispatcher.cast_values(ctx, x0, b)

    return dispatcher(ctx, A, b, x0, max_iter, tol, return_loss)
