#!/usr/bin/env python3
"""Chapter 10, code 4: eigenvalues by Numerov shooting and bisection.

Solve y'' = g(x,lambda)*y, with y(a) = y(b) = 0.
Scan lambda for endpoint sign changes, then refine each bracket.
Example: a uniform vibrating string, y'' = -lambda*y on [0,1].
Exact eigenvalues: lambda_m = (m*pi/(b-a))^2, m = 1,2,... .
Normalize each mode so that the integral of y^2 over [a,b] is 1.

Requires NumPy and Matplotlib. Run: python3 ch10_numerov_eigenvalues.py
SciPy alternatives: brentq can refine a sign-changing shooting bracket.
A finite-difference matrix formulation can use eigh_tridiagonal or
sparse eigsh; this changes the discretization from Numerov shooting.

To solve another problem:
  1. Edit the interval, grid, number of modes, and search controls in main.
  2. Edit coefficient immediately below the driver.
  3. Disable check_exact or replace the analytic comparisons in main.
This version assumes zero-value boundaries at BOTH ends and no y' term.
"""

import numpy as np
import matplotlib.pyplot as plt


def main():
    # ----------------- Problem and search: edit here ------------------
    a = 0.0
    b = 1.0
    n = 200                     # Number of grid intervals
    modes = 3
    lambda_min = 0.0
    lambda_max = 100.0
    scan_step = 1.0
    tolerance = 1.0e-10         # Maximum width of the final lambda bracket
    check_exact = True
    make_plot = True
    # Increase lambda_max if the range contains too few modes.
    # Make scan_step small enough to separate consecutive roots.
    # -----------------------------------------------------------------
    if n < 2 or n != int(n) or modes < 1 or modes != int(modes) or b <= a:
        raise ValueError('Use integer n >= 2, integer modes >= 1, and b > a.')
    if lambda_max <= lambda_min or scan_step <= 0 or tolerance <= 0:
        raise ValueError('Use increasing search interval, positive scan step and tolerance.')
    h = (b-a)/n
    x = np.linspace(a, b, n+1)
    print('Eigenvalues: Numerov shooting and bisection')
    print(f'Grid intervals:    {n:8d}')
    print(f'Grid spacing:      {h:12.6f}')
    print(f'Bracket tolerance: {tolerance:12.4e}')
    values, shapes, residuals, iterations = find_modes(
        coefficient, n, x, modes, lambda_min, lambda_max, scan_step, tolerance)

    print()
    if check_exact:
        print(f"{'mode':>6} {'numerical lambda':>18} {'exact lambda':>18} "
              f"{'abs error':>12} {'bisections':>10}")
    else:
        print(f"{'mode':>6} {'numerical lambda':>18} {'bisections':>10}")
    for j in range(modes):
        m = j+1                 # Physical mode numbers start at 1.
        if check_exact:
            exact = (m*np.pi/(b-a))**2
            print(f'{m:6d} {values[j]:18.10f} {exact:18.10f} '
                  f'{abs(values[j]-exact):12.4e} {iterations[j]:10d}')
        else:
            print(f'{m:6d} {values[j]:18.10f} {iterations[j]:10d}')

    # One file for all modes. Each column of shapes is one eigenfunction.
    filename = 'ch10_numerov_eigenvalues.dat'
    shape_errors = np.zeros(modes)
    with open(filename, 'w', encoding='utf-8') as data_file:
        data_file.write('# Numerov eigenfunctions: x y_1 y_2 ...\n')
        data_file.write('# Each mode has integral(y^2 dx) = 1 (trapezoidal rule).\n')
        for j in range(modes):
            data_file.write(f'# mode {j+1:6d} lambda = {values[j]:.16e}\n')
        print('\nEigenfunctions: columns are x, then the normalized modes.')
        for i in range(n+1):
            print(f'{x[i]:10.6f}', end='')
            for j in range(modes):
                print(f' {shapes[i,j]:16.8e}', end='')
                if check_exact:
                    exact_shape = np.sqrt(2/(b-a))*np.sin((j+1)*np.pi*(x[i]-a)/(b-a))
                    shape_errors[j] = max(shape_errors[j], abs(shapes[i,j]-exact_shape))
            print()
            row = [x[i], *shapes[i]]
            data_file.write(' '.join(f'{value:.16e}' for value in row)+'\n')
    print(f"\n{'mode':>6} {'endpoint residual':>18}")
    for j in range(modes):
        print(f'{j+1:6d} {residuals[j]:18.8e}')
    if check_exact:
        print(f"{'mode':>6} {'largest shape error':>20}")
        for j in range(modes):
            print(f'{j+1:6d} {shape_errors[j]:20.8e}')
    print(f'Data file: {filename}')

    if make_plot:
        fig, axes = plt.subplots(modes, 1, figsize=(7, 2.5*modes), squeeze=False)
        x_plot = np.linspace(a, b, 801)
        for j in range(modes):
            ax = axes[j, 0]
            ax.plot(x, shapes[:,j], 'k-', linewidth=1.2, label='Numerov')
            if check_exact:
                exact_shape = np.sqrt(2/(b-a))*np.sin((j+1)*np.pi*(x_plot-a)/(b-a))
                ax.plot(x_plot, exact_shape, 'k--', linewidth=1, label='Exact solution')
                ax.legend()
            ax.set_xlabel('x')
            ax.set_ylabel(r'$y_m(x)$')
            ax.set_title(f'Mode {j+1}: lambda = {values[j]:.8f}')
            ax.grid(True)
        fig.tight_layout()
        plt.show()


def coefficient(x, lam):
    """Supply g(x,lambda) in y'' = g(x,lambda)*y.

    Uniform string: g = -lambda. For variable dimensionless density
    rho(x), use g = -lambda*rho(x). No first-derivative term is allowed.
    """
    return -lam


def find_modes(fcn, n, x, modes, lambda_min, lambda_max, scan_step, tolerance):
    """Scan the endpoint mismatch y(b;lambda) for sign changes.

    A sign scan can miss two or more roots inside one interval. Reduce
    scan_step and repeat when adapting to another equation. Analytic
    mode numbers assume the scan finds the lowest modes in order.
    """
    values = np.zeros(modes)
    shapes = np.zeros((n+1, modes))
    residuals = np.zeros(modes)
    iterations = np.zeros(modes, dtype=int)
    found = 0
    left = lambda_min
    trial, f_left = numerov(fcn, n, x, left)
    # Handle an eigenvalue exactly at the initial scan endpoint once.
    if f_left == 0:
        values[found] = left
        trial = normalize_mode(n, x, trial)
        shapes[:,found] = trial
        residuals[found] = trial[-1]
        found += 1
    while found < modes and left < lambda_max:
        right = min(left+scan_step, lambda_max)
        if right <= left:
            raise RuntimeError('Scan step cannot advance lambda.')
        _, f_right = numerov(fcn, n, x, right)
        bracketed = (f_left < 0 and f_right > 0) or (f_left > 0 and f_right < 0)
        if f_right == 0 or bracketed:
            if f_right == 0:
                root = right
                count = 0
            else:
                root, count = bisect(fcn, n, x, left, right, f_left, tolerance)
            values[found] = root
            iterations[found] = count
            trial, _ = numerov(fcn, n, x, root)
            trial = normalize_mode(n, x, trial)
            shapes[:,found] = trial
            # Keep normalized endpoint residuals; do not force them to zero.
            residuals[found] = trial[-1]
            found += 1
        left = right
        f_left = f_right
    if found < modes:
        raise RuntimeError('Too few modes found: enlarge range or reduce scan_step.')
    return values, shapes, residuals, iterations


def bisect(fcn, n, x, left, right, f_left, tolerance):
    """Refine a bracket; stop on lambda width, not mismatch magnitude.

    Mismatch magnitude depends on the arbitrary shooting amplitude.
    Local copies preserve the scan endpoints.
    """
    lo = left
    hi = right
    flo = f_left
    iterations = 0
    max_iterations = 100
    while hi-lo > tolerance:
        if iterations >= max_iterations:
            raise RuntimeError('Bisection iteration limit reached.')
        mid = lo+(hi-lo)/2
        if mid <= lo or mid >= hi:
            raise RuntimeError('Lambda bracket cannot shrink; check tolerance.')
        _, fmid = numerov(fcn, n, x, mid)
        iterations += 1
        if fmid == 0:
            return mid, iterations
        if (flo < 0 and fmid > 0) or (flo > 0 and fmid < 0):
            hi = mid
        else:
            lo = mid
            flo = fmid
    return lo+(hi-lo)/2, iterations


def numerov(fcn, n, x, lam):
    """Numerov propagation for y'' = g*y, with c = h^2/12:

    (1-c*g_next)*y_next =
        2*(1+5*c*g_center)*y_center - (1-c*g_previous)*y_previous.

    y[0] = 0 and y[1] = h fix an arbitrary nonzero amplitude.
    No Taylor approximation of a derivative is needed: the homogeneous
    recurrence makes eigenvalues independent of amplitude. Normalize
    the completed mode after finding the root. This starting choice
    applies to a zero-value left boundary only.
    """
    h = (x[-1]-x[0])/n
    c = h*h/12
    g = np.zeros(n+1)
    y = np.zeros(n+1)
    for i in range(n+1):
        g[i] = fcn(x[i], lam)
        if not np.isfinite(g[i]):
            raise RuntimeError('Nonfinite Numerov coefficient.')
    y[0] = 0
    y[1] = h
    for i in range(1, n):
        denominator = 1-c*g[i+1]
        if abs(denominator) <= np.finfo(float).eps*max(1, abs(c*g[i+1])):
            raise RuntimeError('Near-zero Numerov denominator: refine the grid.')
        y[i+1] = (2*(1+5*c*g[i])*y[i]-(1-c*g[i-1])*y[i-1])/denominator
        if not np.isfinite(y[i+1]):
            raise RuntimeError('Nonfinite Numerov solution: check equation, interval and grid.')
    return y, y[-1]


def normalize_mode(n, x, y):
    """Trapezoidal integral of y^2; keep the initial sign."""
    h = (x[-1]-x[0])/n
    norm_value = np.sqrt(h*(0.5*y[0]**2+np.sum(y[1:n]**2)+0.5*y[-1]**2))
    if norm_value <= 0 or not np.isfinite(norm_value):
        raise RuntimeError('Cannot normalize the eigenfunction.')
    return y/norm_value


if __name__ == '__main__':
    main()