//=======================================================================
// Chapter 9, code 2: adaptive RKF45 for n first-order ODEs.
// Translated from the revised Fortran companion code.
// Mathematical example:
//   x1' = x2, x2' = -x1, x3' = -2*x3, x4' = x1+x3.
// Initial values at t = 0: (1,0,1,0).
// Exact: (cos(t), -sin(t), exp(-2*t), sin(t)+(1-exp(-2*t))/2).
//
// To supply another system:
//   1. Set n and all n initial values in the driver.
//   2. Edit equations: dx[j] is the derivative of x[j].
//   3. Choose the interval, first trial step, and tolerance.
//   4. Disable check_exact or supply matching exact values.
// C++ indices start at ZERO: x[0] represents mathematical x1.
// std::vector<double> holds n values; it permits changing n without
// changing the solver. All stages and step control are inside the solver.
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <string>
#include <vector>

void equations(double t, const std::vector<double>& x, std::vector<double>& dx);
bool rkf45_system(void (*f)(double,const std::vector<double>&,std::vector<double>&),
                   int n, double t0, const std::vector<double>& x0,
                   double t_end, double h_first, double tolerance,
                   std::ofstream& data, std::vector<double>& xf,
                   int& accepted, int& rejected);

int main()
{
    //--------------- Initial data: edit here -----------------------------
    const int n = 4;
    double t0 = 0.0;
    double t_end = 10.0;
    std::vector<double> x0(n);
    x0[0] = 1.0;
    x0[1] = 0.0;
    x0[2] = 1.0;
    x0[3] = 0.0;
    double h_first = 0.5;
    double tolerance = 1.0e-8; // Same absolute local tolerance per component
    bool check_exact = true;
    //-------------------------------------------------------------------
    std::ofstream data("ch09_rkf45_system.dat");
    if (!data) {
        std::cerr << "Cannot open ch09_rkf45_system.dat.\n";
        return 1;
    }
    std::cout << "Adaptive RKF45: system of first-order ODEs\n"
              << "Number of equations: " << n << '\n'
              << "Tolerance: " << std::scientific << std::setprecision(4)
              << tolerance << '\n';
    data << std::scientific << std::setprecision(16);
    std::vector<double> xf(n);
    int accepted, rejected;
    if (!rkf45_system(equations,n,t0,x0,t_end,h_first,tolerance,data,
                      xf,accepted,rejected)) return 1;
    data.close();
    std::cout << "Accepted steps: " << std::setw(8) << accepted << '\n'
              << "Rejected trials:" << std::setw(8) << rejected << '\n';
    if (check_exact) {
        // These exact values match ONLY the default system and initial data.
        std::vector<double> exact(n);
        exact[0] = std::cos(t_end);
        exact[1] = -std::sin(t_end);
        exact[2] = std::exp(-2.0*t_end);
        exact[3] = std::sin(t_end)+(1.0-std::exp(-2.0*t_end))/2.0;
        std::cout << std::setw(10) << "component"
                  << ' ' << std::setw(18) << "numerical x"
                  << ' ' << std::setw(18) << "exact x"
                  << ' ' << std::setw(12) << "abs error" << '\n';
        for (int j = 0; j < n; ++j) {
            std::cout << std::setw(10) << j+1
                      << ' ' << std::scientific << std::setprecision(10)
                      << std::setw(18) << xf[j]
                      << ' ' << std::setw(18) << exact[j]
                      << ' ' << std::setprecision(4) << std::setw(12)
                      << std::abs(xf[j]-exact[j]) << '\n';
        }
    }
    std::cout << "Data file: ch09_rkf45_system.dat\n";
    return 0;
}

void equations(double t, const std::vector<double>& x, std::vector<double>& dx)
{
    // One derivative per line, in the same order as the initial values.
    // Input x is read-only (const); fill ALL entries of dx.
    // Time t is available for equations that depend explicitly on time.
    (void)t;
    dx[0] = x[1];
    dx[1] = -x[0];
    dx[2] = -2.0*x[2];
    dx[3] = x[0]+x[2];
}

bool rkf45_system(void (*f)(double,const std::vector<double>&,std::vector<double>&),
                   int n, double t0, const std::vector<double>& x0,
                   double t_end, double h_first, double tolerance,
                   std::ofstream& data, std::vector<double>& xf,
                   int& accepted, int& rejected)
{
    // Integrate n coupled equations with adaptive Fehlberg RKF45.
    // Use one absolute tolerance and the largest component error.
    // This assumes comparable component scales; rescale variables if needed.
    // Six stages give fourth- and fifth-order estimates. Advance with
    // the FIFTH-order value; their difference estimates local error.
    // The tolerance does not bound total accumulated error.
    // xf and the counters are returned through reference arguments (&).
    // Return false if the integration cannot be completed.
    const int max_attempts = 100000;
    const double h_min = 1.0e-12;
    if (n < 1 || t_end <= t0 || h_first <= 0.0 || tolerance <= 0.0) {
        std::cerr << "Use n > 0, t_end > t0, h_first > 0, and tolerance > 0.\n";
        return false;
    }
    double t = t0;
    std::vector<double> x = x0;
    std::vector<double> work(n), dx(n), trial(n);
    std::vector<double> k1(n), k2(n), k3(n), k4(n), k5(n), k6(n);
    double h = h_first;
    accepted = 0;
    rejected = 0;
    int attempts = 0;
    std::cout << std::setw(10) << "t";
    data << "# t";
    for (int j = 0; j < n; ++j) {
        std::string label = "x"+std::to_string(j+1);
        std::cout << ' ' << std::setw(18) << label;
        data << ' ' << label;
    }
    std::cout << ' ' << std::setw(18) << "h" << '\n';
    data << " h\n";
    std::cout << std::fixed << std::setprecision(6) << std::setw(10) << t;
    data << t;
    for (int j = 0; j < n; ++j) {
        std::cout << ' ' << std::scientific << std::setprecision(10)
                  << std::setw(18) << x[j];
        data << ' ' << x[j];
    }
    std::cout << ' ' << std::setw(18) << 0.0 << '\n';
    data << ' ' << 0.0 << '\n';

    while (t < t_end) {
        // Ignore only a machine-rounding gap at the endpoint.
        double roundoff = 64.0*std::numeric_limits<double>::epsilon()
                        * std::max(1.0,std::abs(t_end));
        if (t_end-t <= roundoff) break;
        if (attempts >= max_attempts) {
            std::cerr << "Maximum number of trials exceeded.\n";
            return false;
        }
        h = std::min(h,t_end-t);
        double tf = t+h;
        if (h < h_min || tf <= t) {
            std::cerr << "Step too small to advance time.\n";
            return false;
        }
        ++attempts;

        // Fill the complete intermediate state before evaluating equations.
        f(t,x,dx);
        for (int j = 0; j < n; ++j) {
            k1[j] = h*dx[j];
            work[j] = x[j]+k1[j]/4.0;
        }
        f(t+h/4.0,work,dx);
        for (int j = 0; j < n; ++j) {
            k2[j] = h*dx[j];
            work[j] = x[j]+3.0*k1[j]/32.0+9.0*k2[j]/32.0;
        }
        f(t+3.0*h/8.0,work,dx);
        for (int j = 0; j < n; ++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;
        }
        f(t+12.0*h/13.0,work,dx);
        for (int j = 0; j < n; ++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;
        }
        f(t+h,work,dx);
        for (int j = 0; j < n; ++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;
        }
        f(t+h/2.0,work,dx);
        for (int j = 0; j < n; ++j) k6[j] = h*dx[j];

        // Fifth-order candidates and largest embedded component error.
        double error = 0.0;
        for (int j = 0; j < n; ++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);
            error = std::max(error,component_error);
        }
        if (error <= tolerance) {
            if (h == t_end-t) tf = t_end;
            t = tf;
            x = trial;
            ++accepted;
            // h is the step used to reach this accepted point.
            std::cout << std::fixed << std::setprecision(6) << std::setw(10) << t;
            data << t;
            for (int j = 0; j < n; ++j) {
                std::cout << ' ' << std::scientific << std::setprecision(10)
                          << std::setw(18) << x[j];
                data << ' ' << x[j];
            }
            std::cout << ' ' << std::setw(18) << h << '\n';
            data << ' ' << h << '\n';
        } else {
            ++rejected;    // Keep both t and x unchanged.
        }

        // Embedded error scales approximately as h^5.
        // Safety factor 0.9; limit growth and shrinkage.
        double factor;
        if (error == 0.0)
            factor = 2.0;
        else
            factor = 0.9*std::pow(tolerance/error,0.2);
        if (error <= tolerance)
            h *= std::min(2.0,factor);
        else
            h *= std::max(0.2,factor);
    }
    xf = x;
    return true;
}