from typing import Callable, Optional
import numpy as np
from vicentin.utils.logging import info, debug, warn


def gamma_step(x: np.ndarray, f: Callable, grad_x: np.ndarray, gamma: float):
    return x - gamma * grad_x, gamma


def line_search_step(
    x: np.ndarray, f: Callable, grad_x: np.ndarray, gamma: float
):
    next_x = x - gamma * grad_x
    f_x = f(x)
    iters = 0

    while True:
        armijo = f(next_x) <= f_x + np.dot(grad_x, next_x - x) + np.sum(
            (next_x - x) ** 2
        ) / (2 * gamma)

        if armijo or gamma < 1e-12:
            if not armijo:
                warn(
                    "Line search reached minimum gamma limit without Armijo condition."
                )
            else:
                debug(
                    f"Line search converged: gamma={gamma:.4e} after {iters} shrinks"
                )
            break

        gamma /= 2
        next_x = x - gamma * grad_x
        iters += 1

    gamma *= 1.2  # Heuristic to speed up
    return x - gamma * grad_x, gamma


def gradient_descent(
    f,
    grad_f,
    x0: np.ndarray,
    step_size: Optional[float] = None,
    max_iter: int = 100,
    tol: float = 1e-6,
    epsilon: float = 1e-8,
    return_loss: bool = False,
):
    x = x0.copy()
    loss = []

    if step_size is None:
        step_func = line_search_step
        gamma = 1
        info("Starting Gradient Descent [NumPy] with Line Search")
    else:
        step_func = gamma_step
        gamma = step_size
        info(f"Starting Gradient Descent [NumPy] with Fixed Step Size: {gamma}")

    for i in range(max_iter):
        grad_x = grad_f(x)
        x_new, gamma = step_func(x, f, grad_x, gamma)

        f_new = f(x_new)
        f_val = f_new.item()
        loss.append(f_val)

        norm_f = np.abs(f(x) - f_new)
        norm_grad = np.linalg.norm(grad_x)

        info(
            f"Iter {i}: f(x)={f_val:.6e} | Grad Norm={norm_grad:.4e} | Delta f={norm_f:.4e}"
        )

        x = x_new

        if norm_f < tol:
            info(
                f"Success: Convergence reached via function tolerance ({tol})."
            )
            break
        if norm_grad < epsilon:
            info(
                f"Success: Convergence reached via gradient tolerance ({epsilon})."
            )
            break
    else:
        warn(f"Reached max_iter ({max_iter}) without full convergence.")

    if return_loss:
        return x, loss

    return x
