from typing import Callable, Any

import numpy as np

from vicentin.utils.logging import info, debug, warn


def newton_raphson(
    f: Callable[[np.ndarray], np.ndarray],
    grad_f: Callable[[np.ndarray], np.ndarray],
    x0: Any,
    max_iter: int = 100,
    tol: float = 1e-6,
    return_loss: bool = False,
):

    x = x0.copy()
    loss = []

    info(f"Starting Newton-Raphson [NumPy] (max_iter={max_iter}, tol={tol})")

    for i in range(max_iter):
        f_val = np.atleast_1d(f(x))

        if np.isinf(f_val).any():
            error_msg = f"Function diverged at iteration {i}."
            warn(error_msg)
            raise RuntimeError(error_msg)

        f_norm = float(np.linalg.norm(f_val))
        loss.append(f_norm)

        debug(f"Iter {i}: ||f(x)|| = {f_norm:.4e}")

        if f_norm < tol:
            info(f"Newton-Raphson converged at iteration {i}.")
            break

        grad = grad_f(x)

        if np.any(grad == 0):
            error_msg = f"Gradient is zero at iteration {i}."
            warn(error_msg)
            raise RuntimeError(error_msg)

        if grad.ndim > 1:
            delta_x = np.linalg.solve(grad, f_val)
        else:
            delta_x = f_val / grad

        x = x - delta_x
    else:
        warn(
            f"Newton-Raphson reached max_iter ({max_iter}) without full convergence."
        )

    return (x, loss) if return_loss else x
