m = 5;
b = 2;

dt = 0.2;

Tfinal = 10;
data_length = Tfinal / dt;

%

A = [-b/m 0; 1 0]
B = [1/m; 0]
H = [0 1]
R = 3
Q = diag([0.001 0.0001])
Ad = eye(size(A)) + dt * A
Bd = dt * B
u = zeros(1, data_length);
u(1:data_length/2) = 2;
u(data_length/2:end) = -2;
plot(u)
ylim([-2.5 2.5])
X = zeros(2, data_length);

for i = 1:data_length-1
    w = diag(Q) .* randn(2, 1);
    X(:,i+1) = Ad * X(:,i) + Bd * u(i) + w;
end

%

v = R * randn(size(u));
Z = H * X + v;
plot(Z, 'b.')
X_pred = zeros(size(X));
X_update = zeros(size(X));

P = diag([2, 6]);

X_update(:,1) = [0; 1];

for i = 2:data_length
    X_pred(:, i) = Ad * X_update(:,i-1) + Bd * u(i);
    P = Ad * P * Ad' + Q;
    y = Z(i) - H * X_pred(:,i);
    S = H * P * H' + R;
    K = P' * H' * inv(S);
    X_update(:,i) = X_pred(:,i) + K * y;
    P = (eye(size(A)) - K * H) * P;
end

%

figure
hold on
plot(Z, '.')
plot(X_update(2,:))
plot(X(2,:), ':')
legend("Measurements", "Estimate", "True Position")