% problem 4f, 4g
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);

% construct symbolic h', h'', and h''' for TS methods
syms hsym
degree = length(regression_coef) - 1;
poly_sym = 0;
for k = 1:length(regression_coef)
    poly_sym = poly_sym + regression_coef(k) * hsym^(degree - (k-1));
end

h_prime_sym = -pi.*d^2/(4.*poly_sym).*sqrt(2.*g.*(hsym+e));
h_2prime_sym = diff(h_prime_sym, hsym);
h_2prime = matlabFunction(h_2prime_sym, 'Vars', hsym);
h_3prime_sym = diff(h_2prime_sym, hsym);
h_3prime = matlabFunction(h_3prime_sym, 'Vars', hsym);

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

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

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

    % instantiate empty solution h array
    h_TS1 = zeros(1, length(t));
    h_TS2 = zeros(1, length(t));
    h_TS3 = zeros(1, length(t));
    h_modEuler = zeros(1, length(t));
    h_midRK = zeros(1, length(t));
    h_RK4 = zeros(1, length(t));
    h_RK5 = zeros(1, length(t));
    
    % use initial condition
    h_TS1(1) = h0;
    h_TS2(1) = h0;
    h_TS3(1) = h0;
    h_modEuler(1) = h0;
    h_midRK(1) = h0;
    h_RK4(1) = h0;
    h_RK5(1) = h0;
    
    % 1st-order Taylor series approximation (forward Euler)
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_TS1(i) <= 0
            h_TS1(i+1:end) = NaN;
            break;
        end
        h_TS1(i+1) = h_TS1(i) + dt.*h_dot(t(i), h_TS1(i));
    end
    
    % 2nd order TS approximation
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_TS2(i) <= 0
            h_TS2(i+1:end) = NaN;
            break;
        end
        h_TS2(i+1) = h_TS2(i) + dt .* h_dot(t(i), h_TS2(i)) + 0.5.*dt.^2.*h_2prime(h_TS2(i)).*h_dot(t(i), h_TS2(i));
    end
    
    % 3rd order TS approx
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_TS3(i) <= 0
            h_TS3(i+1:end) = NaN;
            break;
        end
        h_TS3(i+1) = h_TS3(i) + dt.*h_dot(t(i), h_TS3(i)) ...
            + 0.5.*dt^2.*h_2prime(h_TS3(i)).*h_dot(t(i), h_TS3(i)) ...
            + (1/6).*dt^3.*(h_3prime(h_TS3(i)).*h_dot(t(i), h_TS3(i)).^2 ...
            + h_2prime(h_TS3(i)).^2.*h_dot(t(i), h_TS3(i)));
    end
    
    % modified Euler
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_modEuler(i) <= 0
            h_modEuler(i+1:end) = NaN;
            break;
        end
        k1 = h_dot(t(i), h_modEuler(i));
        k2 = h_dot(t(i)+dt, h_modEuler(i)+dt.*k1);
        h_modEuler(i+1) = h_modEuler(i) + (dt/2).*(k1 + k2);
    end
    
    % midpoint RK
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_midRK(i) <= 0
            h_midRK(i+1:end) = NaN;
            break;
        end
        k1 = dt.*h_dot(t(i), h_midRK(i));
        k2 = dt.*h_dot(t(i)+dt/2, h_midRK(i)+k1/2);
        h_midRK(i+1) = h_midRK(i) + k2;
    end
    
    % RK4
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_RK4(i) <= 0
            h_RK4(i+1:end) = NaN;
            break;
        end
        k1 = dt.*h_dot(t(i), h_RK4(i));
        k2 = dt.*h_dot(t(i)+dt/2, h_RK4(i)+k1/2);
        k3 = dt.*h_dot(t(i)+dt/2, h_RK4(i)+k2/2);
        k4 = dt.*h_dot(t(i)+dt, h_RK4(i)+k3);
        h_RK4(i+1) = h_RK4(i) + (1/6).*(k1 + 2.*k2 + 2.*k3 + k4);
    end
    
    % RK5
    for i = 1:length(t)-1
        % if pond is empty, stop
        if h_RK5(i) <= 0
            h_RK5(i+1:end) = NaN;
            break;
        end
        k1 = dt.*h_dot(t(i), h_RK5(i));
        k2 = dt.*h_dot(t(i)+dt/4, h_RK5(i)+k1/4);
        k3 = dt.*h_dot(t(i)+dt/4, h_RK5(i)+k1/8+k2/8);
        k4 = dt.*h_dot(t(i)+dt/2, h_RK5(i)-k2/2+k3);
        k5 = dt.*h_dot(t(i)+3.*dt/4, h_RK5(i)+(3/16).*k1+(9/16).*k4);
        k6 = dt.*h_dot(t(i)+dt, h_RK5(i)-(3/7).*k1+(2/7).*k2+(12/7).*k3-(12/7).*k4+(8/7).*k5);
        h_RK5(i+1) = h_RK5(i) + (1/90).*(7.*k1 + 32.*k3 + 12.*k4 + 32.*k5 + 7.*k6);
    end

    % 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_TS1 = abs(h_TS1(valid_idx) - h_ode45_interp_valid);
    abs_err_TS2 = abs(h_TS2(valid_idx) - h_ode45_interp_valid);
    abs_err_TS3 = abs(h_TS3(valid_idx) - h_ode45_interp_valid);
    abs_err_modEuler = abs(h_modEuler(valid_idx) - h_ode45_interp_valid);
    abs_err_midRK = abs(h_midRK(valid_idx) - h_ode45_interp_valid);
    abs_err_RK4 = abs(h_RK4(valid_idx) - h_ode45_interp_valid);
    abs_err_RK5 = abs(h_RK5(valid_idx) - h_ode45_interp_valid);

    % compute rel errors for each method
    rel_err_TS1 = abs_err_TS1 ./ abs(h_ode45_interp_valid);
    rel_err_TS2 = abs_err_TS2 ./ abs(h_ode45_interp_valid);
    rel_err_TS3 = abs_err_TS3 ./ abs(h_ode45_interp_valid);
    rel_err_modEuler = abs_err_modEuler ./ abs(h_ode45_interp_valid);
    rel_err_midRK = abs_err_midRK ./ abs(h_ode45_interp_valid);
    rel_err_RK4 = abs_err_RK4 ./ abs(h_ode45_interp_valid);
    rel_err_RK5 = abs_err_RK5 ./ 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(4,2);
    title(tlo, sprintf('Errors for Step Size h = %.2f', dt), 'FontSize', 18);
    
    nexttile;
    yyaxis left;
    plot(t_valid, abs_err_TS1, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_TS1) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_TS1, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_TS1) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{1st-Order TS}",...
    "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_TS2, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_TS2) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_TS2, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_TS2) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{2nd-Order TS}",...
    "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_TS3, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_TS3) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_TS3, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_TS3) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{3rd-Order TS}",...
    "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_modEuler, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_modEuler) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_modEuler, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_modEuler) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{Modified Euler}",...
    "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_midRK, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_midRK) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_midRK, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_midRK) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{Midpoint RK (RK2)}",...
    "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_RK4, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_RK4) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_RK4, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_RK4) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{RK4}",...
    "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_RK5, "Linewidth", 2, "Color", "r");
    ylabel("Absolute Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(abs_err_RK5) * 1]);
    ax = gca;
    ax.YColor = 'r';
    hold on;
    yyaxis right;
    plot(t_valid, rel_err_RK5, "Linewidth", 2, "Color", "b");
    ylabel("Relative Error", "FontSize", 16, "Interpreter", "latex");
    ylim([0 max(rel_err_RK5) * 1]);
    ax = gca;
    ax.YColor = 'b';
    hold off;
    set(gca, "FontSize", 12);
    grid;
    title("\textbf{RK5}",...
    "FontSize", 16, "Interpreter", "latex");
    xlabel("time $t$ [s]", "FontSize", 16, "Interpreter", "latex");
    legend('Absolute Error', 'Relative Error', 'Interpreter', 'latex', 'Location', 'best');
end