"""Finite-difference and other numerical helpers."""

from __future__ import annotations

from collections.abc import Callable
from typing import Literal

import numpy as np
from numpy.typing import ArrayLike, NDArray

from .errors import InvalidInput
from .utils import assert_finite, assert_finite_array

FiniteDifferenceMethod = Literal["central", "forward", "backward"]


def finite_difference_derivative(
    f: Callable[[float], float],
    x: float,
    *,
    h: float = 1e-6,
    method: FiniteDifferenceMethod = "central",
) -> float:
    x = assert_finite(x, name="x")
    h = assert_finite(h, name="h")
    if h <= 0.0:
        raise InvalidInput("h must be positive.")

    fx = assert_finite(f(x), name="f(x)")
    if method == "central":
        fx_plus = assert_finite(f(x + h), name="f(x+h)")
        fx_minus = assert_finite(f(x - h), name="f(x-h)")
        return (fx_plus - fx_minus) / (2.0 * h)
    if method == "forward":
        fx_plus = assert_finite(f(x + h), name="f(x+h)")
        return (fx_plus - fx) / h
    if method == "backward":
        fx_minus = assert_finite(f(x - h), name="f(x-h)")
        return (fx - fx_minus) / h
    raise InvalidInput(f"Unknown finite-difference method: {method!r}.")


def finite_difference_gradient(
    f: Callable[[NDArray[np.float64]], float],
    x: NDArray[np.float64] | ArrayLike,
    *,
    h: float = 1e-6,
) -> NDArray[np.float64]:
    x_arr = assert_finite_array(x, name="x").astype(float, copy=False)
    if x_arr.ndim != 1:
        raise InvalidInput(f"x must be 1D; got shape {x_arr.shape}.")
    h = assert_finite(h, name="h")
    if h <= 0.0:
        raise InvalidInput("h must be positive.")

    grad = np.empty_like(x_arr, dtype=float)
    for i in range(x_arr.size):
        step = np.zeros_like(x_arr)
        step[i] = h
        f_plus = assert_finite(float(f(x_arr + step)), name="f(x+h)")
        f_minus = assert_finite(float(f(x_arr - step)), name="f(x-h)")
        grad[i] = (f_plus - f_minus) / (2.0 * h)
    return grad


def finite_difference_jacobian(
    residuals_fn: Callable[[NDArray[np.float64]], NDArray[np.float64] | ArrayLike],
    x: NDArray[np.float64] | ArrayLike,
    *,
    h: float = 1e-6,
) -> NDArray[np.float64]:
    x_arr = assert_finite_array(x, name="x").astype(float, copy=False)
    if x_arr.ndim != 1:
        raise InvalidInput(f"x must be 1D; got shape {x_arr.shape}.")
    h = assert_finite(h, name="h")
    if h <= 0.0:
        raise InvalidInput("h must be positive.")

    r0 = assert_finite_array(residuals_fn(x_arr), name="residuals(x)")
    if r0.ndim != 1:
        raise InvalidInput(f"residuals_fn must return 1D; got shape {r0.shape}.")
    m = int(r0.size)
    n = int(x_arr.size)
    jac = np.empty((m, n), dtype=float)

    for j in range(n):
        step = np.zeros_like(x_arr)
        step[j] = h
        r_plus = assert_finite_array(residuals_fn(x_arr + step), name="residuals(x+h)")
        r_minus = assert_finite_array(residuals_fn(x_arr - step), name="residuals(x-h)")
        if r_plus.shape != r0.shape or r_minus.shape != r0.shape:
            raise InvalidInput("residuals_fn returned inconsistent shapes.")
        jac[:, j] = (r_plus - r_minus) / (2.0 * h)
    if not np.all(np.isfinite(jac)):
        raise InvalidInput("Jacobian contains non-finite values.")
    return jac

