%% Simulate dynamics 

% Point mass in 1D 
% x(k) = A x(k-1) + B u(k-1) + v
% y(k) = C x(k) + w


% Dynamics parameters
k = 0; % spring  N/m
b = 0; % damping N*s/m
M = 4; % mass N*s^2/m = kg
V = 0.05*eye(2);  % std of state noise
W = 0.1*eye(2); % std of sensor noise

DT = 0.01; % time step (sec)
 
% Dynamics Equations
A = [0 1; -k/M -b/M]*DT + eye(2);
B = [0; 1/M];
C = eye(2);

% Simuate the dynamics
N = 25;
x0 = [1;0];

% presample all noise so we can reuse noise in later simulations
Vser = V*randn(2,N);
Wser = W*randn(2,N);

% run dynamics
clear X Y 
X = x0;
Y = x0 + Wser(:,1);
for i=2:N,
    X(:,i) = A*X(:,i-1) + Vser(:,i);
    Y(:,i) = C*X(:,i) +  Wser(:,i);
end

figure(1);% clf;

subplot(2,3,1); hold on; title('position - no control')
plot(X(1,:),'b')
plot(Y(1,:),'r')
plot([0 N],[0 0],'k:')

subplot(2,3,4); hold on; title('velocity - no control')
plot(X(2,:),'b')
plot(Y(2,:),'r')
plot([0 N],[0 0],'k:')


%% LQR

QN = diag([1 1]);
Qk = zeros(2,2);
%Qk = diag([1 0]);
R  = 0.01;


clear P F
P{N} = QN;
for k=N:-1:2,
    F{k} = inv(R+B'*P{k}*B)*B'*P{k}*A;
    Q = Qk;
    % if(k==floor(N/2)) Q=QN; end
    P{k-1} = Q + A'*(P{k}-P{k}*B*inv(R+B'*P{k}*B)*B'*P{k})*A;
end


% Now control with no measurement noise using LQR
clear Xlqr Ulqr
Xlqr = x0;
Ylqr = x0 + Wser(:,1);
for k=2:N,
    Ulqr(:,k) = -F{k}*Xlqr(:,k-1);
    Xlqr(:,k) = A*Xlqr(:,k-1) + B*Ulqr(:,k) + Vser(:,k);
end


subplot(2,3,2); hold on; title('position - LQR')
plot(Xlqr(1,:),'b')
plot(B(1)*Ulqr(1,:),'k')
plot([0 N],[0 0],'k:')

subplot(2,3,5); hold on; title('velocity - LQR')
plot(Xlqr(2,:),'b')
plot(B(2)*Ulqr(1,:),'k')
plot([0 N],[0 0],'k:')





%% KF

clear Pprior K Ppost
Pprior{1} = W;
for k=1:N,
    K{k} = Pprior{k}*C*inv(C*Pprior{k}*C'+ W);
    Ppost{k} = (eye(2)-K{k}*C)*Pprior{k};
    Pprior{k+1} = A*Ppost{k}*A' + V;
end


% Now control with measurement noise using KF+LQR = LQG
clear Xlqg Ylqg Ulqg Xhat
Xlqg = x0;
Ylqg = x0 + Wser(:,1);
Xhat = Ylqg;
for k=2:N,
    % LQR
    Ulqg(:,k) = -F{k}*Xhat(:,k-1);
    Xlqg(:,k) = A*Xlqg(:,k-1) + B*Ulqg(:,k) + Vser(:,k);
    Ylqg(:,k) = C*Xlqg(:,k) +  Wser(:,k);

    % KF
    xprior = A*Xhat(:,k-1) + B*Ulqg(:,k);
    innov = Ylqg(:,k-1) - C*xprior;
    Xhat(:,k) = xprior + K{k}*innov;
end


subplot(2,3,3); hold on; title('position - LQR')
plot(Xlqg(1,:),'b')
plot(Xhat(1,:),'b:')
plot(B(1)*Ulqg(1,:),'k')
plot(Ylqg(1,:),'r')
plot([0 N],[0 0],'k:')

subplot(2,3,6); hold on; title('velocity - LQR')
plot(Xlqg(2,:),'b')
plot(Xhat(2,:),'b:')
plot(B(2)*Ulqg(1,:),'k')
plot(Ylqg(2,:),'r')
plot([0 N],[0 0],'k:')


