% problem 4h_ii - 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));
explicit_times = zeros(1,length(dt_list));
implicit_times = zeros(1,length(dt_list));

% 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 t and h
    t_RK4 = zeros(1,4);
    h_RK4 = zeros(1,4);
    % initial condition
    t_RK4(1) = 0;
    h_RK4(1) = h0;
    
    % loop RK4 3 times for next 3 initial conditions
    for i = 1:3
        k1 = dt.*h_dot(t_RK4(i), h_RK4(i));
        k2 = dt.*h_dot(t_RK4(i)+dt/2, h_RK4(i)+k1/2);
        k3 = dt.*h_dot(t_RK4(i)+dt/2, h_RK4(i)+k2/2);
        k4 = dt.*h_dot(t_RK4(i)+dt, h_RK4(i)+k3);
    
        h_RK4(i+1) = h_RK4(i) + (1/6).*(k1 + 2.*k2 + 2.*k3 + k4);
        t_RK4(i+1) = t_RK4(i) + dt;
    end
    
    % instantiate empty solution h array
    h_explicit = zeros(1, length(t));

    % use initial condition from RK4
    h_explicit(1:4) = h_RK4;

    % loop for explicit 4-step method (Adams-Bashforth)
    tic;
    for k = 4:length(t)
        f_k = h_dot(t(k),h_explicit(k));
        f_km1 = h_dot(t(k-1),h_explicit(k-1));
        f_km2 = h_dot(t(k-2),h_explicit(k-2));
        f_km3 = h_dot(t(k-3),h_explicit(k-3));

        h_explicit(k+1) = h_explicit(k) + dt/24.*(55.*f_k - 59.*f_km1 + 37.*f_km2 - 9.*f_km3);
        
        % if pond is empty, stop
        if h_explicit(k+1) <= 0
            h_explicit(k+1:end) = NaN;
            break;
        end
    end
    explicit_times(j) = toc;

    % instantiate empty solution h array
    h_implicit = zeros(1, length(t));

    % use initial condition from RK4
    h_implicit(1:4) = h_RK4;

    % set tolerance and max iterations for Newton's method
    tol = 10^(-8);
    max_iter = 100;

    % loop for implicit 4-step method (Adams-Moulton)
    tic;
    for k = 4:length(t)-1
        % known values
        f_k = h_dot(t(k),h_implicit(k));
        f_km1 = h_dot(t(k-1),h_implicit(k-1));
        f_km2 = h_dot(t(k-2),h_implicit(k-2));

        % use explicit step as initial guess
        f_km3 = h_dot(t(k-3),h_implicit(k-3));
        h_guess = h_implicit(k) + dt/24.*(55.*f_k - 59.*f_km1 + 37.*f_km2 - 9.*f_km3);
        
        % use fixed-point iteration instead of Newton's method
        h_next = h_guess;
        for iter = 1:max_iter
            h_new = h_implicit(k) + dt/24*(9*h_dot(t(k+1), h_next) + 19*f_k - 5*f_km1 + f_km2);

            if abs(h_new - h_next) < tol
                h_next = h_new;
                break;
            end
            h_next = h_new;
        end

        % save value
        h_implicit(k+1) = h_next;

        % if pond is empty, stop
        if h_implicit(k+1) <= 0
            h_implicit(k+1:end) = NaN;
            break;
        end
    end
    implicit_times(j) = toc;
    dt_record(j) = dt;  % store step size for plotting against computational times

    % 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 for each method
    abs_err_explicit = abs(h_explicit(valid_idx) - h_ode45_interp_valid);
    abs_err_implicit = abs(h_implicit(valid_idx) - h_ode45_interp_valid);

    % compute rel errors for each method
    rel_err_explicit = abs_err_explicit ./ abs(h_ode45_interp_valid);
    rel_err_implicit = abs_err_implicit ./ abs(h_ode45_interp_valid);

    % plot
    t_valid = t(valid_idx);
    f1 = figure();
    set(f1, 'Position', [50, 50, 1200, 850]);  % [left, bottom, width, height]
    tlo = tiledlayout(1,2);
    title(tlo, sprintf('Errors for Step Size h = %.2f', dt), 'FontSize', 18);

    nexttile;
    yyaxis left;
    plot(t_valid, abs_err_explicit, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_explicit) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_explicit, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_explicit) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{4-Step Explicit Method}",...
    "FontSize", 16, "Interpreter", "latex");
    xlabel("time $t$ [s]", "FontSize", 16, "Interpreter", "latex");
    legend('Absolute Error', 'Relative Error', 'Interpreter', 'latex', 'Location', 'best');

    nexttile;
    yyaxis left;
    plot(t_valid, abs_err_implicit, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_implicit) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_implicit, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_implicit) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{4-Step Implicit Method}",...
    "FontSize", 16, "Interpreter", "latex");
    xlabel("time $t$ [s]", "FontSize", 16, "Interpreter", "latex");
    legend('Absolute Error', 'Relative Error', 'Interpreter', 'latex', 'Location', 'best');
end

% plot computational times vs step size for each method
f2 = figure();
loglog(dt_record, explicit_times, "Linewidth", 3, "Color", "g");
hold on;
loglog(dt_record, implicit_times, "Linewidth", 3, "Color", "m");
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");
legend('Explicit', 'Implicit', 'Interpreter', 'latex', 'Location', 'best');