//=======================================================================
// Chapter 10, code 1: a linear boundary value problem.
// Translation of ch10_finite_difference.f90.
// Solve y'' + p(x)*y' + q(x)*y = r(x), with y(a)=alpha, y(b)=beta.
// Method: centered finite differences and the Thomas algorithm.
// Example: y'' - y = 0, y(0) = y(1) = 1.
// Exact solution: cosh(x-0.5)/cosh(0.5).
//
// To solve another problem:
//   1. Edit the interval, endpoint values, and grid size in main.
//   2. Edit coefficients(x,p,q,r) immediately below the driver.
//   3. Disable check_exact or replace exact_solution.
// Compile: g++ -std=c++11 -O2 ch10_finite_difference.cpp -o ch10_finite_difference
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

using Coefficients = void (*)(double,double&,double&,double&);
void coefficients(double x, double& p, double& q, double& r);
double exact_solution(double x, double a, double b, double alpha, double beta);
bool finite_difference(Coefficients fcn, int n, const std::vector<double>& x,
                       double alpha, double beta, std::vector<double>& y);
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 grid: edit here -------------------
    const int n = 20;        // n intervals, n+1 points, n-1 interior unknowns
    double a = 0.0;
    double b = 1.0;
    double alpha = 1.0;      // y(a)
    double beta = 1.0;       // y(b)
    bool check_exact = true;
    //-------------------------------------------------------------------
    if (n < 2 || b <= a) {
        std::cerr << "Use n >= 2 and b > a.\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;
    x[n] = b;
    if (!finite_difference(coefficients,n,x,alpha,beta,y)) return 1;

    std::ofstream data("ch10_finite_difference.dat");
    if (!data) {
        std::cerr << "Cannot open ch10_finite_difference.dat.\n";
        return 1;
    }
    data << std::uppercase << std::scientific << std::setprecision(16);
    std::cout << "Linear boundary value problem: finite differences\n"
              << "Grid intervals:" << std::setw(8) << n << '\n'
              << "Grid spacing:  " << std::fixed << std::setprecision(6)
              << std::setw(12) << h << "\n\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' << std::uppercase;
    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],a,b,alpha,beta);
            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';
    if (check_exact)
        std::cout << "Largest grid error: " << std::scientific << std::setprecision(4)
                  << std::setw(12) << max_error << '\n';
    std::cout << "Data file: ch10_finite_difference.dat\n";
    return 0;
}


void coefficients(double x, double& p, double& q, double& r)
{
    // Supply the coefficients in y'' + p(x)*y' + q(x)*y = r(x).
    // Default example: y'' - y = 0. x is available for variable coefficients.
    (void)x;
    p = 0.0;
    q = -1.0;
    r = 0.0;
}


double exact_solution(double x, double a, double b, double alpha, double beta)
{
    // Exact solution of y'' - y = 0 with the supplied endpoint values.
    // Replace this function when changing the equation, or disable the check.
    return (alpha*std::sinh(b-x)+beta*std::sinh(x-a))/std::sinh(b-a);
}


bool finite_difference(Coefficients fcn, int n, const std::vector<double>& x,
                       double alpha, double beta, std::vector<double>& y)
{
    // Centered differences give a tridiagonal system at interior points.
    // After multiplication by h*h, the coefficients of y(i-1), y(i), y(i+1)
    // are 1-h*p/2, -2+h*h*q, 1+h*p/2; the right-hand side is h*h*r.
    // Only interior values are unknown. Move known endpoint contributions
    // to the right-hand side before solving.
    // C++ interior array entry j corresponds to physical grid point i=j+1.
    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), solution(interior);
    for (int j = 0; j < interior; ++j) {
        double p, q, r;
        fcn(x[j+1],p,q,r);
        lower[j] = 1.0-h*p/2.0;
        diagonal[j] = -2.0+h*h*q;
        upper[j] = 1.0+h*p/2.0;
        rhs[j] = h*h*r;
    }
    // Apply both contributions even when there is only one interior point.
    rhs[0] -= lower[0]*alpha;
    rhs[interior-1] -= upper[interior-1]*beta;
    lower[0] = 0.0;
    upper[interior-1] = 0.0;
    if (!thomas(interior,lower,diagonal,upper,rhs,solution)) return false;
    y[0] = alpha;
    for (int j = 0; j < interior; ++j) y[j+1] = solution[j];
    y[n] = beta;
    return true;
}


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;
}