# Solve a singularly perturbed nonlinear two-point boundary-value problem.
#
# The finite-difference method solves
#
#     y'' = f(x, y, y', epsilon),       x_left < x < x_right,
#
# with Dirichlet boundary conditions
#
#     y(x_left) = y_left,
#     y(x_right) = y_right.
#
# Thus a problem written as
#
#     epsilon*y'' = F(x, y, y', epsilon)
#
# should be supplied using
#
#     f = F / epsilon.
#
# The method uses:
#     * second-order centered differences,
#     * Newton's method,
#     * a tridiagonal linear solver,
#     * continuation in epsilon from 10**e_max to 10**e_min.

from __future__ import annotations

from collections.abc import Callable

import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray


ScalarFunction = Callable[[float, float, float, float], float]


def solve_sbvp(
    f: ScalarFunction,
    fy: ScalarFunction,
    fz: ScalarFunction,
    x_left: float,
    y_left: float,
    x_right: float,
    y_right: float,
    nx: int = 1000,
    n_epsilon: int = 160,
    e_max: float = 0.0,
    e_min: float = -4.0,
    newton_tolerance: float = 1.0e-5,
    max_newton_iterations: int = 100,
    print_progress: bool = True,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Solve a singularly perturbed nonlinear boundary-value problem."""

    if x_right <= x_left:
        raise ValueError("x_right must be greater than x_left.")
    if nx < 1:
        raise ValueError("nx must be at least 1.")
    if n_epsilon < 1:
        raise ValueError("n_epsilon must be at least 1.")
    if e_min > e_max:
        raise ValueError("e_min must not exceed e_max.")
    if newton_tolerance <= 0.0:
        raise ValueError("newton_tolerance must be positive.")
    if max_newton_iterations < 1:
        raise ValueError("max_newton_iterations must be at least 1.")

    # nx interior points plus the two boundary points.
    x = np.linspace(x_left, x_right, nx + 2, dtype=float)
    dx = x[1] - x[0]
    dx2 = dx*dx

    # Initial guess: the line joining the two boundary values.
    y = y_left + (y_right - y_left)*(x - x_left)/(x_right - x_left)

    # Start at 10**e_max and continue down to 10**e_min.
    epsilon_values = 10.0**np.linspace(e_max, e_min, n_epsilon)

    # Work arrays for the tridiagonal Newton system.
    diagonal = np.zeros(nx)
    upper = np.zeros(nx)
    correction = np.zeros(nx)
    factor = np.zeros(nx)

    for epsilon in epsilon_values:
        converged = False

        for iteration in range(1, max_newton_iterations + 1):

            # Assemble the Newton system at the interior grid points.
            for j in range(1, nx + 1):
                k = j - 1
                z = (y[j + 1] - y[j - 1])/(2.0*dx)

                diagonal[k] = 2.0 + dx2*fy(x[j], y[j], z, epsilon)
                upper[k] = -1.0 - 0.5*dx*fz(x[j], y[j], z, epsilon)
                correction[k] = (
                    -2.0*y[j]
                    + y[j + 1]
                    + y[j - 1]
                    - dx2*f(x[j], y[j], z, epsilon)
                )

            # Forward sweep for the tridiagonal Newton system.
            if abs(diagonal[0]) <= np.finfo(float).eps:
                raise ZeroDivisionError(
                    f"Zero pivot in Newton system at epsilon={epsilon:.6e}."
                )

            correction[0] = correction[0]/diagonal[0]
            factor[0] = -(2.0 + upper[0])/diagonal[0]

            for j in range(1, nx):
                pivot = diagonal[j] - upper[j]*factor[j - 1]

                if abs(pivot) <= np.finfo(float).eps:
                    raise ZeroDivisionError(
                        f"Zero pivot in Newton system at epsilon={epsilon:.6e}, "
                        f"interior index {j + 1}."
                    )

                correction[j] = (
                    correction[j] - upper[j]*correction[j - 1]
                )/pivot
                factor[j] = -(2.0 + upper[j])/pivot

            # Back substitution and update.
            delta = correction[nx - 1]
            y[nx] = y[nx] + delta
            error = abs(delta)

            for j in range(nx - 1, 0, -1):
                delta = correction[j - 1] - factor[j - 1]*delta
                y[j] = y[j] + delta
                error = max(error, abs(delta))

            # Reimpose the prescribed boundary values.
            y[0] = y_left
            y[nx + 1] = y_right

            if not np.all(np.isfinite(y)):
                raise FloatingPointError(
                    f"Nonfinite solution encountered at epsilon={epsilon:.6e}."
                )

            if error <= newton_tolerance:
                converged = True
                break

        if print_progress:
            print(
                f"epsilon = {epsilon:.6e}; "
                f"Newton iterations = {iteration}"
            )

        if not converged:
            raise RuntimeError(
                "Newton's method did not converge within "
                f"{max_newton_iterations} iterations for "
                f"epsilon={epsilon:.6e}. Increase n_epsilon or nx, "
                "or choose a less negative e_min."
            )

    return x, y


# ---------------------------------------------------------------------
# Example from the MATLAB program:
#
#     epsilon*y'' = -tanh(y') + y - 1,
#     y(0) = 3,  y(1) = 5.
#
# Because solve_sbvp expects y'' = f, divide the right-hand side and its
# derivatives by epsilon.
# ---------------------------------------------------------------------


def example_f(x: float, y: float, z: float, epsilon: float) -> float:
    # Right-hand side for the example problem.
    del x
    return (-np.tanh(z) + y - 1.0)/epsilon


def example_fy(x: float, y: float, z: float, epsilon: float) -> float:
    # Partial derivative of example_f with respect to y.
    del x, y, z
    return 1.0/epsilon


def example_fz(x: float, y: float, z: float, epsilon: float) -> float:
    # Partial derivative of example_f with respect to z.
    del x, y
    # Use sech(z)**2 = 1 - tanh(z)**2 to avoid overflow in cosh(z).
    sech_squared = 1.0 - np.tanh(z)**2
    return -sech_squared/epsilon


def main() -> None:
    # Solve and plot the example problem.
    x, y = solve_sbvp(
        f=example_f,
        fy=example_fy,
        fz=example_fz,
        x_left=0.0,
        y_left=3.0,
        x_right=1.0,
        y_right=5.0,
        nx=1000,
        n_epsilon=160,
        e_max=0.0,
        e_min=-4.0,
        newton_tolerance=1.0e-5,
        max_newton_iterations=100,
        print_progress=True,
    )

    plt.plot(x, y)
    plt.xlabel("x-axis", fontsize=14, fontweight="bold")
    plt.ylabel("Solution", fontsize=14, fontweight="bold")
    plt.grid(True)
    plt.box(True)
    plt.tick_params(axis="both", labelsize=14)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()


# ---------------------------------------------------------------------
# Running the program
#
# Save this file as sbvp.py and run it from a terminal with
#
#     python sbvp.py
#
# The file must be in the current folder, or its folder must be on the
# Python path. NumPy and Matplotlib are required.
#
# Needed input from the user
#
#     f(x,y,z,epsilon):   right-hand side of y'' = f, where z represents y'
#     fy(x,y,z,epsilon):  derivative of f with respect to y
#     fz(x,y,z,epsilon):  derivative of f with respect to z
#
# The boundary conditions and numerical parameters are specified in main().
# For a new problem, replace example_f, example_fy, and example_fz and modify
# the arguments in the call to solve_sbvp.
#
# Example
#
#     epsilon*y'' = -tanh(y') + y - 1,      0 < x < 1,
#     y(0) = 3,  y(1) = 5.
#
# In the notation used by solve_sbvp,
#
#     f(x,y,z,epsilon)  = (-tanh(z) + y - 1)/epsilon
#     fy(x,y,z,epsilon) = 1/epsilon
#     fz(x,y,z,epsilon) = -sech(z)**2/epsilon
#
# The continuation begins at epsilon = 10**e_max and proceeds to
# epsilon = 10**e_min. If Newton's method does not converge, try increasing
# n_epsilon or nx, or use a less negative value of e_min.
# ---------------------------------------------------------------------