%% Setup the problem
Ts = 0.2; % Sampling time for the discretization

% Define system model and discretize it,
% Pgen _* is the transfer function matrix of the general plant model,
% and Pmass_* is the transfer function from u to p2.
[Pgen_c, Pmass_c] = define_system;
Pgen_d = c2d(Pgen_c, Ts);
Pmass_d = c2d(Pmass_c, Ts);

% Find matrices T1, T2, T3 %%% 
% We would like to find stable T1, T2, T3 such that the set of all 
% stable closed loop transfer functions are given by
% {T1 + T2*Q*T3 | Q stable}     (*)
%
% If P is unsable things are bit more involved, but essentially one needs to
% close the loop with a stabilizing controller and then augment the
% controller with the Q-factor. If this can be done one obtains the form
% (*) above
% This is handled by the following function which finds T1, T2, T3
% and data for the nominal controller so that the actual, Q-augmented,
% controller can be recovered
[T1,T2,T3,Knom_data]= qcvx_q_parametrization( Pgen_d , 3, 3, 1 );

%% Create Q_basis 
NQ = 10; % Number of terms 1, z^-1, z^-2, z^-(N-1) for each element of Q
Qsize = [1 2];
Q_basis = setup_Q_basis(Ts, Qsize, NQ);

%% Setup constraints and cost function
% Define grids t and w over which time and frequency contraints are evaluated
n_sample = 170;
t = Ts*(0:n_sample-1)';
w = linspace(0,(2*pi/Ts)/2,200)';

Smax = 1.3; % Upper bound on sensitivity function
umin = -6; % Lower bound on control signal
umax = 6; % Upper bound on clontrol signal

% Cost function
weight_e = 0.5; % gamma, control error weighting
weight_u = 1.0; % rho, control signal weighting

% Constraints on the step response from r to p1
max_overshoot = 1.02; % Maximum allowed overshoot
t1 = 3.6; % Maximum allowed 63% rise time
t2 = 19; % Time when settling constraints start
T_settling = 4.5; % Required settling time constant
% Lower bound on p1, contstrains rise time and settling time
lb_p1_refstep = (t < t1)*-0.05 + (t >= t1 & t < t2)*0.63 + (t >= t2).*(1 - (max_overshoot-1)*exp(-(t-t2)/T_settling));
% Upper bound on p1, constrains over shoot and settling time
ub_p1_refstep = (t < t2)*max_overshoot + (t >= t2).*(1 + (max_overshoot-1)*exp(-(t-t2)/T_settling));

%% Setup and solve the convex optimization problem    
cvx_begin
cvx_solver sedumi

variable b(numel(Q_basis))
[cl_stepresp, cl_freqresp] = qcvx_compute_cl_response(T1, T2, T3, Q_basis, b, t, w);

minimize weight_u*sum_square( cl_stepresp(:,2,1) ) + weight_e*sum_square( cl_stepresp(:,3,1) )

subject to:
lb_p1_refstep <= cl_stepresp(:,1,1) <= ub_p1_refstep
umin <= cl_stepresp(:,2,1) <= umax
abs( cl_freqresp(2, 3, :) ) <= Smax

cvx_end

optimal_cost = cvx_optval;
fprintf('Optimal cost: %.2f\n', optimal_cost)

%% Compute the Q-factor
% Compuate the transfer fucntion Q according to the optimal b-vector
% that cvx puts out after it has solved the problem
%
% Watch out for the following way of forming Q as it gives numerical
% problems, (step response r->p1 not within bounds)
%
% b = [prob.Q{1}.num{:} prob.Q{2}.num{:}];
% Q = tf(zeros(1,2));
% for k=1:numel(Q_basis)
%     Q = Q + minreal(Q_basis{k}*b(k));
% end
%
% This is an example of numerical problems arising when working with the
% Q-parametrization

Q = ss([tf(b(1:NQ)', 1, 0.2, 'Variable', 'z^-1') tf(b(NQ+1:2*NQ)', 1, 0.2, 'Variable', 'z^-1')]);


%% Recover the controller
% K1 is the controller rom r -> u
% K2 is the controller rom y -> u
[K1d, K2d, ~] = qcvx_recover_controller(Knom_data, Q);

%% Compute the closed loop system
Gcl = T1 + minreal(T2*Q*T3);

%% Plot the results

% Plot stepresponses from reference to p1 and u
figure(28)
clf
subplot(211)
plot(t, [lb_p1_refstep ub_p1_refstep], 'k--') 
hold on
plot(t, step(Gcl(1,1), t), 'b')
title('Step resonse r -> p1')

subplot(212)
plot(t([1 end]), umin*[1 1], 'k--', t([1 end]), umax*[1 1], 'k--') 
hold on
plot(t, step(Gcl(2,1), t), 'b') 
title('Step response r -> u')
set(gca, 'Ylim', umax*[-1.1, 1.1])
xlabel('Time [s]')


% Plot the sensitivity function
figure(29)
clf
loglog(w, abs( squeeze(freqresp(Gcl(2,3), w)) ), 'k')
hold on
loglog(w([2 end]), [1.3 1.3], 'k--')


% Plot the sensitivity function over a denser grid w.
% The problem that we solved only constrained the sensitivty function to be
% below the constraint at the grid points.
% If the order of the Q-factor is high and the frequency grid w is too sparse,
% then there might be violations of the constraint between the grid
% points in w
w_dense = linspace(0,(2*pi/Ts)/2,1000);
loglog(w_dense, abs( squeeze(freqresp(Gcl(2,3), w_dense)) ), 'b')
title('Sensitivity function')
xlabel('Frequency [rad/s]')


%% Plot the Nyquist curve
P_fr =  squeeze(freqresp(Pmass_d, w_dense));
K_fr =  squeeze(freqresp(K2d, w_dense));

figure(30)
plot(-P_fr .* K_fr, 'b') % Minus, since postive feedback is assumed for general plant formulation
hold on
plot(-1 + exp(1i*(0:0.01:7)) / Smax, 'g--')
plot(-1, 0, 'rx', 'MarkerSize', 10)
axis(0.8*[-2.5 2 -2 2])
grid on
xlabel('Re')
ylabel('Im')

title('Nyquist Plot')
