// Little script to illustrate the non-convergence of Repetti and Pesquet's algorithm
// when the metric is not separable.
// Olivier Fercoq - 17 February 2016


// Primal problem :
// min_{x} 2 x_n + max(0, -sum(x(1:n-1)))^2   st:    sum_{1 <= k <= n-1} x_k = 1  
//                          and  x_n - sum_{1 <= k <= n-1} x_k = 0
//                          and  x_n >= 0
//
//   Qualification conditions hold since this is a linear program
// Saddle point problem :
//  max_{y} min_{x, x_n >= 0} -y_1 + sum_{1 <= k <= n-1} x_k (y_1 - y_2) + x_n y_2 + 2 x_n + max(0, -sum(x(1:n-1)))^2
// f(x) = 2 x_n + max(0, -sum(x(1:n-1)))^2
// g(x) = I_{x'>=0}(x_n)
// M^T y = [y_1 - y_2; y_1 - y_2; .. ; y_1 - y_2); y_2]
// M x = [x_1,   x_2, ...,  x_{n-1}, 0    ;
//           -x_1, -x_2, ..., -x_{n-1}, x_n]
// -h^*(y) = -y_1
// h(u) = I_{1}(u_1) + I_{0}(u_2)
//
// Dual problem :
//  max_{y} -y_1 - (y_1 - y_2)^2 / 4  st: y_1 >= y_{2} for k=1:n-1  and  y_2 >= -2
//
// The unique solution to the dual problem is [-2;-2].
// The set of solutions to the primal problem is the set of x's such that
// sum_{1 <= k <= n-1} x_k = 1 and x_n = 1
// For instance [1; 0; ...; 0; 1] is a solution
//
// To be able to compare with Repetti and Pesquet's algorithm, 
// we implement their by interverting the primal and the dual
// With their notation, this leads to :
//  f(y) = - y_1
//  h(y) = 0
//  g(x) = I_{x'>=0}(x_n)
//  l(x) = 2 x_n


clear

d=2;
n=10;

function [val, grad] = f(x)
    tmp = max(0, - sum(x(1:n-1)));
    val = 2 * x(n) + tmp^2;
    grad = [ -tmp * ones(n-1,1); 2];
endfunction

function prox = proxg(x, i)
     if i < n
        // g_i(x_i) = 0
        prox = x(i);
    elseif i == n
        // g_n(x_n) =  I_{x'>=0}(x_n)
        prox = max(0, x(i));
    end
endfunction

function prox = proxhstar(y, W, j)
    // h^*(y) = +y_1
    gradh = [1; 0];
    prox = y(j) - W(j,j) * gradh(j);
endfunction


// M = [1 0; -1 1];
M = sparse([ones(1,n-1), 0;  - ones(1, n-1), 1]);
// K = [1 0; -1 0; 0 1];
[useless, useless2, vM] = sp2adj(M');
[iM, jM] = find(M')
K = sparse([(1:nnz(M))', iM'], vM, [nnz(M), size(M,2)]);

// Dm = [1 0; 0 2];
Dm = diag(sparse(sum(M~=0,2)))

// S = [1 0 0; 0 1./2. 1./2.];
S = [];
for k = 1:d
    S = sysdiag(S, sparse(1/full(Dm(k,k)) * ones(1, full(Dm(k,k)))));
end

Dminvexpanded = diag(diag(S' * Dm * S))
Dminv = inv(Dm)

// M = K * (Dm * S)'

ystar = [-2;-2];
xstar = [1; zeros(n-2,1); 1];

ydstar = (Dm*S)'  * ystar;

y0 = zeros(d,1);
x0 = zeros(n,1);

y = y0;
x = x0;

yd = (Dm*S)' * y0;
xd = x0;
dd = length(yd);

ydl = (Dm*S)' * y0;
xdl = x0;

normM = norm(M);
normK = norm(K);
W = 1. / normM;
U = 0.99 / (2*(n-1)/2 + normM^2 * W);
Wd = 1. / normK;  // this is \tilde \sigma in our paper, \sigma is Wd * Dminv
Ud = 0.99 / (2*(n-1)/2 + normK^2 * Wd);
Udl = 0.99 / (2 + normK^2 * Wd);

Metric = [1/W * speye(d,d), M ; M', 1/U * speye(n,n)]
Metricd = [1/Wd * speye(dd,dd), K ; K', 1/Ud * speye(n,n)]
Metricdl = [1/Wd * speye(dd,dd), K ; K', 1/Udl * speye(n,n)]

[fstar, gradstar] = f(xstar);

KMAX = 10000;
F = zeros(1, KMAX);
N = zeros(1, KMAX);
Lyap = zeros(1, KMAX);

Fd = zeros(1, KMAX);
Nd = zeros(1, KMAX);
Lyapd = zeros(1, KMAX);

Fdl = zeros(1, KMAX);
Ndl = zeros(1, KMAX);
Lyapdl = zeros(1, KMAX);

//JJ = zeros(n, KMAX)
//JJ2 = zeros(nnz(L), KMAX)

for k = 1:KMAX
    // r = rand();
    // [val, i] = min(double(r > cumsum(p)));
    i = grand(1, 1, 'uin', 1, n);

    // Repetti and Pesquet
    j = find(M(:,i));
    
    y_j = y(j);
    
    // h*(y) = -y_1
    y(j) = proxhstar(y + W * M * x, ...
                            W*speye(d,d),  j);

    // f(x) = 2 x_n + sum(x(1:n-1))^2
    [val, grad] = f(x);
    x(i) = proxg(x - U * M(j,:)' * (2 * y(j) - y_j) - U * grad, i);

    F(k) = sum(x(1:n-1))
    Lyap(k) = [y-ystar; x-xstar]' * Metric * [y-ystar; x-xstar]; 
    N(k) = norm([y-ystar]);

    // With duplication (Bianchi et al)

    j = find(K(:,i))
    
    yd_j = yd(j);
    
    // efficiency could be improved
    yd(j) = (Dm*S(:,j))' * proxhstar(S * yd + Wd * Dminv * M * xd,  ...
                                                    Wd * Dminv, 1:d);

    [val, grad_d] = f(xd);
    xd(i) = proxg(xd - Ud .* K(j,:)' * (2 * yd(j) - yd_j) - Ud .* grad_d, i);

    Fd(k) = val;
    Lyapd(k) = [yd-ydstar; xd-xstar]' * Metricd * [yd-ydstar; xd-xstar];
    Nd(k) = norm([S * yd-ystar]);

    // With duplication and long steps (Fercoq and Bianchi)
    
    j = find(K(:,i))
    
    ydl_j = ydl(j);
    
    // efficiency could be improved
    ydl(j) = (Dm*S(:,j))' * proxhstar(S * ydl + Wd * Dminv * M * xdl,  ...
                                                    Wd * Dminv, 1:d);

    [val, grad_dl] = f(xdl);
    xdl(i) = proxg(xdl - Udl .* K(j,:)' * (2 * ydl(j) - ydl_j) - Udl .* grad_dl, i);

    Fdl(k) = val;
    Lyapdl(k) = f(xd) - fstar - gradstar' * (xd-xstar) +  [ydl-ydstar; xdl-xstar]' * Metricdl * [ydl-ydstar; xdl-xstar];
    Ndl(k) = norm([S * ydl-ystar]);
end

scf(0);
clf()
plot(1:KMAX, Lyap, ':' , 1:KMAX, Lyapdl, '-', 1:KMAX, Lyapd, '--', 'thickness', 2)
handle = get("hdl")
xsetech(logflag='nl')
handle.parent.font_size=3;

//title('Checking the decrease (in expectation) of the Lyapunov function') //, we observe that the sequence is stochastic Fejér')
xlabel('iteration', 'fontsize', 3)
ylabel('Lyapunov function', 'fontsize', 3)
// We observe that the sequence is stochastic Fejér

scf(1);
clf()
plot(1:KMAX, N, ':' , 1:KMAX, Ndl, '-', 1:KMAX, Nd, '--', 'thickness', 2)
handle = get("hdl")
xsetech(logflag='nl')
handle.parent.font_size=3;
//title('Distance to the (unique) dual optimum', 'fontsize', 2) 
xlabel('iteration', 'fontsize', 3)
ylabel('||y - y*||', 'fontsize', 3)
