%                                   Example12_5.m                  
%
%
%  This code is also used in Example 12.13.
%
%  Requires SIR_fun2.m, SIR_rhs2.m, SIR2_integer.txt
%
%

  clear all
%  close all

  global N Y0

%
%  Load data which consists of 141 points of infectious data along with the
%  associated times.
%

  load SIR2_integer.txt

%
%  Set parameters and initial conditions.  These values are close to the flu data
%  but with larger I0 and smaller variance to prevent negative values of I(t). 
%

  S0 = 760;
  I0 = 12;
  R0 = 0;
  N = S0 + I0 + R0;

  beta = 0.0022;         % Estimated values for flu data: beta = 2.3040e-3;
  gamma = 0.4470;        %                                gamma = 0.4284;
  beta_mean = beta;
  gamma_mean = gamma;
  params = [beta gamma];

  h = 1e-16;
  tf = 14;
  dt = 0.1;
  t_data = SIR2_integer(:,1);
  I_data = SIR2_integer(:,2);
  t_data = 0:dt:tf;
  Y0 = [S0; I0];

  n = length(t_data);
  p = 2;
  Nchain = 1.2e+4;
  Nchain_red = 1e+4;

%
% Specify the parameters and compute the solution I(t).  When integrating
% the system, we take advantage of the fact that R(t) = N - S(t) - I(t).
%
% Construct synthetic data having the observation variance sigma^2 = 81;
% Note that this is significantly less than the observation variance 
% sigma^2 = 322.8 estimated for the flu data.
%

  q_opt = [beta gamma];
  ode_options = odeset('RelTol',1e-8);
  [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_opt);
  Y(:,3) = N - Y(:,1) - Y(:,2);
  I = Y(:,2);

  sigma = 9;
  sigma2 = sigma^2;
  error = sigma*randn(size(t_data'));
  %I_data = I + error;                     

  %SIR_data = [t_data' I_data];
  %SIR_data_integer = [t_data' I_data_integer];
  %save('SIR2.txt','SIR_data','-ascii')
  %save('SIR2_integer.txt','SIR_data_integer','-ascii')

%
%  Plot the data and model 
%

  figure(1)
  plot(t,Y(:,1),'k-',t,Y(:,2),'r--',t,Y(:,3),'b-.','linewidth',3)
  axis([0 tf 0 800])
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('Responses S(t), I(t), R(t)')
  legend('Susceptible S(t)','Infectious I(t)','Recovered R(t)','Location','East')
%  print -depsc SIR2_responses

  figure(2)
  plot(t,I,'k-',t,I_data,'bx',t,0*t,'k-','linewidth',3)
  axis([0 tf 0 350])
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('Infectious I(t)')
  legend('Model','Data','Location','NorthEast')
%  print -depsc SIR2_model_data

%
% Estimate the parameters beta and gamma and compute the resulting trajectories I(t)
%

  ode_options = odeset('RelTol',1e-6);

  q_init = [beta gamma];
  modelfun = @(q)SIR_fun2(q,Y0,t_data,I_data);
  [q_opt,fval] = fminsearch(modelfun,q_init);

  beta_opt = q_opt(1);
  gamma_opt = q_opt(2);

  [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_opt);
  Y(:,3) = N - Y(:,1) - Y(:,2);
  I = Y(:,2);

%
% Use the complex-step derivative approximation to compute the sensitivity matrix and
% covariance matrix.
%

  Y0 = [S0; I0];
  for i=1:p
    if i==1
      beta_complex = complex(beta_opt,h);
      q = [beta_complex,gamma];
      [t,Ypert] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q);
      I_beta = imag(Ypert(:,2))/h;
      Sens(:,i) = I_beta;
    else
      gamma_complex = complex(gamma_opt,h);
      q = [beta,gamma_complex];
      [t,Ypert] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q);
      I_gamma = imag(Ypert(:,2))/h;
      Sens(:,i) = I_gamma;
    end
  end

  Res = I_data - I;
  sigma2_est = (1/(n-p))*Res'*Res
  V = sigma2_est*inv(Sens'*Sens);

%
% Plot the residuals and 2-sigma intervals
%

  figure(3)
  plot(t,Res,'xb',t,0*t,'k-','linewidth',3)
  axis([0 tf -25 25]);
  hold on
  set(gca,'Fontsize',[22]);
  plot(t,2*sigma*ones(length(t)),'--k','linewidth',3)
  plot(t,-2*sigma*ones(length(t)),'--k','linewidth',3)
  hold off
  xlabel('Time (Days)')
  ylabel('Residuals')
%  print -depsc SIR2_residuals

%
%  Construct the parameters used when sampling the error variance in the Metropolis algorithm.
%

  SS_old = (I-I_data)'*(I-I_data);
  n0 = .001;
  sigma02 = sigma2_est;
  aval = 0.5*(n0 + n);
  bval = 0.5*(n0*sigma02 + SS_old);
  sigma2 = 1/gamrnd(aval,1/bval);

  q_old = [beta;gamma];
  R = chol(V);

%
% Construct a Metropolis chain of length Nchain
%

  for i = 1:Nchain
    z = randn(2,1);
    q_new = q_old + R'*z;
    beta_new = q_new(1,1);
    gamma_new = q_new(2,1);
    [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_new);
    I_new = Y(:,2);
    u_alpha = rand(1);
    SS_new = (I_new-I_data)'*(I_new-I_data);
    term = exp(-.5*(SS_new-SS_old)/sigma2);
    Term(i) = term;
    alpha = min(1,term);
    if u_alpha < alpha
      Q_MCMC(:,i) = [beta_new; gamma_new];
      q_old = q_new;
      SS_old = SS_new;
    else
      Q_MCMC(:,i) = q_old;
    end
    Sigma2(i) = sigma2;
    bval = 0.5*(n0*sigma02 + SS_old);
    sigma2 = 1/gamrnd(aval,1/bval);
  end

  beta_vals = Q_MCMC(1,:);
  gamma_vals = Q_MCMC(2,:);


%
% Use kde to construct marginal densities for beta and r.  Use samples 2e+3 to 10e+3
% to ensure burn-in.
%

  Nchain_lower = 2000;
  beta_vals(1:Nchain_lower) = [];
  gamma_vals(1:Nchain_lower) = [];

  range_beta = max(beta_vals) - min(beta_vals);
  range_gamma = max(gamma_vals) - min(gamma_vals);
  beta_min = min(beta_vals)-range_beta/10;
  beta_max = max(beta_vals)+range_beta/10;
  gamma_min = min(gamma_vals)-range_gamma/10;
  gamma_max = max(gamma_vals)-range_gamma/10;
  %[bandwidth_beta,density_beta,beta_mesh,cdf_beta]=kde(beta_vals);
  %[bandwidth_gamma,density_gamma,gamma_mesh,cdf_gamma]=kde(gamma_vals);

  [density_beta,beta_mesh] = ksdensity(beta_vals);
  [density_gamma,gamma_mesh] = ksdensity(gamma_vals);

%
%  Construct the posterior density by directly approximating Bayes'rule using
%  a left Riemann sum.  We store the samples in SS_quad2 to avoid having to
%  evaluate the ODE system in the 4-layer nested for-loop used to evaluate
%  Bayes' formula. 
%

  Nq_beta = 100;
  Nq_gamma = 100;
  beta_min = 2.15e-3;
  beta_max = 2.25e-3;
  hq_beta = (beta_max - beta_min)/Nq_beta;
  gamma_min = 0.435;
  gamma_max = 0.46;
  hq_gamma = (gamma_max - gamma_min)/Nq_gamma;
  for i = 1:Nq_beta
    beta = beta_min + (i-1)*hq_beta;
    beta_posterior_axis(i) = beta;
    for j=1:Nq_gamma
      gamma = gamma_min + (j-1)*hq_gamma;
      gamma_posterior_axis(j) = gamma;
      q_quad = [beta gamma];
      [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_quad);
      I_quad = Y(:,2);
      SS_quad2(i,j) = (I_quad-I_data)'*(I_quad-I_data); 
    end
  end

  for i1=1:Nq_beta
    for i2=1:Nq_gamma
      SS2 = SS_quad2(i1,i2);
      for j1=1:Nq_beta
        for j2=1:Nq_gamma
          quad_val2(j1,j2) = exp(-.5*(SS_quad2(j1,j2)-SS2)/sigma2_est);
        end
      end  
      joint_posterior(i1,i2) = 1/(hq_beta*hq_gamma*sum(sum(quad_val2)));
    end
  end
  
  marginal_gamma = hq_beta*sum(joint_posterior);
  marginal_beta = hq_gamma*sum(joint_posterior');


%
%  This loop evaluates the marginal posterior distribution for beta with gamma fixed.
%
%  r = 0.4469;
%  for i = 1:Nq_beta
%    beta = beta_min + (i-1)*hq_beta;
%    q_quad = [beta gamma];
%    [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_quad);
%    I_quad = Y(:,2);
%    SS = (I_quad-I_data)'*(I_quad-I_data);
%    for j=1:Nq_beta
%      beta2 = beta_min + (j-1)*hq_beta;
%      q_quad2 = [beta2 gamma];
%      [t,Y] = ode45(@SIR_rhs2,t_data,Y0,ode_options,q_quad2);
%      I_quad2 = Y(:,2);
%      SS_quad = (I_quad2-I_data)'*(I_quad2-I_data);
%      quad_val(j) = exp(-0.5*(SS_quad - SS)/sigma2_est);
%    end
%    posterior(i) = 1/(hq_beta*sum(quad_val));
%  end


%
% Construct the Gaussian sampling distributions for beta and r.
%

  sample_dist_beta = normpdf(beta_posterior_axis,beta_mean,sqrt(V(1,1)));
  %sample_dist_beta = normpdf(beta_posterior_axis,0.002196,sqrt(V(1,1)));
  sample_dist_gamma = normpdf(gamma_posterior_axis,gamma_mean,sqrt(V(2,2)));

%
% Plot results.  Note that the axis limits may change if one uses a different data set.
%

  figure(4)
  plot(beta_vals,'b-','linewidth',1)
  set(gca,'Fontsize',[22]);
  axis([0 Nchain_red 2.17e-3 2.23e-3])
  xlabel('Chain Iteration')
  ylabel('Parameter \beta')
%  print -depsc SIR2_beta_chain

  figure(5)
  plot(gamma_vals,'b-','linewidth',1)
  set(gca,'Fontsize',[22]);
  axis([0 Nchain_red .44 .455])
  xlabel('Chain Iteration')
  ylabel('Parameter \gamma')
%  print -depsc SIR2_gamma_chain

  figure(6)
  plot(Sigma2,'b-')
  axis([0 Nchain_red 40 120])
  set(gca,'Fontsize',[22]);
  xlabel('Chain Iteration')
  ylabel('Obs Error Variance \sigma^2')
%  print -depsc SIR2_sigma2_chain

  figure(7)
  plot(beta_mesh,density_beta,'k--','linewidth',3)
  axis([2.17e-3 2.225e-3 0 6.0e+4])
  hold on
  plot(beta_posterior_axis,marginal_beta,'b-','linewidth',3)
  hold off
  set(gca,'Fontsize',[22]);
  legend('Metropolis','Bayes','Location','NorthEast')
  xlabel('Parameter \beta')
  ylabel('PDF')
%  print -depsc SIR2_marginal_beta
  

  figure(8)
  plot(gamma_mesh,density_gamma,'k--','linewidth',3)
  axis([0.440 0.455 0 200])
  hold on
  plot(gamma_posterior_axis,marginal_gamma,'b-','linewidth',3)
  hold off
  set(gca,'Fontsize',[22]);
  legend('Metropolis','Bayes','Location','NorthEast')
  xlabel('Parameter \gamma')
  ylabel('PDF')
%  print -depsc SIR2_marginal_gamma

  figure(9)
  scatter(beta_vals,gamma_vals,'b')
  box on
  axis([2.17e-3 2.23e-3 0.44 0.457])
  set(gca,'Fontsize',[22]);
  xlabel('Parameter \beta')
  ylabel('Parameter \gamma')
%  print -depsc SIR2_pairwise

  figure(10)
  plot(beta_posterior_axis,marginal_beta,'b-','linewidth',3)
  axis([2.168e-3 2.225e-3 0 6.0e+4])
  hold on
  plot(beta_posterior_axis,sample_dist_beta,'k--','linewidth',3)
  hold off
  set(gca,'Fontsize',[22]);
  legend('Posterior','OLS','Location','Northeast')
  xlabel('Parameter \beta')
  ylabel('PDF')
%  print -depsc SIR2_Bayes_OLS_beta

  figure(11)
  plot(gamma_posterior_axis,marginal_gamma,'b-','linewidth',3)
  axis([0.438 0.456 0 200])
  hold on
  plot(gamma_posterior_axis,sample_dist_gamma,'k--','linewidth',3)
  hold off
  set(gca,'Fontsize',[22]);
  legend('Posterior','OLS','Location','Northeast')
  xlabel('Parameter \gamma')
  ylabel('PDF')
%  print -depsc SIR2_Bayes_OLS_gamma