% Excercise 4: Kalman filter and predictor
%   for pendulum dynamical system

% z(t+1) = A z(t) + noise(U)
% x(t)   = C z(t)   + noise(V)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Ndata　の値を変えて試して見よ
% 時間ステップ数
Ndata = 50;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% 単振動振子のダイナミクスを記述
dt = 0.1;                                % Delta t
g  = 9.8;                                % Gravity
r  = 1.0;                                % Radius
o2 = g/r;                                % Omega^2
a0 = sqrt(1-dt^2*o2);                    % Dumping factor for stabilization
A  = [a0 dt; -dt*o2 a0];                 % Transition Matrix
C  = [1 0];                              % Observation Matrix

ss = size(A,1);                          % state size
os = size(C,1);                          % observation size
U  = 0.8*eye(ss);                        % System Noise Cov. Matrix
V  = 5*eye(os);                          % Observation Noise Cov. Matrix

initz = [10 0]';                        % Expected Value of Initial State
initQ = 10*eye(ss);                     % Cov. of Initial State

% 上で定義した状態空間モデルに従い、ノイズの入ったデータを生成
% x は内部状態、y は観測値
seed = 9;
rand('state', seed);
randn('state', seed);
T     = Ndata;
[z,x] = sample_lds(A, C, U, V, initz, T);

% 観測値 x とその他のパラメータを用いて
% 1.フィルタリング
% 2.連続予測
[zfilt, Qfilt] = kalman_filter(x, A, C, U, V, initz, initQ);
[zpred, Qpred] = kalman_predict(A, C, U, V, initz, initQ, T);

% 1. 2. それぞれについて平均二乗誤差を計算
dfilt = z([1 2],:) - zfilt([1 2],:);
mse_filt = sqrt(sum(sum(dfilt.^2)))

dpred = z([1 2],:) - zpred([1 2],:);
mse_pred = sqrt(sum(sum(dpred.^2)))

% 1. 2. それぞれについてプロット
subplot(2,1,1)
hold on
plot(z(1,:), 'ks-');
plot(x(1,:), 'g*');
plot(zfilt(1,:), 'rx:');
plot(zfilt(1,:)+reshape(Qfilt(1, 1, :),1,T),'b')
plot(zfilt(1,:)-reshape(Qfilt(1, 1, :),1,T),'b')
hold off
legend('true', 'observed', 'filtered mean', 'filtered var', 0)
title(sprintf('Filtering: mse=%f', mse_filt))
xlabel('t')
ylabel('x')

subplot(2,1,2)
hold on
plot(z(1,:), 'ks-');
plot(x(1,:), 'g*');
plot(zpred(1,:), 'rx:');

axis(axis);
plot(zpred(1,:)+reshape(Qpred(1, 1, :),1,T),'b')
plot(zpred(1,:)-reshape(Qpred(1, 1, :),1,T),'b')
hold off
legend('true', 'observed', 'predicted mean', 'predicted var', 0)
title(sprintf('Iterated predict: mse=%f', mse_pred))
xlabel('t')
ylabel('x')