%% Discrete-time LQR by Prof. Ashish R. Hota, EE, IITKGP

nx = 3; nu = 1;
A = [0 1 0;0 0 1;-1 -2 -3];
B = [0;0;1];
Q = eye(3);
R = 10000;
SN =[1 0 0; 0 2 0; 0 0 3];
N = 20;

%% Offline Computation - Store Sequence of State Feedback Gain

S = zeros(nx,nx,N+1);
K = zeros(nu,nx,N);
S(:,:,end) = SN;
for k=1:N
    K(:,:,N-k+1) = inv(R+B'*S(:,:,N-k+2)*B)*B'*S(:,:,N-k+2)*A;
    S(:,:,N-k+1) = A'*S(:,:,N-k+2)*(A-B*K(:,:,N-k+1))+Q;
end

%% Offline Computation - Solve DARE to obtain Steady State Gain

[S_inf,K_inf,eig_inf] = idare(A,B,Q,R,[],[]);

%% Specify Initial State

x0 = [50;-50;70];

%% Optimization Formulation

u = sdpvar(nu,N);
x = sdpvar(nx,N+1);
cost1 = 0.5*x(:,end)'*SN*x(:,end);
for i=1:N
    cost1 = cost1 + 0.5*x(:,i)'*Q*x(:,i) + 0.5*u(:,i)'*R*u(:,i);
end
constraint = x(:,1) == x0;
for i=1:N
    constraint = [constraint,x(:,i+1) == A*x(:,i)+B*u(:,i)];
end
options = sdpsettings('solver','mosek','verbose',1);
solution = optimize(constraint,cost1,options);
solution.problem
% Extract Optimal Input and State Trajectory
u_opt = value(u);
x_opt = value(x);

%% Online Deployment of Time-Varying LQR Solution

u_lqr = zeros(nu,N);
x_lqr = zeros(nx,N+1);
x_lqr(:,1) = x0;
for k=1:N
    u_lqr(:,k) = -K(:,:,k)*x_lqr(:,k);
    x_lqr(:,k+1) = A*x_lqr(:,k) + B*u_lqr(:,k);
end

%% Online Deployment of Steady-State LQR Solution

x_lqr_ss = zeros(nx,N+1);
x_lqr_ss (:,1) = x0;
for k=1:N
    x_lqr_ss(:,k+1) = (A-B*K_inf)*x_lqr_ss(:,k);
end
u_lqr_ss = -K_inf*x_lqr_ss(:,1:end-1);

%% Plot all the states
%plot(1:1:N+1,x_lqr_ss'-x_opt','LineWidth',1.5);
plot(1:1:N+1,x_lqr','LineWidth',1.5);
legend('x_1(k)','x_2(k)','x_3(k)')
xlabel('Time (k)') 
ylabel('States') 
%% Plot optimal input trajectory
%plot(1:1:N,[u_opt'-u_lqr',u_lqr'-u_lqr_ss'],'LineWidth',1.5);
plot(1:1:N,u_lqr','LineWidth',1.5);
xlabel('Time (k)') 
ylabel('Inputs') 