from typing import Callable

import numpy as np

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


def _fletcher_reeves(g_new, g_old, d, y):
    return np.dot(g_new, g_new) / (np.dot(g_old, g_old) + 1e-12)


def _polak_ribiere(g_new, g_old, d, y):
    return np.dot(g_new, y) / (np.dot(g_old, g_old) + 1e-12)


def _hestenes_stiefel(g_new, g_old, d, y):
    return np.dot(g_new, y) / (np.dot(d, y) + 1e-12)


def _dai_yuan(g_new, g_old, d, y):
    return np.dot(g_new, g_new) / (np.dot(d, y) + 1e-12)


BETA_METHODS = {
    "fletcher-reeves": _fletcher_reeves,
    "polak-ribiere": _polak_ribiere,
    "hestenes-stiefel": _hestenes_stiefel,
    "dai-yuan": _dai_yuan,
}


def ncg_line_search(x, f, grad_x, d, gamma):
    f_x = f(x)
    slope = np.dot(grad_x, d)
    next_x = x + gamma * d
    iters = 0
    while True:
        if f(next_x) <= f_x + 0.1 * gamma * slope or gamma < 1e-12:
            debug(f"Line search: gamma={gamma:.4e} after {iters} shrinks")
            break
        gamma /= 2
        next_x = x + gamma * d
        iters += 1
    return next_x, gamma * 1.2


def ncg(
    f: Callable,
    grad_f: Callable,
    x0: np.ndarray,
    method: str = "polak-ribiere",
    max_iter: int = 100,
    tol: float = 1e-6,
    epsilon: float = 1e-8,
    return_loss: bool = False,
):
    x = x0.copy()
    loss = []

    info(f"Starting Nonlinear CG [NumPy] using {method}")

    beta_func = BETA_METHODS[method]
    g_old = grad_f(x)
    d = -g_old
    gamma = 1

    for i in range(max_iter):
        x_new, gamma = ncg_line_search(x, f, g_old, d, gamma)
        f_val = f(x_new)
        loss.append(f_val.item())

        g_new = grad_f(x_new)
        g_norm = np.linalg.norm(g_new)
        f_delta = np.abs(f(x) - f_val)

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

        if f_delta < tol or g_norm < epsilon:
            info("NCG converged successfully.")
            x = x_new
            break

        y = g_new - g_old
        beta = max(0, beta_func(g_new, g_old, d, y))
        d = -g_new + beta * d

        if (i + 1) % x.size == 0:
            debug("NCG: Performing periodic reset of search direction")
            d = -g_new

        x, g_old = x_new, g_new
    else:
        warn(f"NCG reached max_iter ({max_iter})")

    return (x, loss) if return_loss else x
