"""
ch07_adaptive_gauss_kronrod.py

Adaptive numerical integration using embedded Gauss-Kronrod rules.

Two quadrature pairs are included:
    1. Gauss 7-point / Kronrod 15-point (G7-K15)
    2. Gauss 10-point / Kronrod 21-point (G10-K21)

The same non-recursive adaptive subdivision algorithm is used with
either pair.

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

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

Based on integration codes 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 g7k15(fun, a, b):
    """Gauss 7-point / Kronrod 15-point pair on [a,b]."""
    x = [
        0.000000000000000,
        0.207784955007898,
        0.405845151377397,
        0.586087235467691,
        0.741531185599394,
        0.864864423359769,
        0.949107912342758,
        0.991455371120813,
    ]

    wg = [
        0.417959183673469,
        0.381830050505119,
        0.279705391489277,
        0.129484966168870,
    ]

    wk = [
        0.209482141084728,
        0.204432940075298,
        0.190350578064785,
        0.169004726639267,
        0.140653259715525,
        0.104790010322250,
        0.063092092629979,
        0.022935322010529,
    ]

    h = (b - a) / 2.0
    c = (b + a) / 2.0

    f0 = fun(c)
    fl = [0.0] * 8
    fr = [0.0] * 8

    for i in range(1, 8):
        fl[i] = fun(c - h * x[i])
        fr[i] = fun(c + h * x[i])

    g7 = wg[0] * f0
    for i in range(1, 4):
        g7 += wg[i] * (fl[2 * i] + fr[2 * i])
    g7 *= h

    k15 = wk[0] * f0
    for i in range(1, 8):
        k15 += wk[i] * (fl[i] + fr[i])
    k15 *= h

    return g7, k15


def g10k21(fun, a, b):
    """Gauss 10-point / Kronrod 21-point pair on [a,b]."""
    x = [
        0.0000000000000000000000000,
        0.1488743389816312108848260,
        0.2943928627014601981311266,
        0.4333953941292471907992659,
        0.5627571346686046833390001,
        0.6794095682990244062343274,
        0.7808177265864168970637176,
        0.8650633666889845107320967,
        0.9301574913557082260012072,
        0.9739065285171717200779640,
        0.9956571630258080807355273,
    ]

    wk = [
        0.1494455540029169056649365,
        0.1477391049013384913748415,
        0.1427759385770600807970943,
        0.1347092173114733259280540,
        0.1234919762620658510779581,
        0.1093871588022976418992106,
        0.0931254545836976055350655,
        0.0750396748109199527670431,
        0.0547558965743519960313813,
        0.0325581623079647274788190,
        0.0116946388673718742780644,
    ]

    wg = [
        0.2955242247147528701738930,
        0.2692667193099963550912269,
        0.2190863625159820439955349,
        0.1494513491505805931457763,
        0.0666713443086881375935688,
    ]

    h = (b - a) / 2.0
    c = (b + a) / 2.0

    f0 = fun(c)
    fl = [0.0] * 11
    fr = [0.0] * 11

    for i in range(1, 11):
        fl[i] = fun(c - h * x[i])
        fr[i] = fun(c + h * x[i])

    g10 = 0.0
    for i in range(5):
        index = 2 * i + 1
        g10 += wg[i] * (fl[index] + fr[index])
    g10 *= h

    k21 = wk[0] * f0
    for i in range(1, 11):
        k21 += wk[i] * (fl[i] + fr[i])
    k21 *= h

    return g10, k21


def gkinteg(fun, a, b, eps, key):
    """
    Adaptive non-recursive Gauss-Kronrod integration.

    key = 1  -> G7-K15
    key = 2  -> G10-K21
    """
    im = 32

    x = [0.0] * im
    h = [0.0] * im
    tol = [0.0] * im
    level = [0] * im

    result = 0.0
    errest = 0.0
    nfun = 0

    stack_size = 1
    x[0] = a
    h[0] = b - a
    tol[0] = eps
    level[0] = 1

    while stack_size > 0:
        i = stack_size - 1

        if key == 1:
            gauss, kronrod = g7k15(fun, x[i], x[i] + h[i])
            nfun += 15
        elif key == 2:
            gauss, kronrod = g10k21(fun, x[i], x[i] + h[i])
            nfun += 21
        else:
            raise ValueError("gkinteg: key must be 1 or 2.")

        x0 = x[i]
        step = h[i]
        err = tol[i]
        deep = level[i]

        stack_size -= 1

        difference = abs(kronrod - gauss)

        if (200.0 * difference) ** 1.5 <= err:
            result += kronrod
            errest += difference
        else:
            if deep >= im:
                raise RuntimeError(
                    "gkinteg: maximum subdivision depth reached."
                )

            # Push right interval.
            i = stack_size
            x[i] = x0 + step / 2.0
            h[i] = step / 2.0
            tol[i] = err / 2.0
            level[i] = deep + 1
            stack_size += 1

            # Push left interval.
            i = stack_size
            x[i] = x0
            h[i] = step / 2.0
            tol[i] = err / 2.0
            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 for adaptive Gauss-Kronrod integration")
    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)

    print("Adaptive Gauss-Kronrod integration")
    print("Integral of 1/[1+400(x-0.2)^2] from 0 to 1")
    print(f"Exact value          = {exact:18.10e}\n")

    result, errest, nfun = gkinteg(f, a, b, eps, 1)
    print(f"G7-K15 result        = {result:18.10e}")
    print(f"Estimated error      = {errest:12.4e}")
    print(f"Function calls       = {nfun:10d}\n")

    result, errest, nfun = gkinteg(f, a, b, eps, 2)
    print(f"G10-K21 result       = {result:18.10e}")
    print(f"Estimated error      = {errest:12.4e}")
    print(f"Function calls       = {nfun:10d}")

    plt.show()


if __name__ == "__main__":
    main()