%% Generating the system

nx = 5;
nu = 2;
ny = 2;

G = drss(nx,ny,nu);
G.d = 0;
% G = tf([1 0],[1 -1],'Ts',1)*tf(G);
% G = minreal(ss(G));
% nx = 5;

[A,B,C,D] = ssdata(G);

% A = [1.0 2; 3 4]
% B = [5.0 6; 7 8]
% C = [1.0 0; 0 1]

Gx = ss(A,B,eye(size(A)),0,'Ts',1);                    % model with states as explicit outputs


%% Generating the reference signal

N = 50;
r = zeros(ny,N+1);
r(1,floor(N/3):floor(2*N/3)) = 1;

rs = timeseries(r,0:N,'Name','Reference');     % for simulation in Simulink

rN = r(:,end);

%% Setting the optimal control weights

P = 1*diag(ones(ny,1));
Q = 10*diag(ones(ny,1));
R = 1*diag(ones(nu,1));

%% Solving the discrete-time Riccati equation in order to get the sequence 
%of state feedback and feedforward gains.

S = zeros(nx,nx,N+1);   % although I only need S of lenght N (starts at 1 and ends at N), I will make it longer for notational reasons 
v = zeros(nx,1,N+1);

K = zeros(nu,nx,N);
Kv = zeros(nu,nx,N);

S(:,:,end) = C'*P*C;
v(:,:,end) = C'*P*rN;

for k = N:-1:1          % because Matlab indexing start at 1, the Nth time actually corresponds to the (N-1)th time in the formulas 
    D = (B'*S(:,:,k+1)*B+R)\B';
    K(:,:,k) = D*S(:,:,k+1)*A;
    Kv(:,:,k) = D;
    S(:,:,k) = A'*S(:,:,k+1)*(A-B*K(:,:,k))+C'*Q*C;
    v(:,:,k) = (A-B*K(:,:,k))'*v(:,:,k+1) + C'*Q*r(:,k);
end

figure(1)
subplot(3,1,1)
plot(0:N,squeeze(S(1,1,:)),'.-')
xlabel('k')
ylabel('S_{11}')
subplot(3,1,2)
plot(0:N-1,squeeze(K(1,1,:)),'.-')
xlabel('k')
ylabel('K_{11}')
subplot(3,1,3)
plot(0:N-1,squeeze(Kv(1,1,:)),'.-')
xlabel('k')
ylabel('K^v_{11}')


%% Simulation
vs = timeseries(v(:,:,2:end),0:N-1,'Name','v'); % forward shift implemented this way
Ks = timeseries(K,0:N-1,'Name','K');    
Kvs = timeseries(Kv,0:N-1,'Name','Kv');

open('sim_LQ_tracking')
sim('sim_LQ_tracking')