//=======================================================================
// Chapter 10, code 2: a linear BVP with different boundary conditions.
// Translation of ch10_boundary_conditions.f90.
// Solve y'' + p(x)*y' + q(x)*y = r(x) on a uniform grid.
// Method: centered finite differences and the Thomas algorithm.
// Derivative boundaries use second-order ghost-point elimination.
//
// Example: dimensionless heat conduction, T'' = -sin(pi*x), [0,1].
// Default: T(0) = 0, T'(1) = 0 (an insulated right end).
// y denotes T. Exact solutions are supplied for all four boundary keys.
// For another model, edit the inputs and coefficients, and replace the
// analytic solution or set check_exact = false.
// Compile: g++ -std=c++11 -O2 ch10_boundary_conditions.cpp -o ch10_boundary_conditions
//=======================================================================
#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, int key,
                      double left_value, double right_value, double alpha, double beta);
bool finite_difference(Coefficients fcn, int n, const std::vector<double>& x, int key,
                       double left_value, double right_value, 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;    // Number of intervals; n+1 points including endpoints
    double a = 0.0;
    double b = 1.0;
    int key = 2;
    double left_value = 0.0;
    double right_value = 0.0;
    double robin_alpha = 1.0;
    double robin_beta = 1.0;
    bool check_exact = true;
    // Meaning of key and the two boundary data:
    //   1: y(a)  = left_value; y(b)  = right_value
    //   2: y(a)  = left_value; y'(b) = right_value
    //   3: y'(a) = left_value; y(b)  = right_value
    //   4: y(a)  = left_value;
    //      robin_alpha*y(b) + robin_beta*y'(b) = right_value
    // Derivatives mean dy/dx, not the outward normal derivative.
    // For a rod's convecting right end, key 4 represents
    //   k*T'(b) + H*T(b) = H*T_environment.
    // Use dimensionless counterparts of k, H, and temperature here.
    //-------------------------------------------------------------------
    if (n < 2 || b <= a) {
        std::cerr << "Use n >= 2 and b > a.\n";
        return 1;
    }
    if (key < 1 || key > 4) {
        std::cerr << "Choose a boundary key from 1 to 4.\n";
        return 1;
    }
    if (key == 4) {
        if (robin_beta == 0.0) {
            std::cerr << "Use key 1 for a fixed right-end value.\n";
            return 1;
        }
        if (check_exact && std::abs(robin_alpha*(b-a)+robin_beta)
            <= std::numeric_limits<double>::epsilon()
               *std::max(std::abs(robin_alpha*(b-a)),std::abs(robin_beta))) {
            std::cerr << "The supplied example has no unique solution for these Robin data.\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,key,left_value,right_value,
                           robin_alpha,robin_beta,y)) return 1;

    std::ofstream data("ch10_boundary_conditions.dat");
    if (!data) {
        std::cerr << "Cannot open ch10_boundary_conditions.dat.\n";
        return 1;
    }
    data << std::uppercase << std::scientific << std::setprecision(16);
    std::cout << "Linear boundary value problem: boundary-condition choices\n"
              << "Boundary key:  " << std::setw(8) << key << '\n'
              << "Grid intervals:" << std::setw(8) << n << '\n'
              << "Grid spacing:  " << std::fixed << std::setprecision(6)
              << std::setw(12) << h << '\n';
    switch (key) {
    case 1: std::cout << "Fixed values at both ends.\n"; break;
    case 2: std::cout << "Fixed left value; prescribed right derivative.\n"; break;
    case 3: std::cout << "Prescribed left derivative; fixed right value.\n"; break;
    case 4: std::cout << "Fixed left value; right Robin condition.\n"; break;
    }
    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' << 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,key,left_value,right_value,
                                          robin_alpha,robin_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_boundary_conditions.dat\n";
    return 0;
}


void coefficients(double x, double& p, double& q, double& r)
{
    // Supply y'' + p(x)*y' + q(x)*y = r(x).
    // Here y is dimensionless temperature in a uniformly conducting rod.
    const double pi = std::acos(-1.0);
    p = 0.0;
    q = 0.0;
    r = -std::sin(pi*x);
}


double exact_solution(double x, double a, double b, int key,
                      double left_value, double right_value, double alpha, double beta)
{
    // y'' = -sin(pi*x) gives y = sin(pi*x)/pi^2 + c*(x-a) + d.
    // Determine c and d from the selected boundary conditions.
    // Default key 2 on [0,1]: y = sin(pi*x)/pi^2 + x/pi.
    const double pi = std::acos(-1.0);
    double pa = std::sin(pi*a)/(pi*pi);
    double pb = std::sin(pi*b)/(pi*pi);
    double c = 0.0, d = 0.0;
    switch (key) {
    case 1:
        d = left_value-pa;
        c = (right_value-pb-d)/(b-a);
        break;
    case 2:
        d = left_value-pa;
        c = right_value-std::cos(pi*b)/pi;
        break;
    case 3:
        c = left_value-std::cos(pi*a)/pi;
        d = right_value-pb-c*(b-a);
        break;
    case 4:
        d = left_value-pa;
        c = (right_value-alpha*(pb+d)-beta*std::cos(pi*b)/pi)/(alpha*(b-a)+beta);
        break;
    }
    return std::sin(pi*x)/(pi*pi)+c*(x-a)+d;
}


bool finite_difference(Coefficients fcn, int n, const std::vector<double>& x, int key,
                       double left_value, double right_value, double alpha, double beta,
                       std::vector<double>& y)
{
    // Interior equations are identical to Code 1. All n+1 values, including
    // endpoints, are now unknowns in the tridiagonal system.
    // A fixed endpoint gets a row containing just y = prescribed value.
    // For a derivative boundary eliminate a ghost point using
    //   y(-1) = y(1)-2*h*y'(a), or y(n+1) = y(n-1)+2*h*y'(b).
    // The ghost values are algebraic devices: no out-of-range array access.
    // For Robin data, substitute y'(b) = (right_value-alpha*y(n))/beta.
    double h = (x[n]-x[0])/n;
    std::vector<double> lower(n+1,0.0), diagonal(n+1), upper(n+1,0.0), rhs(n+1);
    for (int i = 1; i < n; ++i) {
        double p, q, r;
        fcn(x[i],p,q,r);
        lower[i] = 1.0-h*p/2.0;
        diagonal[i] = -2.0+h*h*q;
        upper[i] = 1.0+h*p/2.0;
        rhs[i] = h*h*r;
    }
    // Left boundary: derivative for key 3, fixed value otherwise.
    if (key == 3) {
        double p, q, r;
        fcn(x[0],p,q,r);
        diagonal[0] = -2.0+h*h*q;
        upper[0] = 2.0;
        rhs[0] = h*h*r+(2.0*h-h*h*p)*left_value;
    } else {
        diagonal[0] = 1.0;
        rhs[0] = left_value;
    }
    // Right boundary: fixed value, derivative, or Robin condition.
    switch (key) {
    case 1:
    case 3:
        diagonal[n] = 1.0;
        rhs[n] = right_value;
        break;
    case 2: {
        double p, q, r;
        fcn(x[n],p,q,r);
        lower[n] = 2.0;
        diagonal[n] = -2.0+h*h*q;
        rhs[n] = h*h*r-(2.0*h+h*h*p)*right_value;
        break;
    }
    case 4: {
        double p, q, r;
        fcn(x[n],p,q,r);
        double factor = (2.0*h+h*h*p)/beta;
        lower[n] = 2.0;
        diagonal[n] = -2.0+h*h*q-factor*alpha;
        rhs[n] = h*h*r-factor*right_value;
        break;
    }
    }
    return thomas(n+1,lower,diagonal,upper,rhs,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)
{
    // 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;
}