% Demo Kalman + Sequential Monte Carlo for linear gaussian model
% use prior proposal and locally optimal proposal

T=100;
sv=1;
sw=sqrt(1);
phi=0.95;

% simulate data according to the model
x=zeros(1,T);
y=zeros(1,T);
x=randn(1);
for k=2:T
    x(1,k)=phi*x(1,k-1)+sv.*randn(1);
end
y=x+sw.*randn(1,T);

% load simulatedstatesobs   data corresponding to lecture notes


% number of samples/particles
N=1000;

% exact inference using Kalman filter
mp=zeros(1,T);
mf=zeros(1,T);
vp=zeros(1,T);
vf=zeros(1,T);
my=zeros(1,T);
vy=zeros(1,T);
loglike=0;

mp(1,1)=0;
vp(1,1)=1;
my(1,1)=mp(1,1);
vy(1,1)=vp(1,1)+sw^2;
loglike=-0.5*log(2*pi*vy(1,1))-0.5*(y(1,1)-my(1,1))^2/vy(1,1);

for k=1:T
    
    % update
    vf(1,k)=sw^2*vp(1,k)/(vp(1,k)+sw^2);
    mf(1,k)=vf(1,k)*(mp(1,k)/vp(1,k)+y(1,k)/sw^2);
    
    % predict
    if (k<T)
        mp(1,k+1)=phi*mf(1,k);
        vp(1,k+1)=phi^2*vf(1,k)+sv^2;
        my(1,k+1)=mp(1,k+1);
        vy(1,k+1)=vp(1,k+1)+sw^2;
        loglike=loglike-0.5*log(2*pi*vy(1,k+1))-0.5*(y(1,k+1)-my(1,k+1))^2/vy(1,k+1);
    end
end



% prior proposal
xs1=zeros(T,N);
lw1=zeros(T,N);
w1=zeros(T,N);
wnorm1=zeros(T,N);
ess1=zeros(1,T);
xmean1=zeros(1,T);
xvar1=zeros(1,T);

% locally optimal proposal
xs2=zeros(T,N);
lw2=zeros(T,N);
w2=zeros(T,N);
wnorm2=zeros(T,N);
ess2=zeros(1,T);
xmean2=zeros(1,T);
xvar2=zeros(1,T);

% SMC using prior
for k=1:T
    if (k==1)
        % sample particles initial distribution
        xs1(k,:)=randn(1,N);
         % compute log weights for numerical stability
        lw1(k,:)=-0.5*log(2*pi*sw^2).*ones(1,N)-0.5*(y(1,k)*ones(1,N)-xs1(k,:)).^2./sw^2;
    else
        % propagate particles according to prior 
        xs1(k,:)=phi.*xs1(k-1,ind)+sv.*randn(1,N);
        % compute log weights for numerical stability
        lw1(k,:)=-0.5*log(2*pi*sw^2).*ones(1,N)-0.5*(y(1,k)*ones(1,N)-xs1(k,:)).^2./sw^2;
    end
    lmax=max(lw1(k,:));
    w1(k,:)=exp(lw1(k,:)-lmax);  % correct only up to a multiplicative factor for unnormalized weights
    wnorm1(k,:)=w1(k,:)./sum(w1(k,:));
    % effective sample size
    ess1(1,k)=1/sum(wnorm1(k,:).^2);
    % perform multinomial resampling 
    W = cumsum(wnorm1(k,:));
    W = W(:)';
    u = rand(N, 1);
    ind = sum(repmat(u, 1, N) > repmat(W, N, 1), 2) + 1;
    % compute estimate of E(Xk|y1:k)
    xmean1(1,k)=sum(xs1(k,ind))/N;
    % compute estimate of Var(Xk|y1:k)
    xvar1(1,k)=sum(xs1(k,ind).^2)/N-xmean1(1,k)^2;
end

% SMC using locally optimal

% variance locally optimal proposal
ss2=sw^2*sv^2/(sv^2+sw^2);
ss=sqrt(ss2);

for k=1:T
    if (k==1)
        % sample particles initial proposal; e.g p(x1|y1)
        xs2(k,:)=ss2*y(1,k)*ones(1,N)/sw^2+ss.*randn(1,N);
         % compute log weights for numerical stability
        lw2(k,:)=-0.5*log(2*pi*(sw^2+sv^2)).*ones(1,N)-0.5*(y(1,k)*ones(1,N)).^2./(sw^2+sv^2);
    else
        % propagate particles according to p(xk|xk-1,yk) 
        xs2(k,:)=ss2.*(phi.*xs2(k-1,ind)./sv^2+y(1,k)/sw^2)+ss.*randn(1,N);
        % compute log weights for numerical stability
        lw2(k,:)=-0.5*log(2*pi*(sw^2+sv^2)).*ones(1,N)-0.5*(y(1,k)*ones(1,N)-phi.*xs2(k-1,ind)).^2./(sw^2+sv^2);
    end
    lmax=max(lw2(k,:));
    w2(k,:)=exp(lw2(k,:)-lmax);  % correct only up to a multiplicative factor for unnormalized weights
    wnorm2(k,:)=w2(k,:)./sum(w2(k,:));
    % effective sample size
    ess2(1,k)=1/sum(wnorm2(k,:).^2);
     % perform multinomial resampling 
    W = cumsum(wnorm2(k,:));
    W = W(:)';
    u = rand(N, 1);
    ind = sum(repmat(u, 1, N) > repmat(W, N, 1), 2) + 1;
    % compute estimate of E(Xk|y1:k)
    xmean2(1,k)=sum(xs2(k,ind))/N;
    % compute estimate of Var(Xk|y1:k)
    xvar2(1,k)=sum(xs2(k,ind).^2)/N-xmean2(1,k)^2;
end

% display ESS
figure(1)
plot(ess1,'r')
hold on
plot (ess2,'b')
print essprioroptimalhighswsmc -depsc
axis([1 T 0 N])
% display variance log weights


% display absolute error exact conditional mean versus SIS
figure(3)
plot(abs(mf-xmean1),'r');
hold on
plot(abs(mf-xmean2),'b');
axis([1 T 0 0.2])
print errorxmeanprioroptkalmansmc -depsc

% display exact variance versus SIS
figure(4)
plot(xvar1,'r');
hold on
plot(xvar2,'b');
plot(vf,'c');
axis([1 T 0 1])
print vmeanprioroptkalmansmc -depsc

hold off