%
%
%                       SIR_Saltelli 
%
%

  clear all

  global N

%
% Set the initial conditions and number in the population.
%
  S0 = 900;
  I0 = 100;
  R0 = 0;
  Y0 = [S0; I0; R0];
  N = 1000;

%
% Construct the time interval.
%

  tf = 5;
  dt = 0.01;
  t_data = 0:dt:tf;
  ind = length(t_data);

%
%  Construct the matrices used to contruct the Saltelli estimators for the Sobol indices.
%

  M = 10;       % Number of random samples
  alpha = 2;    % parameters used for the interaction rate
  beta = 7;
  A = rand(M,4);
  A(:,2) = betarnd(alpha,beta,M,1);
  B = rand(M,4);
  B(:,2) = betarnd(alpha,beta,M,1);
  C1 = B;
  C2 = B;
  C3 = B;
  C4 = B;
  C1(:,1) = A(:,1);
  C2(:,2) = A(:,2);
  C3(:,3) = A(:,3);
  C4(:,4) = A(:,4);

  factor = 1e-3;    % Factor used to scale the magnitude of the numerator and denominator.

  ode_options = odeset('RelTol',1e-6);     % Set the ode relative tolerance

%
%  Solve the system M times.
%

% 
% Integrate the system using ode45.m.
%

  h = waitbar(0,'Please wait...');
  for j=1:M
    params = A(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_A(j) = factor*sum(dt*Y(:,3));
    y_A_end(j) = Y(ind,3);

    params(4) = 0.5;
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_A_end_fixed(j) = Y(ind,3);
    
    params = B(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_B(j) = factor*sum(dt*Y(:,3));

    params = C1(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_C1(j) = factor*sum(dt*Y(:,3));

    params = C2(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_C2(j) = factor*sum(dt*Y(:,3));

    params = C3(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_C3(j) = factor*sum(dt*Y(:,3));

    params = C4(j,:);
    [t,Y] = ode45(@SIR_rhs,t_data,Y0,ode_options,params);
    y_C4(j) = factor*sum(dt*Y(:,3));
    waitbar(j/M)
  end
  close(h)

%
% Compute and report the Sobol indices.
%

  f02 = ((1/M)^2)*sum(y_A)*sum(y_B);
  S1 = ((1/M)*y_A*y_C1' - f02)/((1/M)*y_A*y_A' - f02);
  S2 = ((1/M)*y_A*y_C2' - f02)/((1/M)*y_A*y_A' - f02);
  S3 = ((1/M)*y_A*y_C3' - f02)/((1/M)*y_A*y_A' - f02);
  S4 = ((1/M)*y_A*y_C4' - f02)/((1/M)*y_A*y_A' - f02);

  ST1 = 1 - ((1/M)*y_B*y_C1' - f02)/((1/M)*y_A*y_A' - f02);
  ST2 = 1 - ((1/M)*y_B*y_C2' - f02)/((1/M)*y_A*y_A' - f02);
  ST3 = 1 - ((1/M)*y_B*y_C3' - f02)/((1/M)*y_A*y_A' - f02);
  ST4 = 1 - ((1/M)*y_B*y_C4' - f02)/((1/M)*y_A*y_A' - f02);

  S = [S1 S2 S3 S4]
  ST = [ST1 ST2 ST3 ST4]

  x_dens = 0:40:1000;
  dens = ksdensity(y_A_end,x_dens);
  dens_fixed = ksdensity(y_A_end_fixed,x_dens);

  Sval = Y(:,1);
  Ival = Y(:,2);
  Rval = Y(:,3); 

%
% Plot the solutions.
%

  figure(1)
  plot(t,Sval,'b',t,Ival,'k--',t,Rval,'r-.','linewidth',3)
  set(gca,'Fontsize',[20]);
  legend('Susceptible','Infected','Recovered','Location','Northeast')
  xlabel('Time') 
  ylabel('Number of Individuals')

  figure(2)
  plot(x_dens,dens,'r',x_dens,dens_fixed,'r--','linewidth',3)
  set(gca,'Fontsize',[20]);
  xlabel('Number of Recovered Individuals')
  ylabel('PDF')
  legend(' Random \gamma, k, r, \delta',' Random \gamma, k, r; Fixed \delta=0.5','Location','Northeast')
%  legend(' Random \gamma, k, r, \delta',' Random r, \delta; Fixed \gamma=0.2, k=0.5','Location','South')

  x = 0:.01:1;
  y1 = betapdf(x,2,7);
  y2 = betapdf(x,0.2,15);
  figure(3)
  plot(x,y1,'k',x,y2,'--b','linewidth',3)
  set(gca,'Fontsize',[20]);
  legend('Beta(2,7)','Beta(0.2,15)','Location','Northeast')