% this is a demo for particle MCMC where we use the bootstrap particle filter

rng(10)   % use rng('shuffle') to initialize the seed to a different value each time you run this script

% load exactly the same data produced with the NimbleSMC R package
% those were produced with a = 1 and b=1.
yobs = load('yobs.dat');


% Set up some parameters
N = 1000;          % Number of particles 
T = length(yobs);            % Length of data record, assumes integer observational times
x0 = 0; % assumed deterministic initial condition (you can change this and make it random)

%::::::::: PARTICLE MCMC ::::::::::::::::::::::::::::::
R = 2000; % number of MCMC iterations
burnin = 200; % will be used when ploting histograms
MCMC = zeros(R,2); % will´be used to store posterior draws for "a" and "b"

a0=0.1; b0=2.5; % arbitrary starting parameters
theta_start = [a0,b0];
MCMC(1,:) = theta_start;
loglik_start = bpf(yobs, a0, b0, x0, N); % starting loglikelihood
theta_old = theta_start;
loglik_old = loglik_start;


for r=2:R
    theta_prop = mvnrnd(theta_old,[0.1^2,0.1^2]); % notice for a multivariate Gaussian, the covariance matrix has *variances* on its diagonal, not standard deviations
    a_prop = theta_prop(1);
    b_prop = theta_prop(2);
    loglik = bpf(yobs, a_prop, b_prop, x0, N);
    u = rand; % sample from the uniform U(0,1)
    % We work everything out on logscale to avoid numerical issues.
    % Also notice I am not coding the proposal g() function since the Gaussian is symmetric. 
    % Also notice the prior function is encoded further down, before the bootstrap filter
    if log(u) < min(0,(loglik - loglik_old) + log(prior(theta_prop)) - log(prior(theta_old)))  % recall: the proposal function is Gaussian so it simplifies out 
       % accept and store proposal theta_prop 
       MCMC(r,:) = theta_prop;
       theta_old = theta_prop; % will be used at next iteration
       loglik_old = loglik;    % will be used at next iteration
    else
       MCMC(r,:) = theta_old; 
    end
end

figure
subplot(1,2,1)
plot(MCMC(:,1))
title('MCMC chain for a (incl. burnin)')
subplot(1,2,2)
plot(MCMC(:,2))
title('MCMC chain for b (incl. burnin)')

% plot marginal posteriors by removing burnin
figure
subplot(1,2,1)
histogram(MCMC(burnin:end,1))
title('approx. \pi(a|y_{1:T})')
subplot(1,2,2)
histogram(MCMC(burnin:end,2))
title('approx. \pi(b|y_{1:T})')

    
function out = prior(theta)
    a = theta(1);
    b = theta(2);
    prior_a = normpdf(a,0.5,1);
    prior_b = normpdf(b,1.5,0.5);
    out = prior_a*prior_b; % I decide to assume that the priors for a and b are independent
                           % hence the joint prior is the product of the
                           % individual priors
end

   
function [loglik,x] = bpf(y, a, b, x0, N)
% Bootstrap particle filter
% Input:
%   y - measurements
%   a, and b - model parameters
%   x0 - initial condition of the latent state
%   N - number of particles
% Output:
%   loglik - the estimated loglikelihood value at input parameters
%   x      - states from an approximation to the filtering distribution

T = length(y);
x = zeros(N, T); % Particles (preallocated to zero)
w = zeros(N, T); % Weights (preallocated to zero)
loglik = 0;  % initialize loglikelihood

for(t = 1:T)
    if t==1 
       x(:,1) = a*x0 + randn(N,1);  % will add x0 to N draws from the standard Gaussian
    else
        ind = resampling(w(:,t-1)); % multinomial resampling of particle indeces
        % propagate forward ONLY selected particles having indeces inside ind 
        x(:,t) = a*x(ind,t-1) + randn(N,1);
    end

    % Compute weights (log-scale for numerical stability) using the
    % (log)density of the observations y ~ N(b*x,0.3^2)
    logweights = -1/2*log(2*pi*0.3^2) - 1/(2*0.3^2)*(y(t) - b*x(:,t)).^2;
    
   % NOTICE: the three commented lines below show how things shuld be PROPERLY done, for numerical stability. But are less readable
   % const = max(logweights); % Subtract below the maximum value for numerical stability
   % weights = exp(logweights-const); % this way it is "less likely" to obtain a numerical underflow/overflow
   % loglik = loglik + const + log(sum(weights)) - log(N);  
    weights = exp(logweights); % unnormalised weights
    loglik = loglik + log(sum(weights)) - log(N); 
    w(:,t) = weights/sum(weights); % Save the normalized weights
end

end