%
% Chapter 3, Example 3.14, pages 118-122
% Optimal Reference Shaping for Dynamical Systems
% Theory and Applications
% -Tarunraj Singh
%
% Design of LQR control with the first state penalized
clear

M = eye(2);
K = [1 -1;-1 1];

A = [zeros(2) eye(2);-inv(M)*K zeros(2)];
B = [0;0;inv(M)*[1;0]];
C = [eye(2) zeros(2)];
D = [0;0];

Q = zeros(4);
Q(1,1) = 1;
R = 1;

klqr = lqr(A,B,Q,R);

x0 = [1;1;0;0];
Tf = 15;

sys = ss(A-B*klqr,B,C,D);
[Y,T,X] = initial(sys,x0,Tf);
figure(1)
subplot(211), plot(T,Y(:,1),T,Y(:,2),'r--');
xlabel('Time (sec)')
ylabel('Displacement');
legend('Mass 1','Mass 2');
U = -X*klqr';
subplot(212), plot(T,U);
xlabel('Time (sec)')
ylabel('Control Input');

% Design of Frequency Shaped LQR control
% input to this filter is the state x_1 of the original system
% and the output is the 
[Af,Bf,Cf,Df] = tf2ss([1],[1 0 2]);

Aa = [A zeros(4,2)
      [zeros(2,2) Bf zeros(2,1)] Af];

Ba = [B;[0;0]];
Ca = [C zeros(2,2)];
Da = D;

Df_temp = zeros(1,4);   % to permit multiplcation with the entire x vector
  
Qa = [Df_temp'*Df_temp  Df_temp'*Cf
      Cf'*Df_temp       Cf'*Cf     ];
Qa(1,1) = 1;
  
Ra = R;
     
klqr_freq = lqr(Aa,Ba,Qa,Ra);

sysa = ss(Aa-Ba*klqr_freq,Ba,Ca,Da);
x0a = [x0;0;0];

[Ya,Ta,Xa] = initial(sysa,x0a,Tf);
figure(2);
subplot(211), plot(Ta,Ya(:,1),Ta,Ya(:,2),'r--');
xlabel('Time (sec)')
ylabel('Displacement');
legend('Mass 1','Mass 2');
Ua = -Xa*klqr_freq';
subplot(212), plot(Ta,Ua);
xlabel('Time (sec)')
ylabel('Control Input');

% Residual energy of the two systems
for indx = 1:length(T)
    Res(indx) = X(indx,1:2)*K*X(indx,1:2)';% + X(indx,3:4)*M*X(indx,3:4)';
end;
for indx = 1:length(Ta)
    Res_a(indx) = Xa(indx,1:2)*K*Xa(indx,1:2)';% + Xa(indx,3:4)*M*Xa(indx,3:4)';
end;
figure(3)
plot(T,Res,Ta,Res_a,'r--');
xlabel('Time (sec)')
ylabel('Potential Energy');
legend('LQR','Frequency Shaped LQR');
figure(4)
[Mag,Phase,W] = bode(sys);
[Mag_a,Phase_a] = bode(sysa,W);
[n1,n2,n3] = size(Mag);
Mag = reshape(Mag,n1,n3);
[n1,n2,n3] = size(Mag_a);
Mag_a = reshape(Mag_a,n1,n3);
subplot(211), loglog(W,Mag(1,:),W,Mag_a(1,:),'r--')
xlabel('Frequency (rad/sec)');
ylabel('Magnitude');
legend('LQR','Frequency Shaped LQR');
subplot(212), loglog(W,Mag(2,:),W,Mag_a(2,:),'r--')
xlabel('Frequency (rad/sec)');
ylabel('Magnitude');
legend('LQR','Frequency Shaped LQR');

