from typing import Callable, Optional

import numpy as np
from numpy.linalg import norm, solve

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


def backtrack_line_search(
    f: Callable,
    r: Callable,
    x: np.ndarray,
    delta_x: np.ndarray,
    w: np.ndarray,
    delta_w: np.ndarray,
    alpha: float,
    beta: float,
):
    y = r(x, w)
    t = 1
    iters = 0

    while True:
        f_x = f(x + t * delta_x)
        infeasible = np.isnan(f_x) or np.isinf(f_x)
        feasible = not infeasible

        if feasible:
            r_val = r(x + t * delta_x, w + t * delta_w)

            lipschitz = r_val <= (1 - alpha * t) * y

            if lipschitz:
                debug(f"Line search converged: t={t:.4e} after {iters} shrinks")
                break

        t *= beta
        iters += 1

        if t < 1e-12:
            warn("Line search failed to find feasible step; returning t=0")
            return 0

    return t


def default_solver(
    hess_f: Callable,
    grad_f: Callable,
    x: np.ndarray,
    w: np.ndarray,
    A: np.ndarray,
    b: np.ndarray,
):
    grad = grad_f(x)
    H = hess_f(x)

    min_eig = np.linalg.svd(H, compute_uv=False).min()
    if min_eig < 1e-12:
        debug(
            f"Hessian near-singular (min_svd={min_eig:.2e}). Applying jitter."
        )
        H = H + 1e-8 * np.eye(H.shape[-1])

    g = grad + A.T @ w
    h = A @ x - b

    inv_H_A = solve(H, A.T)
    inv_H_g = solve(H, g)

    S = -A @ inv_H_A

    delta_w = solve(S, A @ inv_H_g - h)
    delta_x = solve(H, -A.T @ delta_w - g)

    decrement_squared = delta_x @ H @ delta_x

    debug(
        f"Linear system solved. Newton decrement squared: {decrement_squared:.4e}"
    )
    return delta_x, delta_w, decrement_squared


def newton_step(
    f: Callable,
    grad_f: Callable,
    hess_f: Callable,
    r: Callable,
    x: np.ndarray,
    w: np.ndarray,
    A: np.ndarray,
    b: np.ndarray,
    alpha: float = 0.25,
    beta: float = 0.5,
    linear_solver: Optional[Callable] = None,
):
    if linear_solver is None:
        linear_solver = default_solver

    delta_x, delta_w, decrement_squared = linear_solver(
        hess_f, grad_f, x, w, A, b
    )

    t = backtrack_line_search(f, r, x, delta_x, w, delta_w, alpha, beta)

    x = x + t * delta_x
    w = w + t * delta_w

    return x, w, decrement_squared


def newton(
    f: Callable,
    grad_f: Callable,
    hess_f: Callable,
    x0: np.ndarray,
    equality: Optional[tuple] = None,
    max_iter: int = 100,
    tol: float = 1e-8,
    epsilon: float = 1e-4,
    alpha: float = 0.25,
    beta: float = 0.5,
    linear_solver: Optional[Callable] = None,
    return_dual: bool = False,
    return_loss: bool = False,
):
    x = x0.copy()
    i = 0
    loss = []

    if equality is None:
        A = np.empty((0, x.size))
        b = np.empty(0)
        w = np.empty(0)
    else:
        A, b = equality
        w = np.zeros(A.shape[0])

        info(f"Initialized with {A.shape[0]} equality constraints.")

    r = lambda x, w: norm(np.concatenate([grad_f(x) + A.T @ w, A @ x - b]))

    y_new = f(x)

    info(f"Starting solver with {max_iter=}")

    while True:
        y = y_new
        if y == np.inf:
            raise ValueError(f"Reached infeasible point: {x}.")

        x, w, decrement_squared = newton_step(
            f, grad_f, hess_f, r, x, w, A, b, alpha, beta, linear_solver
        )

        y_new = f(x)
        r_val = r(x, w)

        loss.append(y_new)

        feasible = np.isclose(A @ x, b).all()

        info(
            f"Iter {i}: f(x)={y_new:.6e}, Residual={r_val:.4e}, Feasible={feasible}"
        )

        if (not feasible and r_val <= epsilon) or (
            feasible and decrement_squared <= 2 * epsilon
        ):
            info("Success: Convergence reached via Newton decrement/residual.")
            break

        if np.abs(y_new - y) < tol:
            info("Success: Convergence reached via function value tolerance.")
            break

        i += 1

        if i >= max_iter:
            warn(
                f"Warning: Reached max_iter ({max_iter}) without full convergence."
            )
            break

    output = [x]

    if return_dual:
        debug("Appending dual solution to return")
        output.append(w)

    if return_loss:
        debug("Appending losses to return")
        output.append(loss)

    if len(output) == 1:
        return output[0]

    return output
