# Solve a linear two-point boundary-value problem by centered differences.
#
# The differential equation is
#
#     y'' + p(x)y' + q(x)y = rhs(x),    xL < x < xR,
#
# with boundary conditions
#
#     a0*y(xL) + b0*y'(xL) = c0,
#     a1*y(xR) + b1*y'(xR) = c1.
#
# The method uses second-order centered differences on a uniform grid and
# solves the resulting tridiagonal linear system with the Thomas algorithm.

from collections.abc import Callable

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


ScalarFunction = Callable[[float], float]
FloatArray = NDArray[np.float64]


def tridiagonal_solve(
    diagonal: FloatArray,
    lower: FloatArray,
    upper: FloatArray,
    right_side: FloatArray,
) -> FloatArray:
    """Solve a tridiagonal linear system using the Thomas algorithm."""

    diagonal = np.asarray(diagonal, dtype=float).copy()
    lower = np.asarray(lower, dtype=float)
    upper = np.asarray(upper, dtype=float)
    right_side = np.asarray(right_side, dtype=float)

    n = right_side.size

    if n == 0:
        raise ValueError("The linear system must contain at least one equation.")

    if not all(array.size == n for array in (diagonal, lower, upper)):
        raise ValueError(
            "All tridiagonal-system arrays must have the same length."
        )

    v = np.zeros(n, dtype=float)
    y = np.zeros(n, dtype=float)

    pivot = diagonal[0]
    if np.isclose(pivot, 0.0):
        raise np.linalg.LinAlgError(
            "A zero pivot occurred in the tridiagonal solver."
        )

    y[0] = right_side[0]/pivot

    # forward sweep
    for i in range(1, n):
        v[i - 1] = upper[i - 1]/pivot
        pivot = diagonal[i] - lower[i]*v[i - 1]

        if np.isclose(pivot, 0.0):
            raise np.linalg.LinAlgError(
                f"A zero pivot occurred at row {i}."
            )

        y[i] = (right_side[i] - lower[i]*y[i - 1])/pivot

    # back substitution
    for i in range(n - 2, -1, -1):
        y[i] = y[i] - v[i]*y[i + 1]

    return y


def lbvp(
    p: ScalarFunction,
    q: ScalarFunction,
    rhs: ScalarFunction,
    xL: float,
    xR: float,
    left_bc: tuple[float, float, float],
    right_bc: tuple[float, float, float],
    nx: int = 100,
) -> tuple[FloatArray, FloatArray]:
    """Solve the linear boundary-value problem on a uniform grid."""

    if nx < 2:
        raise ValueError("nx must be at least 2.")

    if xR <= xL:
        raise ValueError("xR must be greater than xL.")

    a0, b0, c0 = left_bc
    a1, b1, c1 = right_bc

    if a0 == 0.0 and b0 == 0.0:
        raise ValueError("a0 and b0 cannot both be zero.")

    if a1 == 0.0 and b1 == 0.0:
        raise ValueError("a1 and b1 cannot both be zero.")

    x = np.linspace(xL, xR, nx)
    h = x[1] - x[0]
    hh = h*h

    diagonal = np.zeros(nx, dtype=float)
    lower = np.zeros(nx, dtype=float)
    upper = np.zeros(nx, dtype=float)
    forcing = np.zeros(nx, dtype=float)

    # calculate the coefficients at the interior grid points
    for i in range(1, nx - 1):
        diagonal[i] = -2.0 + hh*q(x[i])
        lower[i] = 1.0 - 0.5*h*p(x[i])
        upper[i] = 1.0 + 0.5*h*p(x[i])
        forcing[i] = hh*rhs(x[i])

    # incorporate the left boundary condition
    if b0 == 0.0:
        diagonal[0] = 1.0
        forcing[0] = c0/a0
    else:
        diagonal[0] = (
            -2.0
            + hh*q(x[0])
            + 2.0*h*a0*(1.0 - 0.5*h*p(x[0]))/b0
        )
        upper[0] = 2.0
        forcing[0] = (
            hh*rhs(x[0])
            + 2.0*h*c0*(1.0 - 0.5*h*p(x[0]))/b0
        )

    # incorporate the right boundary condition
    if b1 == 0.0:
        diagonal[nx - 1] = 1.0
        forcing[nx - 1] = c1/a1
    else:
        diagonal[nx - 1] = (
            -2.0
            + hh*q(x[nx - 1])
            - 2.0*h*a1*(1.0 + 0.5*h*p(x[nx - 1]))/b1
        )
        lower[nx - 1] = 2.0
        forcing[nx - 1] = (
            hh*rhs(x[nx - 1])
            - 2.0*h*c1*(1.0 + 0.5*h*p(x[nx - 1]))/b1
        )

    # solve the tridiagonal system
    y = tridiagonal_solve(diagonal, lower, upper, forcing)

    return x, y


# coefficient of y'
def p(x: float) -> float:
    return 100.0*(3.0*x - 1.0)


# coefficient of y
def q(x: float) -> float:
    return 100.0*x


# right hand side of the differential equation
def rhs(x: float) -> float:
    return 0.0


def main() -> None:
    # set the interval
    xL = 0.0
    xR = 1.0

    # set the boundary conditions
    # a0*y(xL) + b0*y'(xL) = c0
    a0 = 1.0
    b0 = 0.0
    c0 = 1.0

    # a1*y(xR) + b1*y'(xR) = c1
    a1 = 1.0
    b1 = 0.0
    c1 = 2.0

    # number of grid points, including the two boundary points
    nx = 100

    x, y = lbvp(
        p=p,
        q=q,
        rhs=rhs,
        xL=xL,
        xR=xR,
        left_bc=(a0, b0, c0),
        right_bc=(a1, b1, c1),
        nx=nx,
    )

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


if __name__ == "__main__":
    main()


# This code is run by entering the command
#
#     python lbvp.py
#
# in a terminal or command window, making sure that this file is in the
# current folder or that its folder is on the Python path. Before running
# the program, the user must enter the problem information in this file.
#
# The code requires NumPy and Matplotlib. If needed, install them using
#
#     python -m pip install numpy matplotlib
#
# Needed Input From User
#
# Differential equation:
#
#     y'' + p(x)y' + q(x)y = rhs(x)
#
# Enter the functions p(x), q(x), and rhs(x) above main().
#
# Interval:
#
#     xL and xR in main()
#
# Boundary conditions:
#
#     a0*y(xL) + b0*y'(xL) = c0
#     a1*y(xR) + b1*y'(xR) = c1
#
# Enter a0, b0, c0, a1, b1, and c1 in main().
#
# Number of grid points:
#
#     nx in main()
#
# Example:
#
#     y'' + 100*(3*x - 1)*y' + 100*x*y = 0,    0 < x < 1,
#
# with
#
#     y(0) = 1,    y(1) = 2.
#
# For this example,
#
#     xL = 0, xR = 1
#     a0 = 1, b0 = 0, c0 = 1
#     a1 = 1, b1 = 0, c1 = 2
#     p(x) = 100*(3*x - 1)
#     q(x) = 100*x
#     rhs(x) = 0
#     nx = 100
#
# Background
#
# This code solves linear second-order two-point boundary-value problems
# using second-order centered differences on a uniform grid. The resulting
# tridiagonal linear system is solved using the Thomas algorithm.
#
# The program supports Dirichlet, Neumann, and Robin boundary conditions,
# provided that a0 and b0 are not both zero and a1 and b1 are not both zero.