%=======================================================================
% 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 ground contact.
% h_max gives enough trajectory points; accuracy may need smaller steps.
% Returned S has one time per row and columns x, y, vx, vy.
%=======================================================================
clear;

%-------------------- Parameters and initial data: edit here -------------
mass = 1;
gravity = 9.81;
drag = 0;              % Try 0.01 for quadratic air resistance
t0 = 0;
t_end = 4;
state0 = [0; 0; 15; 20]; % Initial x, y (m), vx, vy (m/s)
h_first = 0.1;
h_max = 0.05;          % Maximum step for plotting the trajectory
tolerance = 1.0e-8;    % Absolute local tolerance for each component
check_exact = true;    % Comparison applies only when drag = 0
%------------------------------------------------------------------------
if mass <= 0 || gravity < 0 || drag < 0
    error('Use positive mass, nonnegative gravity, and nonnegative drag.');
end
[t,S,h,accepted,rejected] = rkf45_system(@equations,4,t0,state0,t_end, ...
    h_first,h_max,tolerance,mass,gravity,drag);

% Print and save the initial point and every accepted point.
filename = 'ch09_rkf45_projectile.dat';
fid = fopen(filename,'w');
if fid == -1
    error('Cannot open data file: %s',filename);
end
fprintf('Adaptive RKF45: projectile motion\n');
fprintf('Drag coefficient: %.4e\n',drag);
fprintf('Tolerance: %.4e\n',tolerance);
fprintf('%10s %18s %18s %18s %18s %18s\n','t','x','y','vx','vy','h');
fprintf(fid,'# t x y vx vy h\n');
for i = 1:length(t)
    fprintf('%10.6f',t(i));
    fprintf(' %18.10e',S(i,:));
    fprintf(' %18.10e\n',h(i));
    fprintf(fid,'%.16e',t(i));
    fprintf(fid,' %.16e',S(i,:));
    fprintf(fid,' %.16e\n',h(i));
end
fclose(fid);
fprintf('Accepted steps: %8d\n',accepted);
fprintf('Rejected trials:%8d\n',rejected);
if check_exact && drag == 0
    exact = exact_solution(t_end,t0,state0,gravity);
    names = {'x','y','vx','vy'};
    fprintf('%10s %18s %18s %12s\n','variable','numerical','exact','abs error');
    for j = 1:4
        fprintf('%10s %18.10e %18.10e %12.4e\n', ...
                names{j},S(end,j),exact(j),abs(S(end,j)-exact(j)));
    end
end
fprintf('Data file: %s\n',filename);

% Trajectory: actual accepted points, plus a dense exact curve if available.
figure;
if check_exact && drag == 0
    t_plot = linspace(t0,t_end,1000)';
    S_plot = exact_solution(t_plot,t0,state0,gravity);
    plot(S_plot(:,1),S_plot(:,2),'k-','LineWidth',1.2);
    hold on;
    plot(S(:,1),S(:,2),'ko--','MarkerSize',3);
    hold off;
    legend('Exact','RKF45','Location','best');
else
    plot(S(:,1),S(:,2),'ko-','MarkerSize',3);
end
xlabel('Horizontal position x (m)');
ylabel('Vertical position y (m)');
title(sprintf('Projectile trajectory, drag = %g',drag));
axis equal;
grid on;

% h(i) is the accepted step used to reach t(i); omit the initial h = 0.
figure;
semilogy(t(2:end),h(2:end),'ko-','MarkerSize',3);
xlabel('t (s)');
ylabel('Accepted step size h (s)');
title('Adaptive RKF45: step size');
grid on;

%==================== Equations and exact solution ======================
function deriv = equations(t,state,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).
% t is available for forces that depend explicitly on time.
speed = sqrt(state(3)^2+state(4)^2);
deriv = zeros(4,1);
deriv(1) = state(3);
deriv(2) = state(4);
deriv(3) = -(drag/mass)*speed*state(3);
deriv(4) = -gravity-(drag/mass)*speed*state(4);
end

function S = exact_solution(t,t0,state0,gravity)
% Constant gravity without air resistance; arbitrary initial data.
% One time per row, columns x, y, vx, vy.
tau = t(:)-t0;
x = state0(1)+state0(3)*tau;
y = state0(2)+state0(4)*tau-gravity*tau.^2/2;
vx = state0(3)*ones(size(tau));
vy = state0(4)-gravity*tau;
S = [x,y,vx,vy];
end

%==================== Numerical method ==================================
function [t_values,X,h_values,accepted,rejected] = ...
    rkf45_system(f,n,t0,x0,t_end,h_first,h_max,tolerance,mass,gravity,drag)
% Adaptive Fehlberg RKF45; advance with the fifth-order estimate.
% h_max caps output spacing; accuracy may require smaller steps.
% The largest absolute embedded component error controls acceptance.
% One tolerance assumes comparable numerical scales; rescale variables
% or use component tolerances for widely differing magnitudes or units.
% Local error control does not bound accumulated global error.
max_attempts = 100000;
h_min = 1.0e-12;
if n < 1 || length(x0) ~= n || t_end <= t0 || h_first <= 0 || h_max <= 0 || tolerance <= 0
    error('Use n initial values, t_end > t0, and positive steps and tolerance.');
end
t = t0;
x = x0(:);
h = h_first;
accepted = 0;
rejected = 0;
attempts = 0;
t_values = t0;
X = x';
h_values = 0;

while t < t_end
    if t_end-t <= 64*eps(1)*max(1,abs(t_end))
        break;
    end
    if attempts >= max_attempts
        error('Maximum number of trials exceeded.');
    end
    h = min([h,h_max,t_end-t]);
    tf = t+h;
    if h < h_min || tf <= t
        error('Step too small to advance time.');
    end
    attempts = attempts+1;

    % Each k is a column vector and includes h.
    % All equations use the SAME complete intermediate state at each stage.
    k1 = h*f(t,x,mass,gravity,drag);
    k2 = h*f(t+h/4,x+k1/4,mass,gravity,drag);
    k3 = h*f(t+3*h/8,x+3*k1/32+9*k2/32,mass,gravity,drag);
    k4 = h*f(t+12*h/13,x+1932*k1/2197-7200*k2/2197+7296*k3/2197,mass,gravity,drag);
    k5 = h*f(t+h,x+439*k1/216-8*k2+3680*k3/513-845*k4/4104,mass,gravity,drag);
    k6 = h*f(t+h/2,x-8*k1/27+2*k2-3544*k3/2565 ...
               +1859*k4/4104-11*k5/40,mass,gravity,drag);

    trial = x+16*k1/135+6656*k3/12825+28561*k4/56430 ...
             -9*k5/50+2*k6/55;
    component_error = abs(k1/360-128*k3/4275-2197*k4/75240+k5/50+2*k6/55);
    err = max(component_error);
    if err <= tolerance
        if h == t_end-t
            tf = t_end;
        end
        t = tf;
        x = trial;
        accepted = accepted+1;
        t_values(end+1,1) = t;
        X(end+1,:) = x';
        h_values(end+1,1) = h;
    else
        rejected = rejected+1; % Keep both t and x unchanged.
    end

    % Embedded error scales approximately as h^5.
    if err == 0
        factor = 2;
    else
        factor = 0.9*(tolerance/err)^0.2;
    end
    if err <= tolerance
        h = h*min(2,factor);
    else
        h = h*max(0.2,factor);
    end
end
end