//=======================================================================
// Chapter 9, code 0: fixed-step methods for one first-order ODE.
// Solve x' = f(t,x), given x(t0) = x0.
//
// key = 1: Euler                     (first order)
// key = 2: modified Euler / Heun      (second order)
// key = 3: classical Runge-Kutta      (fourth order)
//
// Translated from the revised Fortran companion code.
// Based on Ode11.f90 by Alexander Godunov (February 2010).
// The three methods are kept in separate, short functions.
// Heun uses an Euler predictor and ONE trapezoidal corrector.
//
// To solve another equation:
//   1. Change x0, t0, t_end, and dt in the driver.
//   2. Edit the derivative function f below.
//   3. Set check_exact = false unless the exact formula also matches.
//
// Default: x' = x-t*t+1, x(0) = 0.5, 0 <= t <= 4.
// Exact solution: x(t) = (t+1)*(t+1)-0.5*exp(t).
// Try dt = 0.2, 0.1, and 0.05 with each method key.
// Halving dt reduces the final error asymptotically by factors of
// about 2 (Euler), 4 (Heun), and 16 (RK4).
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>

// Declarations allow the driver to appear before the equation and methods.
double f(double t, double x);
double euler1(double (*f)(double,double), double ti, double xi, double tf);
double heun1(double (*f)(double,double), double ti, double xi, double tf);
double rk4_scalar(double (*f)(double,double), double ti, double xi, double tf);

int main()
{
    //--------------- Initial data: edit here -----------------------------
    double t0 = 0.0;
    double x0 = 0.5;
    double t_end = 4.0;
    double dt = 0.1;
    int key = 3;
    bool check_exact = true;
    //-------------------------------------------------------------------
    const int max_steps = 100000;
    const char* method;
    const char* filename;

    if (dt <= 0.0 || t_end <= t0) {
        std::cerr << "Use dt > 0 and t_end > t0.\n";
        return 1;
    }
    switch (key) {
    case 1:
        method = "Euler";
        filename = "ch09_euler.dat";
        break;
    case 2:
        method = "Modified Euler (Heun)";
        filename = "ch09_heun.dat";
        break;
    case 3:
        method = "Runge-Kutta 4th order";
        filename = "ch09_rk4.dat";
        break;
    default:
        std::cerr << "Choose key = 1, 2, or 3.\n";
        return 1;
    }

    // One plain text file per run; a rerun of that method overwrites it.
    std::ofstream data(filename);
    if (!data) {
        std::cerr << "Cannot open data file: " << filename << '\n';
        return 1;
    }
    std::cout << "Method: " << method << '\n'
              << "Fixed step: " << std::fixed << std::setprecision(5)
              << dt << '\n';
    std::cout << std::setw(10) << "t"
              << ' ' << std::setw(18) << "numerical x";
    if (check_exact) {
        std::cout << ' ' << std::setw(18) << "exact x"
                  << ' ' << std::setw(14) << "abs error";
        data << "# t x_numerical x_exact abs_error\n";
    } else {
        data << "# t x_numerical\n";
    }
    std::cout << '\n';
    data << std::scientific << std::setprecision(16);

    double ti = t0;
    double xi = x0;
    int i = 0;
    // Print the initial point before taking any step.
    while (true) {
        std::cout << std::fixed << std::setprecision(5) << std::setw(10) << ti
                  << ' ' << std::setprecision(12) << std::setw(18) << xi;
        data << ti << ' ' << xi;
        if (check_exact) {
            // Valid for this equation even when t0 or x0 changes.
            double exact = (ti+1.0)*(ti+1.0)
                         + (x0-(t0+1.0)*(t0+1.0))*std::exp(ti-t0);
            double error = std::abs(xi-exact);
            std::cout << ' ' << std::setw(18) << exact
                      << ' ' << std::scientific << std::setprecision(5)
                      << std::setw(14) << error;
            data << ' ' << exact << ' ' << error;
        }
        std::cout << '\n';
        data << '\n';
        if (ti >= t_end) break;
        if (i >= max_steps) {
            std::cerr << "Maximum number of steps exceeded.\n";
            return 1;
        }

        // Indexed times avoid repeated-addition drift.
        // Shorten the last step if dt does not divide the interval.
        ++i;
        double tf = std::min(t_end,t0+static_cast<double>(i)*dt);
        double roundoff = 64.0*std::numeric_limits<double>::epsilon()
                        * std::max(1.0,std::abs(t_end));
        if (t_end-tf <= roundoff) tf = t_end;
        if (tf <= ti) {
            std::cerr << "Step too small to advance time.\n";
            return 1;
        }
        double xf = xi;
        switch (key) {
        case 1: xf = euler1(f,ti,xi,tf); break;
        case 2: xf = heun1(f,ti,xi,tf); break;
        case 3: xf = rk4_scalar(f,ti,xi,tf); break;
        }
        ti = tf;
        xi = xf;
    }
    data.close();
    std::cout << "Steps: " << i << '\n'
              << "Data file: " << filename << '\n';
    return 0;
}

double f(double t, double x)
{
    // Supply the equation here: return the derivative x' = f(t,x).
    return x-t*t+1.0;
}

double euler1(double (*f)(double,double), double ti, double xi, double tf)
{
    // One Euler step: slope at the beginning of the interval.
    double h = tf-ti;
    return xi+h*f(ti,xi);
}

double heun1(double (*f)(double,double), double ti, double xi, double tf)
{
    // Euler predictor, then ONE trapezoidal corrector.
    double h = tf-ti;
    double slope1 = f(ti,xi);
    double predicted = xi+h*slope1;
    double slope2 = f(tf,predicted);
    return xi+0.5*h*(slope1+slope2);
}

double rk4_scalar(double (*f)(double,double), double ti, double xi, double tf)
{
    // Classical RK4. Each k includes h.
    double h = tf-ti;
    double k1 = h*f(ti,xi);
    double k2 = h*f(ti+h/2.0,xi+k1/2.0);
    double k3 = h*f(ti+h/2.0,xi+k2/2.0);
    double k4 = h*f(tf,xi+k3);
    return xi+(k1+2.0*k2+2.0*k3+k4)/6.0;
}