% This program calculates and plots the supply bulk PDF and CDF, p_b(s) and P_b(s), and the supply tail PDF and CDF, p_t(s) and P_t(s),
% based on the estimated parameters.
% It then calculates the combined supply distribution, p_s(s) and P_s(s).
% It then calculates the distribution of water volume, p_v(s) and P_v(s).
% Extension of pb_pt_plt003.m
% 9.2.2023.

Vmax = 137500; % Maximum volume in reservoir [MCM]; greater volume overflows.
Vc = 34700; % Critical volume in reservoir, below which water is not withdrawn.
V1 = 121875; % Initial volume of water in reservoir.
t_max = 40; % 10; % Final time step (in years) at which to calculate volume distribution.

wtbeta = 66745.1384347884; % Estimated scale parameter of the logistic bulk distribution.
wtgamma = 4252.02672073450; % Estimated shape parameter of the logistic bulk distribution.

wtsigma = 2665.50314611762; % Estimated scale parameter of the EVD.
wtmu = 52825.3072977901; % Estimated location parameter of the EVD.

dprime = 55500; % This is the discharge value in MCM, when the reservoir volume exceeds the critical value.

iprintp = 0; % iprintp = 1 causes printing of p_b and p_t; any other values skips the printing.

fprintf(' logistic bulk parameters  wtbeta    wtgamma ')
[wtbeta    wtgamma ]
fprintf(' EVD parameters wtsigma   wtmu')
[wtsigma   wtmu]

alpha = 0.045;
q_alfa_b = -wtgamma*log((1/alpha) - 1) + wtbeta; % alpha quantile of the bulk distribution.

q_alfa_t = wtmu + wtsigma* log(log(1/(1-alpha))); % alpha quantile of the tail distribution.

Delta_q = q_alfa_b - q_alfa_t;

fprintf(' alpha, q_alfa_b')
formatSpec = '%10.5f  %10.7f\n';
fprintf(formatSpec, alpha, q_alfa_b')

fprintf(' alpha, q_alfa_t')
formatSpec = '%10.5f  %10.7f\n';
fprintf(formatSpec, alpha, q_alfa_t')

% [alpha q_alfa_b]

% CALCULATE BULK AND TAIL DISTRIBUTIONS OF SUPPLY:

s1 = 35000;
s2 = 90000;
num_s = 300;
s_vec = linspace(s1,s2,num_s);

zz = exp(-(s_vec - wtbeta)/wtgamma);
p_b = zz./(wtgamma*(1 + zz).^2); % PDF of bulk distribution.
P_b = 1./(1 + zz); % CDF of bulk distribution.

yy_vec = exp((s_vec - wtmu)/wtsigma);
p_t = yy_vec.*exp(-yy_vec)/wtsigma; % Vector of tail PDF values.
P_t = 1 - exp(-yy_vec); % Vector of tail PDF values.

% CALCULATE COMBINED DISTRIBUTION OF SUPPLY:

p_s = ones(num_s,1);
P_s = ones(num_s,1);

for i = 1 : num_s
  p_s(i) = p_t(i);
  P_s(i) = P_t(i);
  if s_vec(i) > q_alfa_t
    Delta_q_idx = i + Delta_q*num_s/(s2 - s1);
    i_shift = min(num_s, ceil(Delta_q_idx));
    p_s(i) = p_b(i_shift);
    P_s(i) = P_b(i_shift);
  end
end

if iprintp == 1
  fprintf('   s/10e8    p_t      p_b     p_s    P_s')
  [ s_vec'/10^8   p_t'      p_b'    p_s'    P_s']
end

% CALCULATE DISTRIBUTION OF WATER VOLUME IN RESERVOIR FOR t = 2:

vol1 = 0;
vol2 = Vmax;
num_vol = 300; % We will calculate p_v for num_vol volume values from vol1 to vol2.
vol_vec = linspace(vol1,vol2,num_vol); % Vector of MCM volume values at which p_v and P_v are calculated.
p_v = zeros(num_vol,t_max);
P_v = zeros(num_vol,t_max);
for iv = 1 : num_vol
  vol = (iv - 1)*Vmax/(num_vol-1); % Current volume [MCM] at which to calculate p_v.
  d2 = 0;
  if vol > Vc
    d2 = dprime;
  end
  if vol - V1 + d2 <= 0
    p_v(iv,2) = 0;
  else
    idx1 = min(ceil((vol - V1 + d2)*iv/Vmax), num_vol); % Index argument for p_s in numerator.
    p_v(iv,2) = p_s(idx1); % Unnormalized numerator of p_v.
  end
end
dvol = Vmax/num_vol;
sum_p_v = sum(p_v(:,2))*dvol;
p_v(:,2) = p_v(:,2)/sum_p_v; % Normalization of p_v(:,2).
% Check normalization of p_v:
sum_p_v = sum(p_v(:,2))*dvol;
fprintf(' Normalization check on p_v(:,2)')
formatSpec = '%10.8f\n';
fprintf(formatSpec, sum_p_v)
% Calculate P_v at t=2:
P_v(1,2) = p_v(1,2)*dvol;
for iv = 2 : num_vol
  P_v(iv,2) = P_v(iv-1,2) + p_v(iv,2)*dvol;
end

% CALCULATE DISTRIBUTION OF WATER VOLUME IN RESERVOIR  FOR t = 3, ... t_max:

for it = 3 : t_max
  for iy = 1 : num_vol % Loop on y in p_v(y,d,t), calculating the numerator integral:
    y_vol = (iy - 1)*Vmax/(num_vol-1); % Current volume, y, [MCM] at which to calculate p_v(y,d,t-1).
    dt = 0;  % Discharge [MCM].
    if y_vol > Vc
      dt = dprime;
    end
    % Calculate numerator integral for p_v(y,it) at current y value. That is, integrate on x in numerator:
    sum = 0;
    for ix = 1 : num_vol
      x_vol = (ix - 1)*Vmax/(num_vol-1); % Current volume, x, [MCM] in argument of p_v in numerator.
      if y_vol - x_vol + dt <= 0
        p_s_val = 0;
      else
        idx1 = min(ceil((y_vol - x_vol + dt)*ix/Vmax), num_vol); % Index argument for p_s in numerator.
        p_s_val = p_s(idx1); % Value of p_s in numerator.
        sum = sum + p_v(ix,it-1)*p_s_val*dvol; % Increment in numerator integral on RHS of p_v(y,t).
      end
    end  % End loop on integral on x in numerator of p_v(y,d,t).
    p_v(iy,it) = sum; % Unnormalized value. This is the integral in the numerator.
  end % End loop on y in p_v(y,d,t), which is not yet normalized.
% Normalize p_v(y,d,t):
  sum = 0;
  for iy = 1 : num_vol
    sum = sum + p_v(iy,it)*dvol;
  end % End normalization loop.
  p_v(:,it) = p_v(:,it)/sum; % Normalized value of p_v(iy,it) at current it value.
  % Check normalization of p_v:
%  it
%     sum_p_v = sum(p_v(:,it))*dvol;
  sum_p_v = 0;
  for isum = 1 : num_vol
    sum_p_v = sum_p_v + p_v(isum,it)*dvol;
  end
  fprintf(' Normalization check on p_v(:t)')
  formatSpec = '%10.2f  %10.8f\n';
  fprintf(formatSpec, it, sum_p_v)
  % Calculate P_v at t:
  P_v(1,it) = p_v(1,it)*dvol;
  for iv = 2 : num_vol
    P_v(iv,it) = P_v(iv-1,it) + p_v(iv,it)*dvol;
  end
end   % End loop on it.


% PLOT FIGURES:

num_alfa = 20;
alfa_vec_x = linspace(0,6e4,num_alfa);
alfa_vec_y = 0.0425*ones(num_alfa,1);

maxpb = max(p_b);
maxpt = max(p_t);
maxpbt = max(maxpb,maxpt);
maxpbt_fac = 0.2;
num_q = 20;
q_vec_yt = linspace(0,maxpbt_fac*maxpbt,num_q);
q_vec_xt = q_alfa_t*ones(num_q,1);
q_vec_yb = linspace(0,maxpbt_fac*maxpbt,num_q);
q_vec_xb = q_alfa_b*ones(num_q,1);

figure
plot(q_vec_xb,q_vec_yb,':k',q_vec_xt,q_vec_yt,':k',s_vec,p_t,'-k',s_vec,p_b,'--k','Linewidth',2)

figure
plot(alfa_vec_x,alfa_vec_y,':k',s_vec,P_t,'-k',s_vec,P_b,'--k','Linewidth',2)

figure
plot(s_vec,p_s,'-k','Linewidth',2)

figure
plot(alfa_vec_x,alfa_vec_y,':k',s_vec,P_s,'-k','Linewidth',2)

figure
plot(vol_vec,p_v(:,2),'-k','Linewidth',2)

figure
plot(vol_vec,P_v(:,2),'-k','Linewidth',2)

% for it = 3 : t_max
%   figure
%   plot(vol_vec,p_v(:,it),'-k','Linewidth',2)
%   title('it = 3')
%   figure
%   plot(vol_vec,P_v(:,it),'-k','Linewidth',2)
% end

figure
plot(vol_vec,p_v(:,3),'-k',vol_vec,p_v(:,4),'--k',vol_vec,p_v(:,5),'-.k','Linewidth',1)
title('it = 3, 4, 5')

figure
plot(vol_vec,P_v(:,3),'-k',vol_vec,P_v(:,4),'--k',vol_vec,P_v(:,5),'-.k','Linewidth',1)
title('it = 3, 4, 5')


figure
plot(vol_vec,p_v(:,10),'-k',vol_vec,p_v(:,20),'--k',vol_vec,p_v(:,30),'-.k',vol_vec,p_v(:,40),':k','Linewidth',1)
title('it = 10, 20, 30, 40')

figure
plot(vol_vec,P_v(:,10),'-k',vol_vec,P_v(:,20),'--k',vol_vec,P_v(:,30),'-.k',vol_vec,P_v(:,40),':k','Linewidth',1)
title('it = 10, 20, 30, 40')