clear all
close all
clc

%% Dynamical model
% Trucks (except leader)
N = 5;

% Distance
D = 1;

% Damping
damping = 0.0; 

% Continuous time model
% state [d;s]
Ac = [zeros(N), diag(ones(N-1,1), -1) - eye(N) - damping*eye(N); zeros(N), zeros(N)];
Bc = [zeros(N); eye(N)];

% Sampling time
Ts = 1;

% Discrete time model
A = eye(2*N) + Ts*Ac;
B = Ts*Bc;

%% Parameters
% Cost function
% J = x(K)' Qf x(K) + sum_k=0^(K-1) x(k)' Q x(k) + u(k)' R u(k)

Kc = 20;
K  = Kc/Ts;

% State cost
costState      = 1; 
costFinalState = 1; 
costInput      = 10; 

% Build matrices
QfN = zeros(2*N,2*N,N);
QN  = zeros(2*N,2*N,N);
RN  = zeros(N,N,N);
for pl = 1:N
    if pl <= 4
        plstates = [pl,pl+1];
    else
        plstates = pl;
    end
    QN(plstates,plstates,pl) = costState*eye(length(plstates));
    QN(plstates,plstates,pl) = costFinalState*eye(length(plstates));
    RN(pl,pl,pl)             = costInput;
end
Q  = sum(QN,3);
Qf = sum(QfN,3);
R  = sum(RN,3);

%% Parameters for the simulation/plots
T            = K;
initialSpeed = 0.5; 

set(groot, 'defaultAxesTickLabelInterpreter','latex');
set(groot, 'defaultLegendInterpreter','latex');
set(groot, 'defaultTextInterpreter','latex');

titleU = ['Uncontrolled, $s_0 = ' num2str(initialSpeed) '$'];
title1 = ['$1$ player, $s_0 = ' num2str(initialSpeed) '$'];
titleN = ['$N$ players, $s_0 = ' num2str(initialSpeed) '$'];

%% Simulation uncontrolled
xU      = nan(2*N,T+1);
xU(:,1) = [zeros(N,1); initialSpeed*ones(N,1)];
uU      = zeros(N,T);
costU   = zeros(N,1); 

for t = 1:T
    xU(:,t+1) = A * xU(:,t) + B * uU(:,t);
    for i = 1:N
        costU(i) = costU(i) + xU(:,t)'*QN(:,:,i)*xU(:,t) + uU(:,t)'*RN(:,:,i)*uU(:,t);
    end
end
for i = 1:N
    costU(i) = costU(i) + xU(:,T+1)'*QfN(:,:,i)*xU(:,T+1);
end

%% Plot trajectories
figure(1)
subplot(2,1,1)
hold on
plot(1:T+1,xU(1:N,:),'.-')
plot(1:T+1,+D*ones(1,T+1),'--r')
plot(1:T+1,-D*ones(1,T+1),'--r')
grid on
xlim([1 T+1])
ylabel('Distances')
title(titleU)
set(gca,'Box','on','Fontsize',14)

subplot(2,1,2)
plot(1:T,uU,'.-')
grid on
xlim([1 T+1])
xlabel('Time')
ylabel('Acceleration')
set(gca,'Box','on','Fontsize',14)

y = (1:N)'*D*ones(1,T+1) + cumsum(xU(1:N,:),1);

figure(2)
hold on
plot(1:T+1,y,'.-')
plot(1:T+1,zeros(1,T+1),'--r')
grid on
xlim([1 T+1])
ylabel('Positions')
xlabel('Time')
title(titleU)
set(gca,'Box','on','Fontsize',14)

%% One-player LQR Optimal controller (finite time)
P     = nan(2*N,2*N,K+1);
Gamma = nan(N,2*N,K);

P(:,:,K+1) = Qf;

for k = K:-1:1
    P(:,:,k) = Q + A' * P(:,:,k+1) * A - ...
        A' * P(:,:,k+1) * B * inv(R + B' * P(:,:,k+1) * B) * B' * P(:,:,k+1) * A;
    Gamma(:,:,k) = - inv(R + B' * P(:,:,k+1) * B) * B' * P(:,:,k+1) * A;
end

disp('Single player')
disp(Gamma(:,:,1))

%% Simulation one-player LQR
x      = nan(2*N,T+1);
x(:,1) = [zeros(N,1); initialSpeed*ones(N,1)];
u      = nan(N,T);
cost   = zeros(N,1); 

for t = 1:T
    if t <= K
        u(:,t) = Gamma(:,:,t) * x(:,t);
    end
    x(:,t+1) = A * x(:,t) + B * u(:,t);
    for i = 1:N
        cost(i) = cost(i) + x(:,t)'*QN(:,:,i)*x(:,t) + u(:,t)'*RN(:,:,i)*u(:,t);
    end
end
for i = 1:N
    cost(i) = cost(i) + x(:,T+1)'*QfN(:,:,i)*x(:,T+1);
end

%% Plot trajectories
figure(3)
subplot(2,1,1)
hold on
plot(1:T+1,x(1:N,:),'.-')
plot(1:T+1,+D*ones(1,T+1),'--r')
plot(1:T+1,-D*ones(1,T+1),'--r')
grid on
xlim([1 T+1])
ylabel('Distances')
title(title1)
set(gca,'Box','on','Fontsize',14)

subplot(2,1,2)
plot(1:T,u,'.-')
grid on
xlim([1 T+1])
xlabel('Time')
ylabel('Acceleration')
set(gca,'Box','on','Fontsize',14)

y = (1:N)'*D*ones(1,T+1) + cumsum(x(1:N,:),1);

figure(4)
hold on
plot(1:T+1,y,'.-')
plot(1:T+1,zeros(1,T+1),'--r')
grid on
xlim([1 T+1])
ylabel('Positions')
xlabel('Time')
title(title1)
set(gca,'Box','on','Fontsize',14)

%% N-player LQR (finite time)
PN = nan(2*N,2*N,K+1,N);
for pl = 1:N
    PN(:,:,K+1,pl) = QfN(:,:,pl);
end

GammaN = nan(N,2*N,K);

Phi = nan(2,2*N);
F = zeros(N);
G = zeros(N,2*N);
S = zeros(2*N);

for k = K:-1:1
    for pl = 1:N
        Phi = (RN(pl,pl,pl) + B(:,pl)' * PN(:,:,k+1,pl) * B(:,pl)) \ B(:,pl)' * PN(:,:,k+1,pl);
        for pl2 = 1:N
            F(pl,pl2) = Phi * B(:,pl2);
            G(pl,:)   = -Phi * A;
        end
        F(pl,pl) = 1;
    end
    
    GammaN(:,:,k) = F\G;
    
    for pl = 1:N
        S = A;
        for pl2 = 1:N
            S = S + B(:,pl2) * GammaN(pl2,:,k);
        end
        PN(:,:,k,pl) = QN(:,:,pl) + GammaN(pl,:,k)' * RN(pl,pl,pl) * GammaN(pl,:,k) + ...
            S' * PN(:,:,k+1,pl) * S;
    end
end

disp('N players')
disp(GammaN(:,:,1))

%% Simulation N-player LQR
xN      = nan(2*N,T+1);
xN(:,1) = [zeros(N,1); initialSpeed*ones(N,1)];
uN      = nan(N,T);
costN   = zeros(N,1); 

for t = 1:T
    if t <= K
        uN(:,t) = GammaN(:,:,t) * xN(:,t);
    end
    xN(:,t+1) = A * xN(:,t) + B * uN(:,t);
    for i = 1:N
        costN(i) = costN(i) + xN(:,t)'*QN(:,:,i)*xN(:,t) + uN(:,t)'*RN(:,:,i)*uN(:,t);
    end
end
for i = 1:N
    costN(i) = costN(i) + xN(:,T+1)'*QfN(:,:,i)*xN(:,T+1);
end

%% Plot
figure(5)
subplot(2,1,1)
hold on
plot(1:T+1,xN(1:N,:),'.-')
plot(1:T+1,+D*ones(1,T+1),'--r')
plot(1:T+1,-D*ones(1,T+1),'--r')
grid on
xlim([1 T+1])
ylabel('Distances')
title(titleN)
set(gca,'Box','on','Fontsize',14)

subplot(2,1,2)
plot(1:T,uN,'.-')
grid on
xlim([1 T+1])
xlabel('Time')
ylabel('Acceleration')
set(gca,'Box','on','Fontsize',14)

yN = (1:N)'*D*ones(1,T+1) + cumsum(xN(1:N,:),1);

figure(6)
hold on
plot(1:T+1,yN,'.-')
plot(1:T+1,zeros(1,T+1),'--r')
grid on
xlim([1 T+1])
xlabel('Time')
ylabel('Positions')
title(titleN)
set(gca,'Box','on','Fontsize',14)

%% Final print out
fprintf('Social cost (1 player)  : %f\n',sum(cost));
fprintf('Social cost (N players) : %f\n',sum(costN));

disp('Individual costs (1 player and N players)')
disp([cost costN])
