"""
ch07_adaptive_simpson.py

Adaptive numerical integration using Simpson's rule.

The routine compares Simpson estimates on one interval and on its two
equal subintervals. Intervals that do not satisfy the local error test
are subdivided further.

The algorithm is non-recursive: interval data are stored explicitly
in arrays that act as a stack.

Example:
    Integral of 1/[1+400(x-0.2)^2] from 0 to 1.

Exact value:
    [atan(16) + atan(4)] / 20

Original adaptive Simpson routine written by Alexander Godunov,
July 2012. Python version prepared for the companion website, 2026.

Matplotlib is used only for visualization.
"""

import math
import matplotlib.pyplot as plt


def f(x):
    """Narrow-peak test function."""
    return 1.0 / (1.0 + 400.0 * (x - 0.2) ** 2)


def adaptive_simpson(fun, a, b, eps):
    """
    Adaptive non-recursive Simpson integration.

    Returns
    -------
    result : numerical approximation
    errest : accumulated estimate of absolute error
    nfun   : number of function evaluations
    """
    im = 32

    tol = [0.0] * im
    x = [0.0] * im
    h = [0.0] * im
    fa = [0.0] * im
    fm = [0.0] * im
    fb = [0.0] * im
    s = [0.0] * im
    level = [0] * im

    result = 0.0
    errest = 0.0
    i = 0

    x[i] = a
    h[i] = (b - a) / 2.0
    fa[i] = fun(a)
    fm[i] = fun(a + h[i])
    fb[i] = fun(b)
    tol[i] = 15.0 * eps
    level[i] = 1
    s[i] = h[i] * (fa[i] + 4.0 * fm[i] + fb[i]) / 3.0
    nfun = 3

    stack_size = 1

    while stack_size > 0:
        i = stack_size - 1

        f1 = fun(x[i] + h[i] / 2.0)
        f3 = fun(x[i] + 3.0 * h[i] / 2.0)
        nfun += 2

        s1 = h[i] * (fa[i] + 4.0 * f1 + fm[i]) / 6.0
        s2 = h[i] * (fm[i] + 4.0 * f3 + fb[i]) / 6.0

        x0 = x[i]
        f0 = fa[i]
        f2 = fm[i]
        f4 = fb[i]
        step = h[i]
        err = tol[i]
        s0 = s[i]
        deep = level[i]

        stack_size -= 1

        if abs(s1 + s2 - s0) <= err:
            result += s1 + s2
            errest += abs(s1 + s2 - s0) / 15.0
        else:
            if deep >= im:
                raise RuntimeError(
                    "adaptive_simpson: maximum subdivision depth reached."
                )

            # Push right subinterval.
            i = stack_size
            x[i] = x0 + step
            fa[i] = f2
            fm[i] = f3
            fb[i] = f4
            h[i] = step / 2.0
            tol[i] = err / 2.0
            s[i] = s2
            level[i] = deep + 1
            stack_size += 1

            # Push left subinterval.
            i = stack_size
            x[i] = x0
            fa[i] = f0
            fm[i] = f1
            fb[i] = f2
            h[i] = step / 2.0
            tol[i] = err / 2.0
            s[i] = s1
            level[i] = deep + 1
            stack_size += 1

    return result, errest, nfun


def plot_integrand(a, b):
    nplot = 2000
    xplot = [a + i * (b - a) / nplot for i in range(nplot + 1)]
    yplot = [f(x) for x in xplot]

    plt.figure()
    plt.plot(xplot, yplot, linewidth=1.5)
    plt.xlabel("x")
    plt.ylabel("f(x)")
    plt.title("Narrow peak: 1/[1+400(x-0.2)^2]")
    plt.grid(True)


def main():
    a = 0.0
    b = 1.0
    eps = 1.0e-8
    exact = (math.atan(16.0) + math.atan(4.0)) / 20.0

    plot_integrand(a, b)

    result, errest, nfun = adaptive_simpson(f, a, b, eps)

    print("Adaptive Simpson integration")
    print("Integral of 1/[1+400(x-0.2)^2] from 0 to 1")
    print(f"Exact value          = {exact:15.7e}")
    print(f"Calculated value     = {result:15.7e}")
    print(f"Estimated error      = {errest:12.4e}")
    print(f"Function calls       = {nfun:10d}")

    plt.show()


if __name__ == "__main__":
    main()