%% Final Course Project - ECE 602 - Winter 2019
%% Students: Mahsa Parsapour and Zinat Saneeinia
clc;clear;close all;
%%
T = 10; s = 8; h = 1; b = 4; g = 4; 
M = 1000; % number of training samples for w 

mu = 3*rand(T,1)-1.5; 
sig = randn(T)/4;
Sigma = sig*sig';
%% (a) Affine Controller
[F, u0, P_xw, P_uw, x_tilde, u_tilde] = affineController_SupplyChain(mu,Sigma,M, s, h, b, g);
M = 500; % number of testing samples for w 
W = -exp(mu*ones(1,M) + sqrtm(Sigma)*randn(T,M));
x_af = P_xw*W + x_tilde*ones(1,M);
u_af = P_uw*W + u_tilde*ones(1,M);
cost_affine = (sum(max(h*x_af(1:end-1,:),-b*x_af(1:end-1,:))) + max(g*x_af(end,:),-b*x_af(end,:)) + sum(abs(s*u_af)));
%% (b) Naive Greedy Controller
cost_naive = zeros(1,M);
d_mean = exp(mu+.5*diag(Sigma)); % expected value of log normal distribution
for i =1:M
    x_ng = [0; d_mean + W(:,i)];  % Equation from page 7
    u_ng = d_mean - x_ng(1:end-1);% Equation from page 7
    cost_naive(i) = sum(max(h*x_ng(1:end-1),-b*x_ng(1:end-1))) + max(g*x_ng(end),-b*x_ng(end)) + sum(abs(s*u_ng));
end
%% (c) Greedy Controller
cost_greedy = zeros(1, M);
for i =1:M
    d_hat    = zeros(T, 1); 
    d_hat(1) = mu(1);
    for t=2:T 
        delta    = log(-W(1:t-1,i)) - mu(1:t-1);
        v        = mu(t) + Sigma(t,1:t-1)*(Sigma(1:t-1,1:t-1)\delta);
        Lambda   = Sigma(t,t) - Sigma(t,1:t-1)*(Sigma(1:t-1,1:t-1)\Sigma(1:t-1,t));
        d_hat(t) = exp(v+0.5*Lambda); % expected value of log normal distribution
    end    
    x_g = [0; d_hat + W(:,i)]; 
    u_g = d_hat - x_g(1:end-1);
    cost_greedy(i) = sum(max(h*x_g(1:end-1),-b*x_g(1:end-1))) + max(g*x_g(end),-b*x_g(end)) + sum(abs(s*u_g));
end
%% (d) Certainty-Equivalent MPC
cost_mpc = MPC_Controller(mu,Sigma,M, s, h, b, g,W);
%% (e) Prescient Controller
cost_pc = Prescient_Controller(mu,M, s, h, b, g,W);
%% Results
figure(1);
subplot(5,1,1); histogram(cost_naive); xlabel('Naive Greedy Controller'); hold on; line(mean(cost_naive)*ones(1,2), [0 250]);
subplot(5,1,2); histogram(cost_greedy);xlabel('Greedy Controller');       hold on; line(mean(cost_greedy)*ones(1,2), [0 250]);
subplot(5,1,3); histogram(cost_mpc);   xlabel('CE-MPC');                  hold on; line(mean(cost_mpc)*ones(1,2), [0 250]);
subplot(5,1,4); histogram(cost_affine);xlabel('Affine Controller');       hold on; line(mean(cost_affine)*ones(1,2), [0 250]);
subplot(5,1,5); histogram(cost_pc);    xlabel('Prescient Controller');    hold on; line(mean(cost_pc)*ones(1,2), [0 250]);

Controller = {'Naive Greedy';'Greedy';'CE-MPC';'Affine';'Prescient'};
Mean = [mean(cost_naive);mean(cost_greedy);mean(cost_mpc);mean(cost_affine);mean(cost_pc)];
Std = [std(cost_naive);std(cost_greedy);std(cost_mpc);std(cost_affine);std(cost_pc)];
Result = table(Controller,Mean,Std)
%% Functions
function [F, u0, P_xw, P_uw, x_tilde, u_tilde] = affineController_SupplyChain(mu,Sigma,M, s, h, b, g)
T = length(mu); 
G = tril(ones(T+1)) - eye(T+1); 
G = G(:,1:end-1); 
H = G; 
x0 = zeros(T+1,1); 

W = -exp(mu*ones(1,M) + sqrtm(Sigma)*randn(T,M));

cvx_begin quiet
    variable Q(T,T+1) 
    variable r(T) 
       
    P_xw = (eye(T+1)+H*Q)*G;
    P_uw = Q*G;
    x_tilde = (eye(T+1) + H*Q)*x0 + H*r;
    u_tilde = Q*x0 + r;
    
    x = P_xw*W + x_tilde*ones(1,M); % equation (10)
    u = P_uw*W + u_tilde*ones(1,M); % equation (10)

    minimize(sum(sum(max(h*x(1:end-1,:),-b*x(1:end-1,:))) + max(g*x(end,:),-b*x(end,:)) + sum(abs(s*u))))
    Qupper = [Q;zeros(1,T+1)];
    Qupper = triu(Qupper) - diag(diag(Qupper));
    Qupper == 0;
cvx_end

F  = (eye(T)+Q*H)\Q; % equation (17)
u0 = (eye(T)+Q*H)\r; % equation (18)
end

function cost_mpc = MPC_Controller(mu,Sigma,M, s, h, b, g,W)
T = length(mu);
cost_mpc = zeros(1,M); 
for i = 1:M
    x_mpc = zeros(T+1,1); 
    x_mpc(1) = 0;
    u_mpc = zeros(T,1);
    for t=1:T 
        v = mu(t:T) + Sigma(t:T,1:t-1)*(Sigma(1:t-1,1:t-1)\(log(-W(1:t-1,i)) - mu(1:t-1)));
        V = Sigma(t:T,t:T) - Sigma(t:T,1:t-1)*(Sigma(1:t-1,1:t-1)\Sigma(1:t-1,t:T));

        d_bar = exp(v + 0.5*diag(V));       % mean of conditional distr.
        cvx_begin quiet 
            variable u(T-t+1) 
            G = tril(ones(T-t+2)) - eye(T-t+2); 
            G = G(:,1:end-1); 
            H = G; 
            xinit = x_mpc(t)*ones(T-t+2,1); 
            x = -G*d_bar + H*u + xinit; 
            minimize sum(max(h*x(1:end-1),-b*x(1:end-1))) + max(g*x(end),-b*x(end)) + sum(abs(s*u))
        cvx_end 

        u_mpc(t) = u(1);
        x_mpc(t+1) = x_mpc(t) + u_mpc(t) + W(t,i);
    end  
        cost_mpc(i) = sum(max(h*x_mpc(1:end-1),-b*x_mpc(1:end-1))) + max(g*x_mpc(end),-b*x_mpc(end)) + sum(abs(s*u_mpc));
end
end

function cost_pc = Prescient_Controller(mu,M, s, h, b, g,W)
T = length(mu);
G = tril(ones(T+1)) - eye(T+1); 
G = G(:,1:end-1); 
H = G; 
cost_pc = zeros(1,M); 
for i = 1:M
    cvx_begin quiet
        variable u_pc(T) 
        x_pc = G*W(:,i) + H*u_pc;
        minimize sum(max(h*x_pc(1:end-1),-b*x_pc(1:end-1))) + max(g*x_pc(end),-b*x_pc(end)) + sum(abs(s*u_pc))
    cvx_end
    
    cost_pc(i) = sum(max(h*x_pc(1:end-1),-b*x_pc(1:end-1))) + max(g*x_pc(end),-b*x_pc(end)) + sum(abs(s*u_pc));
end
end