"""
ch07_oscillatory.py

Numerical integration of a rapidly oscillating function.

The program evaluates

        1
        /
       |  exp(-x) cos(omega*x) dx,       omega = 50,
       /
       0

using:
    1. Composite Simpson rule
    2. Composite quadratic Filon rule

Simpson resolves the complete oscillatory integrand numerically.
Filon approximates only the slowly varying amplitude by a quadratic
polynomial and integrates the oscillatory factor analytically.

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

Matplotlib is used only for visualization.
"""

import math
import matplotlib.pyplot as plt


def f(x):
    """Slowly varying amplitude."""
    return math.exp(-x)


def composite_simpson_oscillatory(fun, a, b, omega, n):
    """Composite Simpson rule for f(x) cos(omega*x)."""
    if n < 2 or n % 2 != 0:
        raise ValueError(
            "composite_simpson_oscillatory: n must be even."
        )

    h = (b - a) / n

    sum_odd = 0.0
    for i in range(1, n, 2):
        x = a + i * h
        sum_odd += fun(x) * math.cos(omega * x)

    sum_even = 0.0
    for i in range(2, n - 1, 2):
        x = a + i * h
        sum_even += fun(x) * math.cos(omega * x)

    return h * (
        fun(a) * math.cos(omega * a)
        + fun(b) * math.cos(omega * b)
        + 4.0 * sum_odd
        + 2.0 * sum_even
    ) / 3.0


def filon_quadratic(fun, a, b, omega, n):
    """Composite quadratic Filon rule."""
    if n < 2 or n % 2 != 0:
        raise ValueError(
            "filon_quadratic: n must be a positive even integer."
        )

    if abs(omega) < 1.0e-12:
        raise ValueError("filon_quadratic: omega is too close to zero.")

    h = (b - a) / n
    theta = omega * h

    c0 = 2.0 * math.sin(theta) / omega

    c2 = (
        2.0 * h * h * math.sin(theta) / omega
        + 4.0 * h * math.cos(theta) / (omega * omega)
        - 4.0 * math.sin(theta) / (omega ** 3)
    )

    s1 = (
        -2.0 * h * math.cos(theta) / omega
        + 2.0 * math.sin(theta) / (omega * omega)
    )

    integral = 0.0

    for k in range(0, n - 1, 2):
        x0 = a + k * h
        x1 = x0 + h
        x2 = x0 + 2.0 * h

        f0 = fun(x0)
        f1 = fun(x1)
        f2 = fun(x2)

        a0 = f1
        a1 = (f2 - f0) / (2.0 * h)
        a2 = (f0 - 2.0 * f1 + f2) / (2.0 * h * h)

        pair_integral = (
            math.cos(omega * x1) * (a0 * c0 + a2 * c2)
            - math.sin(omega * x1) * a1 * s1
        )

        integral += pair_integral

    return integral


def plot_integrand(a, b, omega):
    nplot = 4000
    xplot = [a + i * (b - a) / nplot for i in range(nplot + 1)]
    amplitude = [f(x) for x in xplot]
    oscillatory = [
        f(x) * math.cos(omega * x)
        for x in xplot
    ]

    plt.figure()
    plt.plot(xplot, oscillatory, linewidth=1.2, label="exp(-x) cos(50x)")
    plt.plot(xplot, amplitude, "--", linewidth=1.0, label="+ exp(-x)")
    plt.plot(xplot, [-y for y in amplitude], "--",
             linewidth=1.0, label="- exp(-x)")
    plt.xlabel("x")
    plt.ylabel("f(x) cos(omega x)")
    plt.title("Rapidly oscillating integrand and amplitude envelope")
    plt.legend()
    plt.grid(True)


def main():
    a = 0.0
    b = 1.0
    omega = 50.0

    exact = (
        math.exp(-b) * (-math.cos(omega * b) + omega * math.sin(omega * b))
        - math.exp(-a) * (-math.cos(omega * a) + omega * math.sin(omega * a))
    ) / (1.0 + omega * omega)

    n_simpson = [100, 200, 400, 800, 1600]
    n_filon = [20, 40, 80]

    plot_integrand(a, b, omega)

    print("Rapidly oscillating integral")
    print(f"Integral of exp(-x) cos(omega*x), omega = {omega:6.1f}")
    print(f"Exact value          = {exact:15.7e}\n")

    print("Composite Simpson rule")
    print("        n            Integral       Absolute error")
    print("----------------------------------------------------")

    for n in n_simpson:
        integral = composite_simpson_oscillatory(f, a, b, omega, n)
        error_value = abs(integral - exact)
        print(f"{n:9d}{integral:20.7e}{error_value:20.7e}")

    print("\nComposite quadratic Filon rule")
    print("        n            Integral       Absolute error")
    print("----------------------------------------------------")

    for n in n_filon:
        integral = filon_quadratic(f, a, b, omega, n)
        error_value = abs(integral - exact)
        print(f"{n:9d}{integral:20.7e}{error_value:20.7e}")

    plt.show()


if __name__ == "__main__":
    main()