"""ch11_newton.py

Find one root of f(x) = 0 by Newton's method.

Origin:
    Python version of Newton1.f90 by Alexander Godunov (January 2010).
    The Newton update and residual stopping test match ch11_newton.f90.
    The old large-step substitution is omitted because it was neither
    a Newton step nor a bisection step.

Example:
    f(x) = x + cos(x), starting from x0 = -1.
    Reference root = -0.73908513321516... .

The program reports the iteration count, root estimate, and f(root).
Newton's method is local: a different initial guess may converge to
another root or fail. Interpret a small residual in the context of
the function and the physical problem.
"""

import math


def main():
    """Set up the example and print its result."""
    # Initial guess and requested residual tolerance
    x0 = -1.0
    eps = 1.0e-7

    print("A single root of f(x)")
    print("Method - Newton")
    print(f"Initial guess = {x0:.3f}")

    root, flag = newton1(f, x0, eps)

    # Report invalid input or a zero derivative before printing a root
    if flag == 0:
        print("invalid tolerance or zero derivative at an iterate")
        return

    print(f"residual tolerance = {eps:.5e}")
    print(f"iterations         = {abs(flag)}")
    if flag < 0:
        print("iteration limit reached before the requested residual")

    print(f"root               = {root:.5e}")
    # Check the residual of the reported estimate
    print(f"f(root)            = {f(root, 0):.5e}")


def newton1(fun, x0, eps):
    """Locate one root near x0 using Newton's update x_new=x-f(x)/f'(x).

    Input:
        fun  - function evaluating f(x) for k=0 and f'(x) for k=1
        x0   - initial guess
        eps  - requested maximum residual abs(f(root)) (positive)

    Output:
        root - most recent Newton estimate when flag is nonzero
        flag - convergence indicator and iteration count
               flag > 0: abs(f(root)) <= eps; flag is the count
               flag = 0: invalid tolerance or zero derivative
               flag < 0: iteration limit reached; abs(flag) is the count

    Notes:
        This is the ordinary Newton method without damping or bracketing.
        The maximum number of iterations is 200. If the initial guess
        already satisfies the residual test, flag=1 denotes success even
        though no Newton update was needed.
    """
    max_iterations = 200
    if eps <= 0.0:
        return math.nan, 0
    x = x0

    # A sufficiently accurate initial guess needs no Newton update.
    fx = fun(x, 0)
    if abs(fx) <= eps:
        return x, 1

    # Repeatedly use the value and derivative at the current estimate.
    for i in range(1, max_iterations + 1):
        dfx = fun(x, 1)
        if dfx == 0.0:
            return math.nan, 0

        xnew = x - fx / dfx
        fx = fun(xnew, 0)

        # Stop on the residual, as in the original Fortran program.
        if abs(fx) <= eps:
            return xnew, i
        x = xnew

    return xnew, -max_iterations


def f(x, k):
    """Evaluate the example function for k=0 and its derivative for k=1."""
    # To use another function, change both expressions consistently.
    if k == 0:
        return x + math.cos(x)
    if k == 1:
        return 1.0 - math.sin(x)
    raise ValueError("The function selector k must be 0 or 1.")


if __name__ == "__main__":
    main()