%=======================================================================
% Chapter 9, code 3: a second-order ODE solved with adaptive RKF45.
% Damped spring: mass*x'' + damping*x' + spring*x = 0.
% Introduce velocity v = x':
%   x' = v
%   v' = -(spring/mass)*x - (damping/mass)*v
% Store position in y(1), velocity in y(2).
%
% To use another physical model:
%   1. Edit parameters and initial data below.
%   2. Edit the acceleration dy(2) in equations.
%   3. Disable check_exact or replace exact_solution.
% All stages and adaptive decisions are inside rkf45_system.
% h_max gives enough points for plotting; tolerance controls local error.
% Returned Y has one time per row: Y(:,1) is position, Y(:,2) velocity.
%=======================================================================
clear;

%-------------------- Parameters and initial data: edit here -------------
mass = 1;
spring = 1;
damping = 0.2;
t0 = 0;
t_end = 20;
y0 = [1; 0];           % Initial position and velocity (column vector)
h_first = 0.1;
h_max = 0.1;           % Maximum step for plotting the oscillations
tolerance = 1.0e-8;    % Absolute local tolerance for each component
check_exact = true;    % Analytic comparison for an underdamped spring
%------------------------------------------------------------------------
if mass <= 0 || spring <= 0 || damping < 0
    error('Use positive mass and spring, and nonnegative damping.');
end
if check_exact && damping^2 >= 4*mass*spring
    error('Exact comparison requires an underdamped spring; disable check_exact.');
end
[t,Y,h,accepted,rejected] = rkf45_system(@equations,2,t0,y0,t_end, ...
    h_first,h_max,tolerance,mass,damping,spring);

% Print and save the initial point and every accepted point.
filename = 'ch09_rkf45_spring.dat';
fid = fopen(filename,'w');
if fid == -1
    error('Cannot open data file: %s',filename);
end
fprintf('Adaptive RKF45: damped spring\n');
fprintf('Tolerance: %.4e\n',tolerance);
fprintf('%10s %18s %18s %18s\n','t','position','velocity','h');
fprintf(fid,'# t position velocity h\n');
for i = 1:length(t)
    fprintf('%10.6f %18.10e %18.10e %18.10e\n',t(i),Y(i,1),Y(i,2),h(i));
    fprintf(fid,'%.16e %.16e %.16e %.16e\n',t(i),Y(i,1),Y(i,2),h(i));
end
fclose(fid);
fprintf('Accepted steps: %8d\n',accepted);
fprintf('Rejected trials:%8d\n',rejected);
if check_exact
    exact = exact_solution(t_end,t0,y0,mass,damping,spring);
    fprintf('%10s %18s %18s %12s\n','variable','numerical','exact','abs error');
    fprintf('%10s %18.10e %18.10e %12.4e\n', ...
            'position',Y(end,1),exact(1),abs(Y(end,1)-exact(1)));
    fprintf('%10s %18.10e %18.10e %12.4e\n', ...
            'velocity',Y(end,2),exact(2),abs(Y(end,2)-exact(2)));
end
fprintf('Data file: %s\n',filename);

% Position and velocity, compared with dense exact curves.
if check_exact
    t_plot = linspace(t0,t_end,2000)';
    Y_plot = exact_solution(t_plot,t0,y0,mass,damping,spring);
end
figure;
labels = {'Position x','Velocity v'};
for j = 1:2
    subplot(2,1,j);
    if check_exact
        plot(t_plot,Y_plot(:,j),'k-','LineWidth',1.2);
        hold on;
        plot(t,Y(:,j),'ko--','MarkerSize',3);
        hold off;
        if j == 1
            legend('Exact','RKF45','Location','best');
        end
    else
        plot(t,Y(:,j),'ko-','MarkerSize',3);
    end
    ylabel(labels{j});
    grid on;
    if j == 1
        title('Damped spring: position and velocity');
    else
        xlabel('t');
    end
end

figure;
semilogy(t(2:end),h(2:end),'ko-','MarkerSize',3);
xlabel('t');
ylabel('Accepted step size h');
title('Adaptive RKF45: step size');
grid on;

%==================== Equations and exact solution ======================
function dy = equations(t,y,mass,damping,spring)
% y(1) = position, y(2) = velocity. Parameters are passed explicitly.
% For another force, change the acceleration dy(2).
% t is available for forces that depend explicitly on time.
dy = zeros(2,1);
dy(1) = y(2);
dy(2) = -(spring/mass)*y(1)-(damping/mass)*y(2);
end

function Y = exact_solution(t,t0,y0,mass,damping,spring)
% Exact underdamped solution, including zero damping.
% One time per row: column 1 is position, column 2 is velocity.
tau = t(:)-t0;
gamma = damping/(2*mass);
omega = sqrt(spring/mass-gamma^2);
a = y0(1);
b = (y0(2)+gamma*a)/omega;
decay = exp(-gamma*tau);
x = decay.*(a*cos(omega*tau)+b*sin(omega*tau));
v = -gamma*x + decay.*omega.*(-a*sin(omega*tau)+b*cos(omega*tau));
Y = [x,v];
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,damping,spring)
% 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,damping,spring);
    k2 = h*f(t+h/4,x+k1/4,mass,damping,spring);
    k3 = h*f(t+3*h/8,x+3*k1/32+9*k2/32,mass,damping,spring);
    k4 = h*f(t+12*h/13,x+1932*k1/2197-7200*k2/2197+7296*k3/2197,mass,damping,spring);
    k5 = h*f(t+h,x+439*k1/216-8*k2+3680*k3/513-845*k4/4104,mass,damping,spring);
    k6 = h*f(t+h/2,x-8*k1/27+2*k2-3544*k3/2565 ...
               +1859*k4/4104-11*k5/40,mass,damping,spring);

    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