clear all;
% name random variables
y1 = RV({'y1'},1);
y2 = RV({'y2'},1);
y3 = RV({'y3'},1);
u1 = RV({'u1'},1);
u2 = RV({'u2'},1);

% create f(y_t| y_{t-3}, u_{t-1})
fy.class = 'mlnorm<ldmat>';
fy.rv    = RVjoin([y1,y2,y3]);
fy.rvc   = RVtimes([y1,y2,y3,u1,u1,u2,u2], [-1, -1, -1, 0, -1, 0, -1]);
fy.A     = [0.8 , 0.2 , 0 , -0.3 , 0.4 , 0 , 0;...
           -0.2 , 0.5 , -0.8 , 0.2 , 0.5 , -0.2 , -0.5;...
            0 , 1.1 , -0.5 , 0 , 0 , -0.2 , 0.3];
fy.const = [0;0;0];
fy.R     = 0.1*eye(3);

DS.class = 'PdfDS';
DS.pdf = fy;

% create ARX estimator
A1.class = 'ARX';
A1.rv = RVjoin([y1,y2]);
A1.rgr = RVtimes([y1,y2,u1,u1],[-1, -1, 0, -1]) ; % correct structure is {y,y}
A1.log_level ='logbounds,logevidence';
A1.frg = 0.99;

A2=A1;
A2.rv = RVjoin([y2,y3]);
A2.rgr = RVtimes([y2,y3,u2,u2],[-1, -1, 0, -1]) ; % correct structure is {y,y}

C1.class = 'LQG_ARX';
C1.ARX = A1;
C1.Qu = 0.1;
C1.Qy = 0.1*eye(2);
C1.yreq = [0;1]; %y2=1
C1.horizon = 1;

C2=C1;
C2.ARX = A2;
C2.yreq = [1;0]; %y2=1

P1.class = 'ARXAgent';
P1.name = 'P1';
P1.lqg_arx = C1;
P1.lqg_arx.class = 'LQG_ARX';
P1.merger.class = 'merger_mix';
P1.merger.method = 'geometric';
%P1.merger.dbg_file = 'mp.it';
P1.merger.ncoms = 1;
P1.merger.stop_niter= 5;
P1.neighbours = {};%{'P2'};

P2=P1;
P2.name = 'P2';
P2.lqg_arx = C2;
P2.neighbours = {};

exper.Ndat = 10;
exper.burnin = 3;
exper.burn_pdf.class = 'enorm<ldmat>';
exper.burn_pdf.mu = [0;0];
exper.burn_pdf.R  = 0.01*eye(2);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% MONTE CARLO %%%%%%%%%%%%%%%%%%%

Ntrials = 100;
loss_non_coop = zeros(1,Ntrials);
for i=1:Ntrials
    M= arena(DS,{P1,P2},exper);

    Y = [M.DS_y1 M.DS_y2 M.DS_y3];
    Yreq = ones(size(M.DS_y1))*[0 1 0];
    loss_non_coop(i) = trace((Y-Yreq)'*0.01*(Y-Yreq)) + M.DS_u1'*C1.Qu*M.DS_u1 + M.DS_u2'*C1.Qu*M.DS_u2;
    if loss_non_coop(i)>100
        %keyboard
    end
end
mean(loss_non_coop)

loss_coop = zeros(1,Ntrials);
for i=1:Ntrials
    P1.neighbours = {'P2'};
    P2.neighbours = {'P1'};
    M= arena(DS,{P1,P2},exper);

    Y = [M.DS_y1 M.DS_y2 M.DS_y3];
    Yreq = ones(size(M.DS_y1))*[0 1 0];
    loss_coop(i) = trace((Y-Yreq)'*0.01*(Y-Yreq)) + M.DS_u1'*C1.Qu*M.DS_u1 + M.DS_u2'*C1.Qu*M.DS_u2;
    if loss_coop(i)>100
        %keyboard
    end
end
mean(loss_coop)