% problem 4h_iii - RK4
clc, clear

% data
hi = [6,5,4,3,2,1,0];
Ai = [1.17,0.97,0.67,0.45,0.32,0.18,0].*10.^4;

% 5th degree least squares polynomial regression
regression_coef = polyfit(hi,Ai,5);
p5_regression = @(h) polyval(regression_coef, h);

% define parameters & inital condition
h0 = 6;
d = 0.25;
e = 1;
g = 9.81;
time_span = [0,19].*3600;

% define ODE
h_dot = @(t,h) -pi.*d.^2/(4.*p5_regression(h)).*sqrt(2.*g.*(h+e));

% solve using ode45
[t_ode45, h_ode45] = ode45(h_dot, time_span, h0);

% define final time
t_final = 68000; % little larger from part 4d

% define list of step sizes
dt_list = [1000,100,10,1,0.1];

% instantiate empty vectors to store step sizes and computation times
dt_record = zeros(1,length(dt_list));
pc_times = zeros(1,length(dt_list));

% instantiate tiled plot
f1 = figure();
set(f1, 'Position', [50, 50, 1000, 800]);
tlo = tiledlayout(3, 2);
title(tlo, 'Predictor-Corrector Errors', '(RK-4)',...
    'FontSize', 18, 'Interpreter', 'latex');

% main loop over all step sizes
for j = 1:length(dt_list)
    % get step size and time vector
    dt = dt_list(j);
    t = 0:dt:t_final;

    % get first 4 initial conditions using RK4
    % instantiate empty vectors to store h
    h_pc = zeros(1,length(t));
    % initial condition
    h_pc(1) = h0;
    % loop RK4 3 times for next 3 initial conditions
    for i = 1:3
        k1 = dt.*h_dot(t(i), h_pc(i));
        k2 = dt.*h_dot(t(i)+dt/2, h_pc(i)+k1/2);
        k3 = dt.*h_dot(t(i)+dt/2, h_pc(i)+k2/2);
        k4 = dt.*h_dot(t(i)+dt, h_pc(i)+k3);
    
        h_pc(i+1) = h_pc(i) + (1/6).*(k1 + 2.*k2 + 2.*k3 + k4);
    end
    
    % loop for predictor-corrector method
    tic;
    for k = 4:length(t)-1
        f_k = h_dot(t(k),h_pc(k));
        f_km1 = h_dot(t(k-1),h_pc(k-1));
        f_km2 = h_dot(t(k-2),h_pc(k-2));
        f_km3 = h_dot(t(k-3),h_pc(k-3));

        % predictor (using 4-step explicit)
        h_pred = h_pc(k) + dt/24.*(55.*f_k - 59.*f_km1 + 37.*f_km2 - 9.*f_km3);
        
        % corrector (using 4-step implicit)
        f_pred = h_dot(t(k+1),h_pred);
        h_pc(k+1) = h_pc(k) + dt/24 * (9*f_pred + 19*f_k - 5*f_km1 + f_km2);
        
        % if pond is empty, stop
        if h_pc(k+1) <= 0
            h_pc(k+1:end) = NaN;
            break;
        end
    end
    pc_times(j) = toc;
    dt_record(j) = dt;

    % interpolate ode45 to match over time vectors
    h_ode45_interp = interp1(t_ode45, h_ode45, t);

    % find indices where pond isn't empty (singularity produces plot divergence)
    valid_idx = find(h_ode45_interp >= 0);
    h_ode45_interp_valid = h_ode45_interp(valid_idx);

    % compute abs errors
    abs_err = abs(h_pc(valid_idx) - h_ode45_interp_valid);

    % compute rel errors
    rel_err = abs_err ./ abs(h_ode45_interp_valid);

    % plot
    t_valid = t(valid_idx);
    nexttile;
    
    yyaxis left;
    plot(t_valid, abs_err, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title(sprintf("dt = %.1f", dt),...
    "FontSize", 16, "Interpreter", "latex");
    xlabel("time $t$ [s]", "FontSize", 16, "Interpreter", "latex");
    legend('Absolute Error', 'Relative Error', 'Interpreter', 'latex', 'Location', 'best');
end

%saveas(f1, "4h_iii_rk4_errors", "fig");
%saveas(f1, "4h_iii_rk4_errors", "epsc");

% plot computational times vs step size
f2 = figure();
loglog(dt_record, pc_times, "Linewidth", 3, "Color", "g");
hold on;
set(gca, "FontSize", 12);
grid;
title("\textbf{Computation Time vs Step Size}",...
"FontSize", 16, "Interpreter", "latex");
xlabel("Step Size $dt$ [s]", "FontSize", 16, "Interpreter", "latex");
ylabel("Computation Time [s]", "FontSize", 16, "Interpreter", "latex");

%saveas(f2, "4h_iii_rk4_comp_time_vs_step_size", "fig");
%saveas(f2, "4h_iii_rk4_comp_time_vs_step_size", "epsc");