"""
ch07_gauss_legendre_2d.py

Two-dimensional integration using a Gauss-Legendre product rule.

The program evaluates

        1   1
        /   /
       |   |        1
   I = |   |   -------------  dy dx
       |   |   1 + x^2 + y^2
       /   /
        0   0

An n-point Gauss-Legendre rule is used in each direction. The
tensor-product rule therefore requires n^2 function evaluations.

Based on the Fortran implementation by Alexander Godunov.
Python version prepared for the companion website, 2026.

The numerical routine uses only the Python standard library.
NumPy and Matplotlib are used only for visualization.
"""

import math
import numpy as np
import matplotlib.pyplot as plt


def f(x, y):
    """Two-dimensional integrand."""
    return 1.0 / (1.0 + x * x + y * y)


def legendre_pair(n, z):
    if n == 1:
        return z, 1.0

    p0 = 1.0
    p1 = z

    for j in range(2, n + 1):
        p2 = ((2.0 * j - 1.0) * z * p1 - (j - 1.0) * p0) / j
        p0 = p1
        p1 = p2

    return p1, p0


def gauss_legendre_nodes_weights(n, a, b):
    """Generate n Gauss-Legendre nodes and weights on [a,b]."""
    tol = 1.0e-14
    max_iter = 100

    if n < 1:
        raise ValueError(
            "gauss_legendre_nodes_weights: n must be positive."
        )

    node = [0.0] * n
    weight = [0.0] * n

    midpoint = 0.5 * (a + b)
    halfwidth = 0.5 * (b - a)
    m = (n + 1) // 2

    for i in range(1, m + 1):
        z = math.cos(math.pi * (i - 0.25) / (n + 0.5))
        converged = False

        for _ in range(max_iter):
            pn, pnm1 = legendre_pair(n, z)
            dp = 1.0 if n == 1 else n * (z * pn - pnm1) / (z * z - 1.0)

            zold = z
            z = zold - pn / dp

            if abs(z - zold) <= tol:
                converged = True
                break

        if not converged:
            raise RuntimeError(
                "Gauss-Legendre root iteration did not converge."
            )

        pn, pnm1 = legendre_pair(n, z)
        dp = 1.0 if n == 1 else n * (z * pn - pnm1) / (z * z - 1.0)
        w = 2.0 / ((1.0 - z * z) * dp * dp)

        left = i - 1
        right = n - i

        node[left] = midpoint - halfwidth * z
        node[right] = midpoint + halfwidth * z

        weight[left] = halfwidth * w
        weight[right] = halfwidth * w

    return node, weight


def gauss_legendre_2d(fun, ax, bx, ay, by, n):
    """Two-dimensional Gauss-Legendre tensor-product rule."""
    xnode, xweight = gauss_legendre_nodes_weights(n, ax, bx)
    ynode, yweight = gauss_legendre_nodes_weights(n, ay, by)

    integral = 0.0
    nfun = 0

    for i in range(n):
        for j in range(n):
            integral += (
                xweight[i]
                * yweight[j]
                * fun(xnode[i], ynode[j])
            )
            nfun += 1

    return integral, nfun


def plot_integrand(ax, bx, ay, by):
    xplot = np.linspace(ax, bx, 120)
    yplot = np.linspace(ay, by, 120)
    xgrid, ygrid = np.meshgrid(xplot, yplot)
    zgrid = 1.0 / (1.0 + xgrid ** 2 + ygrid ** 2)

    fig = plt.figure()
    ax3d = fig.add_subplot(111, projection="3d")
    ax3d.plot_surface(xgrid, ygrid, zgrid, linewidth=0)
    ax3d.set_xlabel("x")
    ax3d.set_ylabel("y")
    ax3d.set_zlabel("f(x,y)")
    ax3d.set_title("Integrand: 1/(1+x^2+y^2)")


def main():
    ax = 0.0
    bx = 1.0
    ay = 0.0
    by = 1.0

    reference = 6.3951035187031102e-01
    nvalues = [2, 4, 8, 16, 32]

    plot_integrand(ax, bx, ay, by)

    print("Two-dimensional Gauss-Legendre integration")
    print("Integral of 1/(1+x^2+y^2) over [0,1] x [0,1]")
    print(f"Reference value      = {reference:15.7e}\n")

    print("        n     Function calls          Integral       Absolute error")
    print("-------------------------------------------------------------------")

    for n in nvalues:
        integral, nfun = gauss_legendre_2d(f, ax, bx, ay, by, n)
        error_value = abs(integral - reference)

        print(
            f"{n:9d}{nfun:19d}"
            f"{integral:20.7e}{error_value:20.7e}"
        )

    plt.show()


if __name__ == "__main__":
    main()