program main
!=======================================================================
! 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(1), state(2), state(3), state(4).
!
! 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 contact with the ground.
! Choose t_end appropriately or explore ground detection as an exercise.
!=======================================================================
implicit none
integer, parameter :: n = 4, data_unit = 20
double precision mass, gravity, drag
double precision t0, t_end, h_first, h_max, tolerance, tau
double precision state0(n), statef(n), exact(n)
integer accepted, rejected, j
logical check_exact
character(len=10) names(n)
external equations

!-------------------- Parameters and initial data: edit here -------------
mass = 1.0d0
gravity = 9.81d0
drag = 0.0d0            ! Try 0.01d0 for quadratic air resistance
t0 = 0.0d0
t_end = 4.0d0
state0(1) = 0.0d0       ! Initial horizontal position (m)
state0(2) = 0.0d0       ! Initial vertical position (m)
state0(3) = 15.0d0      ! Initial horizontal velocity (m/s)
state0(4) = 20.0d0      ! Initial vertical velocity (m/s)
h_first = 0.1d0
h_max = 0.05d0          ! Maximum step: enough output points for plotting
tolerance = 1.0d-8      ! Absolute local tolerance for each component
check_exact = .true.    ! Comparison applies only when drag = 0
!------------------------------------------------------------------------
if (mass <= 0.0d0 .or. gravity < 0.0d0 .or. drag < 0.0d0) &
    stop 'Use positive mass, nonnegative gravity, and nonnegative drag.'

open(unit=data_unit, file='ch09_rkf45_projectile.dat', &
     status='replace', action='write')
write(*,'(a)') 'Adaptive RKF45: projectile motion'
write(*,'(a,es12.4)') 'Drag coefficient: ', drag
write(*,'(a,es12.4)') 'Tolerance: ', tolerance
call rkf45_system(equations,n,t0,state0,t_end,h_first,h_max,tolerance, &
                  mass,gravity,drag,data_unit,statef,accepted,rejected)
close(data_unit)
write(*,'(a,i8)') 'Accepted steps: ', accepted
write(*,'(a,i8)') 'Rejected trials:', rejected

if (check_exact .and. drag == 0.0d0) then
    ! Analytic solution for constant gravity without air resistance.
    tau = t_end-t0
    exact(1) = state0(1)+state0(3)*tau
    exact(2) = state0(2)+state0(4)*tau-gravity*tau**2/2.0d0
    exact(3) = state0(3)
    exact(4) = state0(4)-gravity*tau
    names(1) = 'x'
    names(2) = 'y'
    names(3) = 'vx'
    names(4) = 'vy'
    write(*,'(a10,2(1x,a18),1x,a12)') 'variable','numerical','exact','abs error'
    do j = 1,n
        write(*,'(a10,2(1x,es18.10),1x,es12.4)') &
            names(j),statef(j),exact(j),abs(statef(j)-exact(j))
    end do
end if
write(*,'(a)') 'Data file: ch09_rkf45_projectile.dat'
end program main

subroutine equations(t,state,deriv,n,mass,gravity,drag)
!=======================================================================
! state(1) = x, state(2) = y, state(3) = vx, state(4) = vy.
! The first two derivatives are velocities; the last two are accelerations.
! For another force, change deriv(3) and deriv(4).
! Time t is available for a force that depends explicitly on time.
!=======================================================================
implicit none
integer n
double precision t, state(n), deriv(n), mass, gravity, drag, speed
speed = sqrt(state(3)**2+state(4)**2)
deriv(1) = state(3)
deriv(2) = state(4)
deriv(3) = -(drag/mass)*speed*state(3)
deriv(4) = -gravity-(drag/mass)*speed*state(4)
end subroutine equations

subroutine rkf45_system(fcn,n,t0,x0,t_end,h_first,h_max,tolerance,mass,gravity,drag,data_unit, &
                        xf,accepted,rejected)
!=======================================================================
! Solve the four position/velocity equations with adaptive Fehlberg RKF45.
! All step decisions are inside this subroutine.
! h_max caps the output spacing; the tolerance can require smaller steps.
! fcn(t,x,dx,n,mass,gravity,drag) returns the derivatives in dx.
! Input: initial data, interval, step, tolerance, physical parameters,
!        and opened data_unit.
! Output: xf(n), accepted and rejected counts.
! Screen and file output include the initial point and accepted points.
!
! Advance with the fifth-order estimate. The largest absolute embedded
! component error controls acceptance, using ONE absolute tolerance.
! This is suitable for components of comparable scales. Variables with
! different units or magnitudes may require individual scaled tolerances.
! The tolerance controls local error, not accumulated global error.
!=======================================================================
implicit none
integer n, data_unit, accepted, rejected, attempts, j
double precision t0, x0(n), t_end, h_first, h_max, tolerance, xf(n)
double precision mass, gravity, drag
double precision t, tf, h, err, component_error, factor
double precision x(n), work(n), dx(n), trial(n)
double precision k1(n), k2(n), k3(n), k4(n), k5(n), k6(n)
integer, parameter :: max_attempts = 100000
double precision, parameter :: h_min = 1.0d-12
external fcn

if (n < 1 .or. t_end <= t0 .or. h_first <= 0.0d0 .or. h_max <= 0.0d0 .or. tolerance <= 0.0d0) &
    stop 'Use n > 0, t_end > t0, positive steps, and positive tolerance.'
! Output columns: time, x, y, vx, vy, and the accepted step size.
write(*,'(a10,5(1x,a18))') 't','x','y','vx','vy','h'
write(data_unit,'(a)') '# t x y vx vy h'

t = t0
h = h_first
x = x0
accepted = 0
rejected = 0
attempts = 0
write(*,'(f10.6,5(1x,es18.10))') t,x,0.0d0
write(data_unit,'(6(1x,es24.16))') t,x,0.0d0

do while (t < t_end)
    if (t_end-t <= 64.0d0*epsilon(t)*max(1.0d0,abs(t_end))) exit
    if (attempts >= max_attempts) stop 'Maximum number of trials exceeded.'
    h = min(h,h_max,t_end-t)
    tf = t+h
    if (h < h_min .or. tf <= t) stop 'Step too small to advance time.'
    attempts = attempts+1

    ! Stage 1, then prepare the state for stage 2. Every k includes h.
    call fcn(t,x,dx,n,mass,gravity,drag)
    do j = 1,n
        k1(j) = h*dx(j)
        work(j) = x(j)+k1(j)/4.0d0
    end do
    ! Stage 2: all equations use the SAME intermediate state work.
    call fcn(t+h/4.0d0,work,dx,n,mass,gravity,drag)
    do j = 1,n
        k2(j) = h*dx(j)
        work(j) = x(j)+3.0d0*k1(j)/32.0d0+9.0d0*k2(j)/32.0d0
    end do
    call fcn(t+3.0d0*h/8.0d0,work,dx,n,mass,gravity,drag)
    do j = 1,n
        k3(j) = h*dx(j)
        work(j) = x(j)+1932.0d0*k1(j)/2197.0d0 &
                  -7200.0d0*k2(j)/2197.0d0+7296.0d0*k3(j)/2197.0d0
    end do
    call fcn(t+12.0d0*h/13.0d0,work,dx,n,mass,gravity,drag)
    do j = 1,n
        k4(j) = h*dx(j)
        work(j) = x(j)+439.0d0*k1(j)/216.0d0-8.0d0*k2(j) &
                  +3680.0d0*k3(j)/513.0d0-845.0d0*k4(j)/4104.0d0
    end do
    call fcn(t+h,work,dx,n,mass,gravity,drag)
    do j = 1,n
        k5(j) = h*dx(j)
        work(j) = x(j)-8.0d0*k1(j)/27.0d0+2.0d0*k2(j) &
                  -3544.0d0*k3(j)/2565.0d0+1859.0d0*k4(j)/4104.0d0 &
                  -11.0d0*k5(j)/40.0d0
    end do
    call fcn(t+h/2.0d0,work,dx,n,mass,gravity,drag)
    do j = 1,n
        k6(j) = h*dx(j)
    end do

    ! Fifth-order candidate and largest embedded component error.
    err = 0.0d0
    do j = 1,n
        trial(j) = x(j)+16.0d0*k1(j)/135.0d0+6656.0d0*k3(j)/12825.0d0 &
                   +28561.0d0*k4(j)/56430.0d0-9.0d0*k5(j)/50.0d0+2.0d0*k6(j)/55.0d0
        component_error = abs(k1(j)/360.0d0-128.0d0*k3(j)/4275.0d0 &
                         -2197.0d0*k4(j)/75240.0d0+k5(j)/50.0d0+2.0d0*k6(j)/55.0d0)
        err = max(err,component_error)
    end do
    if (err <= tolerance) then
        if (h == t_end-t) tf = t_end
        t = tf
        x = trial
        accepted = accepted+1
        write(*,'(f10.6,5(1x,es18.10))') t,x,h
        write(data_unit,'(6(1x,es24.16))') t,x,h
    else
        rejected = rejected+1  ! Keep both t and x unchanged.
    end if
    if (err == 0.0d0) then
        factor = 2.0d0
    else
        factor = 0.9d0*(tolerance/err)**0.2d0
    end if
    if (err <= tolerance) then
        h = h*min(2.0d0,factor)
    else
        h = h*max(0.2d0,factor)
    end if
end do
xf = x
end subroutine rkf45_system