%=======================================================================
% Chapter 9, code 0: fixed-step methods for one first-order ODE.
% Solve x' = f(t,x), given x(t0) = x0.
% Based on the revised Fortran companion code.
%
% key = 1: Euler                     (first order)
% key = 2: modified Euler / Heun      (second order)
% key = 3: classical Runge-Kutta      (fourth order)
%
% 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^2+1, x(0) = 0.5, 0 <= t <= 4.
% Exact: x(t) = (t+1)^2-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).
%=======================================================================
clear;

%-------------------- Initial data: edit here ----------------------------
t0 = 0;
x0 = 0.5;
t_end = 4;
dt = 0.1;
key = 3;
check_exact = true;
%------------------------------------------------------------------------
max_steps = 100000;
if dt <= 0 || t_end <= t0
    error('Use dt > 0 and t_end > t0.');
end
switch key
    case 1
        method = 'Euler';
        filename = 'ch09_euler.dat';
    case 2
        method = 'Modified Euler (Heun)';
        filename = 'ch09_heun.dat';
    case 3
        method = 'Runge-Kutta 4th order';
        filename = 'ch09_rk4.dat';
    otherwise
        error('Choose key = 1, 2, or 3.');
end

% Build indexed times to avoid repeated-addition drift.
% Include t_end, shortening the last step if necessary.
if ceil((t_end-t0)/dt) > max_steps
    error('Maximum number of steps exceeded.');
end
t = t0 + (0:floor((t_end-t0)/dt))'*dt;
roundoff = 64*eps(1)*max(1,abs(t_end));
if t_end-t(end) <= roundoff
    t(end) = t_end;
else
    t(end+1) = t_end;
end
if any(diff(t) <= 0)
    error('Step too small to advance time.');
end
steps = length(t)-1;
x = zeros(size(t));
x(1) = x0;

for i = 1:steps
    switch key
        case 1
            x(i+1) = euler1(@f,t(i),x(i),t(i+1));
        case 2
            x(i+1) = heun1(@f,t(i),x(i),t(i+1));
        case 3
            x(i+1) = rk4_scalar(@f,t(i),x(i),t(i+1));
    end
end

% One plain text file per run; a rerun of that method overwrites it.
fid = fopen(filename,'w');
if fid == -1
    error('Cannot open data file: %s',filename);
end
fprintf('Method: %s\n',method);
fprintf('Fixed step: %.5f\n',dt);
if check_exact
    % Valid for this equation even when t0 or x0 changes.
    exact = (t+1).^2 + (x0-(t0+1)^2)*exp(t-t0);
    abs_error = abs(x-exact);
    fprintf('%10s %18s %18s %14s\n','t','numerical x','exact x','abs error');
    fprintf(fid,'# t x_numerical x_exact abs_error\n');
    for i = 1:length(t)
        fprintf('%10.5f %18.12f %18.12f %14.5e\n', ...
                t(i),x(i),exact(i),abs_error(i));
        fprintf(fid,'%.16e %.16e %.16e %.16e\n', ...
                t(i),x(i),exact(i),abs_error(i));
    end
else
    fprintf('%10s %18s\n','t','numerical x');
    fprintf(fid,'# t x_numerical\n');
    for i = 1:length(t)
        fprintf('%10.5f %18.12f\n',t(i),x(i));
        fprintf(fid,'%.16e %.16e\n',t(i),x(i));
    end
end
fclose(fid);
fprintf('Steps: %d\n',steps);
fprintf('Data file: %s\n',filename);

% Plot a dense exact curve and the actual numerical points.
figure;
if check_exact
    t_plot = linspace(t0,t_end,1000);
    x_plot = (t_plot+1).^2 + (x0-(t0+1)^2)*exp(t_plot-t0);
    plot(t_plot,x_plot,'k-','LineWidth',1.2);
    hold on;
    plot(t,x,'ko--','MarkerSize',4);
    hold off;
    legend('Exact',method,'Location','best');
else
    plot(t,x,'ko-','MarkerSize',4);
    legend(method,'Location','best');
end
xlabel('t');
ylabel('x(t)');
title(sprintf('%s, step = %g',method,dt));
grid on;

%==================== Equation and numerical methods ====================
function dx = f(t,x)
% Supply the derivative x' = f(t,x) here.
dx = x-t*t+1;
end

function xf = euler1(f,ti,xi,tf)
% One Euler step: slope at the beginning of the interval.
h = tf-ti;
xf = xi+h*f(ti,xi);
end

function xf = heun1(f,ti,xi,tf)
% Euler predictor, then ONE trapezoidal corrector.
h = tf-ti;
slope1 = f(ti,xi);
predicted = xi+h*slope1;
slope2 = f(tf,predicted);
xf = xi+0.5*h*(slope1+slope2);
end

function xf = rk4_scalar(f,ti,xi,tf)
% Classical RK4. Each k includes h.
h = tf-ti;
k1 = h*f(ti,xi);
k2 = h*f(ti+h/2,xi+k1/2);
k3 = h*f(ti+h/2,xi+k2/2);
k4 = h*f(tf,xi+k3);
xf = xi+(k1+2*k2+2*k3+k4)/6;
end