%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% continuous time mass spring system
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

m  = 4;
nx = 2*m;
nu = 3; % in {1,..,m}

Ac = zeros(nx, nx);
for ii=1:m Ac(ii,m+ii) = 1.0; end
for ii=1:m Ac(m+ii,ii) = -2.0; end
for ii=1:m-1 Ac(m+ii,ii+1) = 1.0; end
for ii=1:m-1 Ac(m+ii+1,ii) = 1.0; end

Bc = zeros(nx, nu);
for ii=1:nu Bc(m+ii,ii) = 1.0; end

Cc = eye(nx);
Dc = zeros(nx, nu);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% discretization
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% continuout time state space system
sysc = ss(Ac, Bc, Cc, Dc);

% sampling time
Ts = 0.5;

% convert to discrete time
sys = c2d(sysc, Ts);

%sys
%sysd

% extract discrete time state space matrices
A = sys.A
B = sys.B

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% simulate the discrete time system
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% simulation steps
ns = 50;

x = zeros(nx, ns+1);
u = zeros(nu, ns);

% initial state
x0 = zeros(nx,1);
x0(1) = 2.5;
x0(2) = 2.5;

x(:,1) = x0;
for ii=1:ns
	x(:,ii+1) = A*x(:,ii) + B*u(:,ii);
end

% plot
figure(1)
% position
subplot(3,1,1);
plot(1:ns, x(1:m,1:ns))
axis([1,ns,-3,3])
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LQR
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% input deviation penalty
R = 2.0*eye(nu);
% state deviation penalty
Q = 1.0*eye(nx);

% discrete time algebraic riccati equation
[P, L, K] = dare(A, B, Q, R)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% closed loop simulation with LQR controller
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% simulation steps
ns = 50;

x = zeros(nx, ns+1);
u = zeros(nu, ns);

% closed loop cost
clcost = 0;

x(:,1) = x0;
for ii=1:ns
	u(:,ii) = - K * x(:,ii);
	x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
	clcost = clcost + 0.5*x(:,ii)'*Q*x(:,ii) + 0.5*u(:,ii)'*R*u(:,ii);
end

% plot
figure(2)
% position
subplot(3,1,1);
plot(1:ns, x(1:m,1:ns))
axis([1,ns,-3,3])
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])

clcost

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% closed loop simulation with LQR controller
% and control saturation
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% simulation steps
ns = 50;

x = zeros(nx, ns+1);
u = zeros(nu, ns);

% closed loop cost
clcostost = 0;

% control bounds
u_min = -0.5*ones(nu,1);
u_max =  0.5*ones(nu,1);

x(:,1) = x0;
for ii=1:ns
	u(:,ii) = max(u_min, min(u_max, - K * x(:,ii)));
	x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
	clcostost = clcostost + 0.5*x(:,ii)'*Q*x(:,ii) + 0.5*u(:,ii)'*R*u(:,ii);
end

% plot
figure(3)
% position
subplot(3,1,1);
plot(1:ns, x(1:m,1:ns))
axis([1,ns,-3,3])
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])

clcost

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% closed loop simulation with MPC controller
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% control horizon
N = 12;

% cast MPC problem as QP
HH = zeros(N*(nx+nu)+nx, N*(nx+nu)+nx);
for ii=1:N
	HH((ii-1)*(nu+nx)+[1:nx],(ii-1)*(nu+nx)+[1:nx]) = Q;
	HH((ii-1)*(nu+nx)+nx+[1:nu],(ii-1)*(nu+nx)+nx+[1:nu]) = R;
end
ii = N+1;
HH((ii-1)*(nu+nx)+[1:nx],(ii-1)*(nu+nx)+[1:nx]) = Q;
%HH

gg = zeros(N*(nx+nu)+nx,1);
%gg

AA = zeros((N+1)*nx, N*(nx+nu)+nx);
AA(1:nx,1:nx) = eye(nx);
for ii=1:N
	AA(ii*nx+[1:nx],(ii-1)*(nu+nx)+[1:nx]) = -A;
	AA(ii*nx+[1:nx],(ii-1)*(nu+nx)+nx+[1:nu]) = -B;
	AA(ii*nx+[1:nx],(ii-1)*(nu+nx)+nx+nu+[1:nx]) = eye(nx);
end
%AA

bb = zeros((N+1)*nx,1);
%bb

lb = zeros(N*(nx+nu)+nx,1);
for ii=1:N
	lb((ii-1)*(nu+nx)+[1:nx]) = -inf*ones(nx,1);
	lb((ii-1)*(nu+nx)+nx+[1:nu]) = -0.5*ones(nu,1);
end
ii = N+1;
lb((ii-1)*(nu+nx)+[1:nx]) = -inf*ones(nx,1);

ub = zeros(N*(nx+nu)+nx,1);
for ii=1:N
	ub((ii-1)*(nu+nx)+[1:nx]) = inf*ones(nx,1);
	ub((ii-1)*(nu+nx)+nx+[1:nu]) = 0.5*ones(nu,1);
end
ii = N+1;
ub((ii-1)*(nu+nx)+[1:nx]) = inf*ones(nx,1);

% simulation steps
ns = 50;

x = zeros(nx, ns+1);
u = zeros(nu, ns);

% closed loop cost
clcost = 0;

% control bounds
u_min = -0.5*ones(nu,1);
u_max =  0.5*ones(nu,1);

% initial guess for the qp solver
v0 = zeros(N*(nx+nu)+nx,1);

% solution vector
vs = zeros(N*(nx+nu)+nx,1);

% make data sparse
HH = sparse(HH);
AA = sparse(AA);

x(:,1) = x0;
bb(1:nx) = x(:,1);
for ii=1:ns
	vs = quadprog(HH, gg, AA, bb, [], [], lb, ub, v0); % use quadprog in matlab
	u(:,ii) = vs(nx+(1:nu));
	x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
	clcost = clcost + 0.5*x(:,ii)'*Q*x(:,ii) + 0.5*u(:,ii)'*R*u(:,ii);
	bb(1:nx) = x(:,ii+1); % update initial state in the QP
end

% plot
figure(4)
% position
subplot(3,1,1);
plot(1:ns, x(1:m,1:ns))
axis([1,ns,-3,3])
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])

clcost