//=======================================================================
// Chapter 9, code 4: two second-order ODEs, solved with adaptive RKF45.
// Projectile motion in a vertical plane; upward y is positive.
// Gravity-only model: x'' = 0, y'' = -gravity.
// Introduce vx = x' and vy = y':
//   x' = vx, y' = vy, vx' = 0, vy' = -gravity.
// Store x, y, vx, vy in state[0], state[1], state[2], state[3].
//
// Optional extension: quadratic air resistance, F_drag = -drag*speed*v.
// Set drag > 0 to activate it; drag = 0 gives the gravity-only model.
// SI units: mass in kg, gravity in m/s^2, drag in kg/m.
// Air is stationary; gravity and drag coefficient are constant.
//
// To use another physical model:
//   1. Edit parameters and initial data below.
//   2. Edit the two acceleration equations in equations.
//   3. Disable check_exact or supply a matching exact solution.
// Integration ends at t_end, without testing for ground contact.
// h_max gives enough points for plotting; accuracy may need smaller steps.
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

void equations(double t, const std::vector<double>& state,
               std::vector<double>& deriv, double mass, double gravity, double drag);
bool rkf45_system(void (*f)(double,const std::vector<double>&,std::vector<double>&,double,double,double),
                   int n, double t0, const std::vector<double>& x0,
                   double t_end, double h_first, double h_max, double tolerance,
                   double mass, double gravity, double drag,
                   std::ofstream& data, std::vector<double>& xf,
                   int& accepted, int& rejected);

int main()
{
    //--------------- Parameters and initial data: edit here --------------
    const int n = 4;
    double mass = 1.0;
    double gravity = 9.81;
    double drag = 0.0;       // Try 0.01 for quadratic air resistance
    double t0 = 0.0;
    double t_end = 4.0;
    std::vector<double> state0(n);
    state0[0] = 0.0;        // Initial horizontal position (m)
    state0[1] = 0.0;        // Initial vertical position (m)
    state0[2] = 15.0;       // Initial horizontal velocity (m/s)
    state0[3] = 20.0;       // Initial vertical velocity (m/s)
    double h_first = 0.1;
    double h_max = 0.05;    // Maximum step for plotting the trajectory
    double tolerance = 1.0e-8;
    bool check_exact = true; // Comparison applies only when drag = 0
    //-------------------------------------------------------------------
    if (mass <= 0.0 || gravity < 0.0 || drag < 0.0) {
        std::cerr << "Use positive mass, nonnegative gravity, and nonnegative drag.\n";
        return 1;
    }
    std::ofstream data("ch09_rkf45_projectile.dat");
    if (!data) {
        std::cerr << "Cannot open ch09_rkf45_projectile.dat.\n";
        return 1;
    }
    std::cout << "Adaptive RKF45: projectile motion\n"
              << "Drag coefficient: " << std::scientific << std::setprecision(4)
              << drag << '\n'
              << "Tolerance: " << tolerance << '\n';
    data << std::scientific << std::setprecision(16);
    std::vector<double> statef(n);
    int accepted, rejected;
    if (!rkf45_system(equations,n,t0,state0,t_end,h_first,h_max,tolerance,
                      mass,gravity,drag,data,statef,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 && drag == 0.0) {
        // Constant gravity without air resistance, for arbitrary initial data.
        double tau = t_end-t0;
        std::vector<double> exact(n);
        exact[0] = state0[0]+state0[2]*tau;
        exact[1] = state0[1]+state0[3]*tau-gravity*tau*tau/2.0;
        exact[2] = state0[2];
        exact[3] = state0[3]-gravity*tau;
        const char* names[] = {"x","y","vx","vy"};
        std::cout << std::setw(10) << "variable"
                  << ' ' << std::setw(18) << "numerical"
                  << ' ' << std::setw(18) << "exact"
                  << ' ' << std::setw(12) << "abs error" << '\n';
        for (int j = 0; j < n; ++j) {
            std::cout << std::setw(10) << names[j]
                      << ' ' << std::scientific << std::setprecision(10)
                      << std::setw(18) << statef[j]
                      << ' ' << std::setw(18) << exact[j]
                      << ' ' << std::setprecision(4) << std::setw(12)
                      << std::abs(statef[j]-exact[j]) << '\n';
        }
    }
    std::cout << "Data file: ch09_rkf45_projectile.dat\n";
    return 0;
}

void equations(double t, const std::vector<double>& state,
               std::vector<double>& deriv, double mass, double gravity, double drag)
{
    // state[0] = x, state[1] = y, state[2] = vx, state[3] = vy.
    // The first two derivatives are velocities; the last two are accelerations.
    // For another force, change deriv[2] and deriv[3].
    // t is available for forces that depend explicitly on time.
    (void)t;
    double speed = std::sqrt(state[2]*state[2]+state[3]*state[3]);
    deriv[0] = state[2];
    deriv[1] = state[3];
    deriv[2] = -(drag/mass)*speed*state[2];
    deriv[3] = -gravity-(drag/mass)*speed*state[3];
}

bool rkf45_system(void (*f)(double,const std::vector<double>&,std::vector<double>&,double,double,double),
                   int n, double t0, const std::vector<double>& x0,
                   double t_end, double h_first, double h_max, double tolerance,
                   double mass, double gravity, double drag,
                   std::ofstream& data, std::vector<double>& xf,
                   int& accepted, int& rejected)
{
    // Integrate the four position/velocity equations with adaptive Fehlberg RKF45.
    // h_max caps output spacing; accuracy may require smaller steps.
    // 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 || h_max <= 0.0 || tolerance <= 0.0) {
        std::cerr << "Use n > 0, t_end > t0, positive steps, and positive tolerance.\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;
    // Output: time, x, y, vx, vy, and the step used to reach the point.
    std::cout << std::setw(10) << "t" << ' ' << std::setw(18) << "x"
              << ' ' << std::setw(18) << "y"
              << ' ' << std::setw(18) << "vx"
              << ' ' << std::setw(18) << "vy"
              << ' ' << std::setw(18) << "h" << '\n';
    data << "# t x y vx vy 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,std::min(h_max,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,mass,gravity,drag);
        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,mass,gravity,drag);
        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,mass,gravity,drag);
        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,mass,gravity,drag);
        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,mass,gravity,drag);
        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,mass,gravity,drag);
        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;
}