function [z_new] = z_step(x_node, u, w, rho, x_edge, u_edge, lambda)
% 
% Z_STEP(x_node, u, w, rho, x_edge, u_edge, lambda): perform the update of the z variables at node i
% There are j in N(i) z_ij variable z at node i.
%   [spaceDim+1, N(i)]
% 
% X_NODE: current X variable (hyperplane parameters) at the node evaluated
%   [spaceDim+1, 1]
% 
% U: the running sum of the residual, from the current node to each edge
%   [spaceDim+1, N(i)]
% 
% W: the weight of each arc i-j
%   [N(i), 1]
% 
% RHO: regularization parameter
% 
% X_EDGE: Xs at the nodes linked to the currently evaluated node
%   [[spaceDim+1, N(i)]
%
% U_EDGE: the running sum of the residual, from each edge to the current node

% Check the input
m = length(x_node);
spaceDim = m-1;
[p, N_i] = size(u);
[q, r] = size(w);
[s, t] = size(x_edge);

assert(spaceDim > 0);
assert(rho > 0);
assert(m == p);
assert(q == N_i);
assert(r == 1);
assert(s == m);
assert(t == N_i);
assert(N_i > 0);
% (Note, if N_i == 0, there is no need for a z update, since there is no z)
assert(lambda >= 0);

x_rep = repmat(x_node, [1, N_i]);
matrix = x_rep + u - (x_edge + u_edge);
scale = rho .* sqrt(diag(matrix'*matrix)); % we do all the norm 2 at once, and we extract the diagonal.
theta = max(1-lambda*w./scale, 0.5);

theta = repmat(theta', [m, 1]);

z_new = theta .* (x_rep + u) + (1 - theta) .* (x_edge  + u_edge);

return;


end