//=======================================================================
// Chapter 10, code 0: shooting for a second-order boundary value problem.
// Translation of ch10_shooting.f90.
// Solve y'' = f(x,y,y'), with y(a) = alpha and y(b) = beta.
// Secant iteration finds the unknown initial slope. Each shooting trial
// is an initial value problem, integrated with adaptive Fehlberg RKF45.
//
// Example: y'' = y, y(0) = y(1) = 1.
// Exact solution: cosh(x-0.5)/cosh(0.5).
// To solve another problem:
//   1. Edit the interval, boundary values, and slope guesses in main.
//   2. Edit f(x,y,dy) immediately below the driver.
//   3. Disable check_exact or replace exact_solution.
// C++ indices start at zero: state[0] = y, state[1] = y'.
// Compile: g++ -std=c++11 -O2 ch10_shooting.cpp -o ch10_shooting
//=======================================================================
#include <algorithm>
#include <array>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

using State = std::array<double,2>;
using Equation = double (*)(double,double,double);

double f(double x, double y, double dy);
double exact_solution(double x);
bool shoot2(Equation fcn, double a, double b, double alpha, double beta,
            double slope1, double slope2, double h_first, double h_max,
            double ivp_tolerance, double boundary_tolerance, int capacity,
            std::vector<double>& grid, std::vector<State>& solution,
            double& slope, double& residual, int& trials, int& accepted, int& rejected);
bool equations(Equation fcn, double x, const State& y, State& dy);
bool rkf45_system(Equation fcn, double t0, const State& x0, double t_end,
                   double h_first, double h_max, double tolerance, int capacity,
                   std::vector<double>& grid, std::vector<State>& solution,
                   State& xf, int& accepted, int& rejected);

int main()
{
    //-------------------- Problem and controls: edit here ---------------
    const int capacity = 100001;
    double a = 0.0;
    double b = 1.0;
    double alpha = 1.0;                 // y(a)
    double beta = 1.0;                  // y(b)
    double slope1 = 0.0;                // First guess for y'(a)
    double slope2 = 1.0;                // Second guess for y'(a)
    double h_first = 0.05;
    double h_max = 0.05;                // Maximum spacing between output points
    double ivp_tolerance = 1.0e-10;      // Absolute LOCAL tolerance for RKF45
    double boundary_tolerance = 1.0e-8;  // Required |y(b)-beta|
    bool check_exact = true;
    //-------------------------------------------------------------------
    // Vectors allocate trajectory storage on the heap.
    std::vector<double> x(capacity);
    std::vector<State> y(capacity);
    double slope = 0.0, residual = 0.0;
    int trials = 0, accepted = 0, rejected = 0;

    std::cout << std::uppercase << std::scientific << std::setprecision(4)
              << "Boundary value problem: shooting with secant iteration and RKF45\n"
              << "RKF45 tolerance:   " << std::setw(12) << ivp_tolerance << '\n'
              << "Boundary tolerance:" << std::setw(12) << boundary_tolerance << '\n';
    if (!shoot2(f,a,b,alpha,beta,slope1,slope2,h_first,h_max,
                ivp_tolerance,boundary_tolerance,capacity,x,y,
                slope,residual,trials,accepted,rejected)) return 1;

    // Print and save only the converged trajectory.
    std::ofstream data("ch10_shooting.dat");
    if (!data) {
        std::cerr << "Cannot open ch10_shooting.dat.\n";
        return 1;
    }
    data << std::uppercase << std::scientific << std::setprecision(16);
    std::cout << '\n' << std::setw(10) << "x"
              << ' ' << std::setw(16) << "numerical y"
              << ' ' << std::setw(16) << "dy/dx";
    if (check_exact) {
        std::cout << ' ' << std::setw(16) << "exact y"
                  << ' ' << std::setw(12) << "abs error";
        data << "# x y dy_dx exact_y abs_error\n";
    } else {
        data << "# x y dy_dx\n";
    }
    std::cout << '\n';
    double max_error = 0.0;
    for (int i = 0; i <= accepted; ++i) {
        std::cout << std::fixed << std::setprecision(6) << std::setw(10) << x[i]
                  << std::scientific << std::setprecision(8)
                  << ' ' << std::setw(16) << y[i][0]
                  << ' ' << std::setw(16) << y[i][1];
        data << ' ' << std::setw(24) << x[i]
             << ' ' << std::setw(24) << y[i][0]
             << ' ' << std::setw(24) << y[i][1];
        if (check_exact) {
            double exact = exact_solution(x[i]);
            double error = std::abs(y[i][0]-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' << "Shooting trials:      " << std::setw(8) << trials << '\n'
              << "Initial slope:        " << std::fixed << std::setprecision(10)
              << std::setw(16) << slope << '\n'
              << "Boundary residual:    " << std::scientific << std::setprecision(6)
              << std::setw(16) << residual << '\n'
              << "Accepted steps:       " << std::setw(8) << accepted << '\n'
              << "Rejected trials:      " << std::setw(8) << rejected << '\n';
    if (check_exact)
        std::cout << "Largest sampled error:" << std::setw(16) << max_error << '\n';
    std::cout << "Step counts refer to the final shooting trial.\n"
              << "Data file: ch10_shooting.dat\n";
    return 0;
}


double f(double x, double y, double dy)
{
    // The original second-order equation: y'' = f(x,y,y').
    // x and dy are available for equations that depend on them.
    (void)x;
    (void)dy;
    return y;
}


double exact_solution(double x)
{
    // Analytic solution for the supplied example and boundary values.
    return std::cosh(x-0.5)/std::cosh(0.5);
}


bool shoot2(Equation fcn, double a, double b, double alpha, double beta,
            double slope1, double slope2, double h_first, double h_max,
            double ivp_tolerance, double boundary_tolerance, int capacity,
            std::vector<double>& grid, std::vector<State>& solution,
            double& slope, double& residual, int& trials, int& accepted, int& rejected)
{
    // Treat s = y'(a) as an unknown. Each trial starts with y(a) = alpha,
    // y'(a) = s. Its mismatch R(s) = y(b;s)-beta must vanish.
    // Secant iteration is not bracketed; convergence depends on the guesses.
    // A small boundary residual alone does not establish solution accuracy:
    // tighten the RKF45 tolerance and compare trajectories for a new problem.
    const int max_trials = 30;
    if (boundary_tolerance <= 0.0 || slope1 == slope2) {
        std::cerr << "Use positive boundary tolerance and two distinct slope guesses.\n";
        return false;
    }
    std::cout << std::setw(8) << "trial"
              << ' ' << std::setw(18) << "initial slope"
              << ' ' << std::setw(18) << "boundary residual" << '\n';
    slope = slope1;
    double previous_slope = slope1, previous_residual = 0.0;
    for (trials = 1; trials <= max_trials; ++trials) {
        State initial = {{alpha,slope}}, final;
        if (!rkf45_system(fcn,a,initial,b,h_first,h_max,ivp_tolerance,
                           capacity,grid,solution,final,accepted,rejected)) return false;
        residual = final[0]-beta;
        std::cout << std::scientific << std::setprecision(10)
                  << std::setw(8) << trials
                  << ' ' << std::setw(18) << slope
                  << ' ' << std::setw(18) << residual << '\n';
        if (std::abs(residual) <= boundary_tolerance) return true;
        if (trials == max_trials) break;
        double next_slope;
        if (trials == 1) {
            next_slope = slope2;
        } else {
            double denominator = residual-previous_residual;
            double scale = std::max({std::abs(residual),std::abs(previous_residual),
                                     std::numeric_limits<double>::min()});
            if (std::abs(denominator) <= std::numeric_limits<double>::epsilon()*scale) {
                std::cerr << "Secant residuals are indistinguishable; change the guesses.\n";
                return false;
            }
            next_slope = slope-residual*(slope-previous_slope)/denominator;
            if (!std::isfinite(next_slope)) {
                std::cerr << "Nonfinite secant slope.\n";
                return false;
            }
            if (next_slope == slope) {
                std::cerr << "Secant slope cannot advance; check tolerances.\n";
                return false;
            }
        }
        previous_slope = slope;
        previous_residual = residual;
        slope = next_slope;
    }
    std::cerr << "Shooting did not converge within the trial limit.\n";
    return false;
}


bool equations(Equation fcn, double x, const State& y, State& dy)
{
    // Two first-order equations: y[0]' = y[1], y[1]' = f(x,y[0],y[1]).
    dy[0] = y[1];
    dy[1] = fcn(x,y[0],y[1]);
    if (!std::isfinite(dy[0]) || !std::isfinite(dy[1])) {
        std::cerr << "Nonfinite derivative in trial integration.\n";
        return false;
    }
    return true;
}


bool rkf45_system(Equation fcn, double t0, const State& x0, double t_end,
                   double h_first, double h_max, double tolerance, int capacity,
                   std::vector<double>& grid, std::vector<State>& solution,
                   State& xf, int& accepted, int& rejected)
{
    // Adaptive Fehlberg RKF45. All step decisions are inside this routine.
    // Save initial and accepted points; trial trajectories produce no output.
    // Advance with the fifth-order estimate. The largest absolute embedded
    // component error controls acceptance. One absolute tolerance assumes
    // components of comparable magnitudes and controls LOCAL error.
    const int max_attempts = 100000;
    if (capacity < 2 || t_end <= t0 || h_first <= 0.0 || h_max <= 0.0 || tolerance <= 0.0) {
        std::cerr << "Use capacity >= 2, b > a, positive steps, and positive tolerance.\n";
        return false;
    }
    double t = t0, h = h_first;
    State x = x0, work, dx, trial, k1, k2, k3, k4, k5, k6;
    accepted = 0;
    rejected = 0;
    int attempts = 0;
    grid[0] = t;
    solution[0] = x;
    while (t < t_end) {
        if (attempts >= max_attempts) {
            std::cerr << "Maximum number of trials exceeded.\n";
            return false;
        }
        h = std::min({h,h_max,t_end-t});
        double tf = t+h;
        if (tf <= t) {
            std::cerr << "Step too small to advance the independent variable.\n";
            return false;
        }
        ++attempts;
        // Every k includes h. All equations use the SAME intermediate state.
        if (!equations(fcn,t,x,dx)) return false;
        for (int j = 0; j < 2; ++j) {
            k1[j] = h*dx[j];
            work[j] = x[j]+k1[j]/4.0;
        }
        if (!equations(fcn,t+h/4.0,work,dx)) return false;
        for (int j = 0; j < 2; ++j) {
            k2[j] = h*dx[j];
            work[j] = x[j]+3.0*k1[j]/32.0+9.0*k2[j]/32.0;
        }
        if (!equations(fcn,t+3.0*h/8.0,work,dx)) return false;
        for (int j = 0; j < 2; ++j) {
            k3[j] = h*dx[j];
            work[j] = x[j]+1932.0*k1[j]/2197.0
                      -7200.0*k2[j]/2197.0+7296.0*k3[j]/2197.0;
        }
        if (!equations(fcn,t+12.0*h/13.0,work,dx)) return false;
        for (int j = 0; j < 2; ++j) {
            k4[j] = h*dx[j];
            work[j] = x[j]+439.0*k1[j]/216.0-8.0*k2[j]
                      +3680.0*k3[j]/513.0-845.0*k4[j]/4104.0;
        }
        if (!equations(fcn,t+h,work,dx)) return false;
        for (int j = 0; j < 2; ++j) {
            k5[j] = h*dx[j];
            work[j] = x[j]-8.0*k1[j]/27.0+2.0*k2[j]
                      -3544.0*k3[j]/2565.0+1859.0*k4[j]/4104.0-11.0*k5[j]/40.0;
        }
        if (!equations(fcn,t+h/2.0,work,dx)) return false;
        for (int j = 0; j < 2; ++j) k6[j] = h*dx[j];

        double err = 0.0;
        for (int j = 0; j < 2; ++j) {
            trial[j] = x[j]+16.0*k1[j]/135.0+6656.0*k3[j]/12825.0
                       +28561.0*k4[j]/56430.0-9.0*k5[j]/50.0+2.0*k6[j]/55.0;
            double component_error = std::abs(k1[j]/360.0-128.0*k3[j]/4275.0
                                     -2197.0*k4[j]/75240.0+k5[j]/50.0+2.0*k6[j]/55.0);
            if (!std::isfinite(trial[j]) || !std::isfinite(component_error)) {
                std::cerr << "Nonfinite RKF45 result: check model and slope guesses.\n";
                return false;
            }
            err = std::max(err,component_error);
        }
        if (err <= tolerance) {
            if (h == t_end-t) tf = t_end;
            t = tf;
            x = trial;
            ++accepted;
            if (accepted+1 > capacity) {
                std::cerr << "Increase trajectory storage capacity.\n";
                return false;
            }
            grid[accepted] = t;
            solution[accepted] = x;
        } else {
            ++rejected;  // Keep both t and x unchanged.
        }
        double factor = (err == 0.0) ? 2.0 : 0.9*std::pow(tolerance/err,0.2);
        if (err <= tolerance) h *= std::min(2.0,factor);
        else h *= std::max(0.2,factor);
    }
    xf = x;
    return true;
}