%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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

%Ac
%Bc

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

% load control package in octave
pkg load control 

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

% 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

% octave using gnuplot
graphics_toolkit('gnuplot');

% plot
figure(1)
% position
subplot(3,1,1);
plot(1:ns, x(1:m,1:ns))
title('free response')
axis([1,ns,-3,3])
ylabel('position')
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
ylabel('velocity')
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])
ylabel('force')
xlabel('simulation step')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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
closedloopcpost = 0;

x(:,1) = x0;
for ii=1:ns
	u(:,ii) = - K * x(:,ii);
	x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
	closedloopcpost = closedloopcpost + 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))
title('LQR')
axis([1,ns,-3,3])
ylabel('position')
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
ylabel('velocity')
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])
ylabel('force')
xlabel('simulation step')

closedloopcpost

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 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
closedloopcpost = 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);
	closedloopcpost = closedloopcpost + 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))
title('LQR with saturation')
axis([1,ns,-3,3])
ylabel('position')
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
ylabel('velocity')
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])
ylabel('force')
xlabel('simulation step')

closedloopcpost

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

if 0

% 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
closedloopcpost = 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);

x(:,1) = x0;
bb(1:nx) = x(:,1);
for ii=1:ns
	vs = qp(v0, HH, gg, AA, bb, lb, ub); % use quadprog in matlab
	u(:,ii) = vs(nx+(1:nu));
	x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
	closedloopcpost = closedloopcpost + 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))
title('MPC')
axis([1,ns,-3,3])
ylabel('position')
% velocity
subplot(3,1,2);
plot(1:ns, x(m+1:end,1:ns))
axis([1,ns,-3,3])
ylabel('velocity')
% inputs
subplot(3,1,3);
plot(1:ns, u)
axis([1,ns,-1.5,1.5])
ylabel('force')
xlabel('simulation step')

closedloopcpost

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% timings using qp
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if 1

Nmax_qp = 10;
times_qp = zeros(Nmax_qp,1);

for nn=1:Nmax_qp

	N = nn;

	% 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 = 20;

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

	% closed loop cost
	closedloopcpost = 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);

	tic; % start simulation

	x(:,1) = x0;
	bb(1:nx) = x(:,1);
	for ii=1:ns
		vs = qp(v0, HH, gg, AA, bb, lb, ub); % use quadprog in matlab
		u(:,ii) = vs(nx+(1:nu));
		x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
		closedloopcpost = closedloopcpost + 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

	times_qp(nn) = toc/ns; % stop simulation

end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% timings using hpipm
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% compule mex interface
mex HPIPM_d_solve_ipm2_hard_ocp_qp.c /opt/hpipm/lib/libhpipm.a /opt/blasfeo/lib/libblasfeo.a -I/opt/hpipm/include -I/opt/blasfeo/include

Nmax_hpipm = 100;
times_hpipm = zeros(Nmax_hpipm,1);

for nn=1:Nmax_hpipm

	N = nn;

	% number of bounds
	nb = 3;

	% number of general constraints
	ng = 0;
	ngN = 0;

	% extra (zero) matrices for MPC formulation
	b = zeros(nx,1);
	q = zeros(nx,1);
	r = zeros(nx,1);
	S = zeros(nx,nu);
	C  = zeros(ng, nx);
	D  = zeros(ng, nu);
	QN = Q;
	qN = q;
	CN = zeros(ng, nx);
	lg = zeros(ng,1);
	ug = zeros(ng,1);
	lgN = zeros(ngN,1);
	ugN = zeros(ngN,1);

	% repeat matrices
	hA = repmat(A, 1, N);
	hB = repmat(B, 1, N);
	hb = repmat(b, 1, N);
	hQ = repmat(Q, 1, N);
	hS = repmat(S, 1, N);
	hR = repmat(R, 1, N);
	hq = repmat(q, 1, N);
	hr = repmat(r, 1, N);
	hC = [zeros(ng,nx), repmat(C, 1, N-1)];
	hD = [repmat(D, 1, N)];
	hlb = repmat(u_min, 1, N+1);
	hub = repmat(u_max, 1, N+1);
	hlg = repmat(lg, 1, N);
	hug = repmat(ug, 1, N);

	% memory for solution
	hx = zeros(nx, N+1); 
	hu = zeros(nu, N);
	hpi = zeros(nx,N);
	hlam = zeros(2*(nb+ng)*N+2*(nb+ngN),1);

	% solver arguments
	free_x0 = 0; % consider x0 as optimization variable
	warm_start = 0; % read initial guess from x and u
	time_invariant = 0; % read a different matrix for each stage

	mu0 = 2;        % max element in cost function as estimate of max multiplier
	kk = -1;		% actual number of performed iterations
	k_max = 20;		% maximim number of iterations
	tol = 1e-8;		% tolerance in the duality measure
	infos = zeros(5, k_max);
	inf_norm_res = zeros(1, 4);


	% simulation steps
	ns = 20;

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

	% closed loop cost
	closedloopcpost = 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);

	tic; % start simulation

	x(:,1) = x0; % initial state
	hx(:,1) = x(:,1);
	for ii=1:ns
		HPIPM_d_solve_ipm2_hard_ocp_qp(kk, k_max, mu0, tol, N, nx, nu, nb, ng, ngN, time_invariant, free_x0, warm_start, hA, hB, hb, hQ, QN, hR, hS, hq, qN, hr, hlb, hub, hC, hD, hlg, hug, CN, lgN, ugN, hx, hu, hpi, hlam, inf_norm_res, infos);
		u(:,ii) = hu(:,1);
		x(:,ii+1) = A * x(:,ii) + B * u(:,ii);
		closedloopcpost = closedloopcpost + 0.5*x(:,ii)'*Q*x(:,ii) + 0.5*u(:,ii)'*R*u(:,ii);
		hx(:,1) = x(:,ii+1);
	end

	times_hpipm(nn) = toc/ns; % stop simulation

end

figure(5)
loglog(1:Nmax_qp, times_qp, 'b');
hold on
loglog(1:Nmax_hpipm, times_hpipm, 'r');
hold off
xlabel('N')
ylabel('time [s]')
