clear;
close all;

% Initialize seed
seed1 = 3455755; rand('seed', seed1); randn('seed',seed1);

T=100;
H = 1;
A = 0.95;
Omega = 5.00;
Sigma = 1.00;
betaBar = 0.5;
G = 1;

betaTrue = zeros(T,1);
betaTrue(1) = betaBar;
for t=2:T
  betaTrue(t) = betaBar + A*(betaTrue(t-1)-betaBar) + sqrt(Sigma)*randn;
end
z = H*(betaTrue-betaBar) + sqrt(Omega)*randn(T,1);

yHat_filter = zeros(T,1);
yHat = zeros(T,1);
yHat_smooth = zeros(T,1);

V = zeros(T,1);
V_smooth = zeros(T,1);

J = zeros(T,1);
K = zeros(T,1);
R = zeros(T,1);

y0Hat = 0;
V0 = Sigma/(1-A^2);

R(1) = A*V0*A' + G*Sigma*G';
K(1) = R(1)*H'*inv(Omega+H*R(1)*H');

yHat(1) = A*y0Hat + K(1)*(z(1)-H*A*y0Hat);
V(1) = (1-K(1)*H)*R(1);

for t=2:T
  yHat_filter(t) = A*yHat(t-1);
  
  R(t) = A*V(t-1)*A' + G*Sigma*G';
  K(t) = R(t)*H'*inv(Omega+H*R(t)*H');
  
  yHat(t) = A*yHat(t-1) + K(t)*(z(t)-betaBar - H*A*yHat(t-1));
  V(t) = (1-K(t)*H)*R(t);
end

yHat_smooth(T) = yHat(T);
V_smooth(T) = V(T);

for t=T-1:-1:1
  J(t) = V(t)*A'*inv(R(t+1));
  yHat_smooth(t) = yHat(t) + J(t)*(yHat_smooth(t+1)-yHat_filter(t+1));
  V_smooth(t) = V(t) + J(t)*(V_smooth(t+1)-R(t+1))*J(t)';
end

subplot(3,1,1);
plot(1:1:T,[yHat_filter+betaBar,betaTrue]);
title('Filtered beta with true beta')
subplot(3,1,2);
plot(1:1:T,[yHat+betaBar,betaTrue]);
title('Updated beta with true beta')
subplot(3,1,3);
plot(1:1:T,[yHat_smooth+betaBar,betaTrue]);
title('Smoothed beta with true beta')

figure;
subplot(3,1,1);
plot(1:1:T,[yHat_filter+betaBar,yHat_filter+betaBar+2*sqrt(R),...
            yHat_filter+betaBar-2*sqrt(R)]);
title('Filtered beta +/- 2 standard deviations')
subplot(3,1,2);
plot(1:1:T,[yHat+betaBar,yHat+betaBar+2*sqrt(V),yHat+betaBar-2*sqrt(V)]);
title('Updated beta +/- 2 standard deviations')
subplot(3,1,3);
plot(1:1:T,[yHat_smooth+betaBar,yHat_smooth+betaBar+2*sqrt(V_smooth),...
           yHat_smooth+betaBar-2*sqrt(V_smooth)]);
title('Smoothed beta +/- 2 standard deviations')