clear;clc;

% Parameters and constants
Qin = 7;
nbuckets = 3;
r = 2.5;
A = pi*r^2;
H = 30;
rh = 0.1;
Ah = pi*rh^2;
g = 981;
rho = 1;
mu = 0.01;
dt = 50;
total_time = 250;
t_vec = [0:round(total_time/dt)]';
time = t_vec*dt;

% Pre-allocate
h = nan(length(t_vec),nbuckets);
Qout = nan(length(t_vec),nbuckets);

% Reynolds number table
Re_table = log10([1 10 100 1000 10000 100000]);
CD_table = [0.04 0.28 0.73 0.74 0.64 0.61];

% Initial conditions
h(1,:) = 0; % First bucket empty at t=0
Qout(1,:) = 0;

for t = 2:length(t_vec)
    h_in = Qin/A;
    
    % Calculate height at the end of the time step. h2 is not h at t=2, it
    % is the height of the water inside the bucket at the end of the current
    % time step. Similarly, h(t-1) is the height at the beginning of the 
    % current time step ( As you can imagine, h(t-1) is h2 of the prior time step).
    % this is because our dt has a finite value (it's not infinitesimally small), 
    
    % Step 1 define an anonymous function where h2 is the variable we want
    % to find.
    h2_fun = @(h2) h2 - h(t-1) - dt*(h_in -...
                                    (interp1(Re_table,...
                                                CD_table,...
                                                max(log10((2*rh*rho)/mu *sqrt(2*g*h2)),1))...
                                    *Ah...
                                    *sqrt(2*g*h2))/A);
                                
    % Estimate an initial value for h2. We do not know it yet, that will be
    % the job of the fzero function, but we need to provide an estimated
    % value just to get the root finding algorithm an initial point. To do
    % this we will choose h(t-1)+h_in. In other words, our initial guess
    % just ignores the outflow.
    x0 = h(t-1)+h_in*dt; % Initial guess.
    
    % Find the h2 that makes the equation equal to zero.
    h(t) = fzero(h2_fun,x0);
    
    % Now we take the average of the heights
    h_mean = (h(t-1)+h(t))/2;
    
    % Once we have a new estimated h(t), the h_mean, we re-calculate h. 
    Re(t) = (2*rh*rho)/mu *sqrt(2*g*h_mean);
    CD(t) = max(interp1(Re_table,CD_table,log10(Re(t))),0);
    Qout(t) = CD(t)*Ah*sqrt(2*g*h_mean);
    h_out(t) = Qout(t)/A;
   
    h(t) = min(h(t-1)+dt*(h_in-h_out(t)),H);
end

% Plots
plot(time,h,'-k','LineWidth',1)
xlabel('Time (s)')
ylabel('Height (cm)')
ylim([0 H+5])
legend('Crank-Nicolson method')
legend boxoff
hold all


