//=======================================================================
// Chapter 10, code 3: a nonlinear boundary value problem.
// Translation of ch10_nonlinear_bvp.f90.
// Solve y'' = f(x,y,y'), with y(a) = alpha and y(b) = beta.
// Method: centered finite differences and Newton iteration.
// Each Newton correction is found with the Thomas algorithm.
//
// Example on [0,1]: y'' = y*y - sin(pi*x)^2 - pi^2*sin(pi*x),
// y(0) = y(1) = 0. Exact solution: y = sin(pi*x).
// The forcing was chosen so that the exact solution is known.
// To solve another problem:
//   1. Edit the interval, boundary values, grid, and initial guess.
//   2. Edit equation below, including its two partial derivatives.
//   3. Disable check_exact or replace exact_solution.
// Compile: g++ -std=c++11 -O2 ch10_nonlinear_bvp.cpp -o ch10_nonlinear_bvp
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

using Equation = void (*)(double,double,double,double&,double&,double&);
void equation(double x, double y, double v, double& f, double& dfdy, double& dfdv);
double exact_solution(double x);
bool newton_bvp(Equation fcn, int n, const std::vector<double>& x,
                std::vector<double>& y, double tolerance, int max_iterations,
                int& iterations, double& residual_norm);
bool thomas(int n, const std::vector<double>& lower, std::vector<double>& diagonal,
            const std::vector<double>& upper, std::vector<double>& rhs,
            std::vector<double>& y);

int main()
{
    //-------------------- Problem and controls: edit here ---------------
    const int n = 20;          // Intervals; n+1 points including endpoints
    double a = 0.0;
    double b = 1.0;
    double alpha = 0.0;        // y(a)
    double beta = 0.0;         // y(b)
    double tolerance = 1.0e-10; // Maximum absolute discrete-equation residual
    int max_iterations = 30;
    bool check_exact = true;
    //-------------------------------------------------------------------
    if (n < 2 || b <= a) {
        std::cerr << "Use n >= 2 and b > a.\n";
        return 1;
    }
    if (tolerance <= 0.0 || max_iterations < 1) {
        std::cerr << "Use positive tolerance and iteration limit.\n";
        return 1;
    }
    double h = (b-a)/n;
    std::vector<double> x(n+1), y(n+1);
    for (int i = 0; i <= n; ++i) {
        x[i] = a+i*h;
        // Initial guess: a straight line between the endpoint values.
        // To supply another guess, change this expression.
        y[i] = alpha+(beta-alpha)*(x[i]-a)/(b-a);
    }
    x[n] = b;
    y[0] = alpha;
    y[n] = beta;
    int iterations = 0;
    double residual_norm = 0.0;
    std::cout << std::uppercase
              << "Nonlinear boundary value problem: finite differences and Newton iteration\n"
              << "Grid intervals:  " << std::setw(8) << n << '\n'
              << "Grid spacing:    " << std::fixed << std::setprecision(6)
              << std::setw(12) << h << '\n'
              << "Newton tolerance:" << std::scientific << std::setprecision(4)
              << std::setw(12) << tolerance << '\n';
    if (!newton_bvp(equation,n,x,y,tolerance,max_iterations,iterations,residual_norm))
        return 1;

    std::ofstream data("ch10_nonlinear_bvp.dat");
    if (!data) {
        std::cerr << "Cannot open ch10_nonlinear_bvp.dat.\n";
        return 1;
    }
    data << std::uppercase << std::scientific << std::setprecision(16);
    std::cout << '\n' << std::setw(10) << "x"
              << ' ' << std::setw(16) << "numerical y";
    if (check_exact) {
        std::cout << ' ' << std::setw(16) << "exact y"
                  << ' ' << std::setw(12) << "abs error";
        data << "# x numerical_y exact_y abs_error\n";
    } else {
        data << "# x numerical_y\n";
    }
    std::cout << '\n';
    double max_error = 0.0;
    for (int i = 0; i <= n; ++i) {
        std::cout << std::fixed << std::setprecision(6) << std::setw(10) << x[i]
                  << std::scientific << std::setprecision(8)
                  << ' ' << std::setw(16) << y[i];
        data << ' ' << std::setw(24) << x[i] << ' ' << std::setw(24) << y[i];
        if (check_exact) {
            double exact = exact_solution(x[i]);
            double error = std::abs(y[i]-exact);
            max_error = std::max(max_error,error);
            std::cout << ' ' << std::setw(16) << exact
                      << std::setprecision(4) << ' ' << std::setw(12) << error;
            data << ' ' << std::setw(24) << exact << ' ' << std::setw(24) << error;
        }
        std::cout << '\n';
        data << '\n';
    }
    data.close();
    std::cout << '\n' << "Newton iterations: " << std::setw(8) << iterations << '\n'
              << "Final residual:    " << std::scientific << std::setprecision(4)
              << std::setw(12) << residual_norm << '\n';
    if (check_exact)
        std::cout << "Largest grid error:" << std::setw(12) << max_error << '\n';
    std::cout << "Data file: ch10_nonlinear_bvp.dat\n";
    return 0;
}


void equation(double x, double y, double v, double& f, double& dfdy, double& dfdv)
{
    // Supply y'' = f(x,y,v), where v = y', and the partial derivatives
    // df/dy and df/dv. Treat x, y, and v as independent variables.
    // For an equation independent of v, set dfdv = 0.
    (void)v;
    const double pi = std::acos(-1.0);
    double s = std::sin(pi*x);
    f = y*y-s*s-pi*pi*s;
    dfdy = 2.0*y;
    dfdv = 0.0;
}


double exact_solution(double x)
{
    // Exact solution for the supplied equation and default boundary values.
    const double pi = std::acos(-1.0);
    return std::sin(pi*x);
}


bool newton_bvp(Equation fcn, int n, const std::vector<double>& x,
                std::vector<double>& y, double tolerance, int max_iterations,
                int& iterations, double& residual_norm)
{
    // At each interior point, define
    // R_i = (y[i-1]-2*y[i]+y[i+1])/h^2 - f(x[i],y[i],v_i),
    // v_i = (y[i+1]-y[i-1])/(2*h).
    // Newton solves J*correction = -R. J is tridiagonal.
    // Multiplication by h*h gives the coefficients
    // lower = 1+h*(df/dv)/2; diagonal = -2-h*h*(df/dy);
    // upper = 1-h*(df/dv)/2; rhs = -h*h*R_i.
    // Endpoint corrections are zero. Only interior values change.
    // Stop on the residual of the DISCRETE equations. Grid refinement
    // tests discretization error. Full Newton steps are used; convergence
    // and the solution reached can depend on the initial guess.
    const int interior = n-1;
    double h = (x[n]-x[0])/n;
    std::vector<double> lower(interior), diagonal(interior), upper(interior);
    std::vector<double> rhs(interior), correction(interior);
    double change = 0.0;
    std::cout << std::setw(10) << "iteration"
              << ' ' << std::setw(18) << "max residual"
              << ' ' << std::setw(18) << "max correction" << '\n';
    for (iterations = 0; iterations <= max_iterations; ++iterations) {
        residual_norm = 0.0;
        for (int i = 1; i < n; ++i) {
            double v = (y[i+1]-y[i-1])/(2.0*h);
            double f, dfdy, dfdv;
            fcn(x[i],y[i],v,f,dfdy,dfdv);
            if (!std::isfinite(f) || !std::isfinite(dfdy) || !std::isfinite(dfdv)) {
                std::cerr << "Nonfinite equation or partial derivative; check the initial guess.\n";
                return false;
            }
            double residual = (y[i-1]-2.0*y[i]+y[i+1])/(h*h)-f;
            residual_norm = std::max(residual_norm,std::abs(residual));
            // Interior array entry j corresponds to physical grid point i=j+1.
            int j = i-1;
            lower[j] = 1.0+h*dfdv/2.0;
            diagonal[j] = -2.0-h*h*dfdy;
            upper[j] = 1.0-h*dfdv/2.0;
            rhs[j] = -h*h*residual;
        }
        std::cout << std::scientific << std::setprecision(10)
                  << std::setw(10) << iterations
                  << ' ' << std::setw(18) << residual_norm
                  << ' ' << std::setw(18) << change << '\n';
        if (residual_norm <= tolerance) return true;
        if (iterations == max_iterations) break;
        lower[0] = 0.0;
        upper[interior-1] = 0.0;
        if (!thomas(interior,lower,diagonal,upper,rhs,correction)) return false;
        change = 0.0;
        for (double value : correction) {
            if (!std::isfinite(value)) {
                std::cerr << "Nonfinite Newton correction; check the initial guess.\n";
                return false;
            }
            change = std::max(change,std::abs(value));
        }
        double scale = 1.0;
        for (double value : y) scale = std::max(scale,std::abs(value));
        if (change <= std::numeric_limits<double>::epsilon()*scale) {
            std::cerr << "Newton iteration cannot advance; check tolerance and initial guess.\n";
            return false;
        }
        for (int j = 0; j < interior; ++j) y[j+1] += correction[j];
    }
    std::cerr << "Newton iteration did not converge within the iteration limit.\n";
    return false;
}


bool thomas(int n, const std::vector<double>& lower, std::vector<double>& diagonal,
            const std::vector<double>& upper, std::vector<double>& rhs,
            std::vector<double>& y)
{
    // Forward elimination and backward substitution.
    // diagonal and rhs are overwritten. No pivoting is used.
    // A near-zero pivot may require a solver with pivoting or indicate a
    // singular boundary value problem.
    for (int i = 0; i < n; ++i) {
        double scale = std::max({std::abs(lower[i]),std::abs(diagonal[i]),std::abs(upper[i])});
        if (std::abs(diagonal[i]) <= std::numeric_limits<double>::epsilon()*scale) {
            std::cerr << "Thomas algorithm: zero or near-zero pivot.\n";
            return false;
        }
        if (i == n-1) break;
        double multiplier = lower[i+1]/diagonal[i];
        diagonal[i+1] -= multiplier*upper[i];
        rhs[i+1] -= multiplier*rhs[i];
    }
    y[n-1] = rhs[n-1]/diagonal[n-1];
    for (int i = n-2; i >= 0; --i)
        y[i] = (rhs[i]-upper[i]*y[i+1])/diagonal[i];
    return true;
}