
function betaGS = CleaningStep(survivor, yt, gram, lambda, uu)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%
%%% ClearningStep:  perform the cleaning step of graphlet screening
%%%
%%% INPUT :
%%% survivor: the survivors after the screening step, a logical vector
%%% yt:      X' * y, NOT y
%%% lambda:    a tuning parameter of the cleaning step, tied to the sparsity.
%%% uu:   a tuning parameter of the cleaning step, tied to the detected
%%%       minimal signal strength
%%%
%%% OUTPUT:
%%% betaGS  The GC-step estimate for beta, a column vector of 0 and
%%% numbers larger than uu in magnitude.
%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

[p0, p] = size(gram);

if nargin < 5
    error('Too few input arguments')
end

%% options for quadprog
if  verLessThan('matlab', '7.12.0')
    options = optimset('algorithm', 'interior-point', 'LargeScale','off','display', 'off'); 
elseif  verLessThan('matlab', '7.14.0')
    options = optimset('Algorithm', 'interior-point-convex', 'LargeScale','off','display', 'off');
else
    options = optimset('Algorithm','interior-point-convex','Display','off');
end


%% survivor must be a logical vector.
survivor = logical(survivor);
nSurvivor = sum(survivor);

%% extract the survivors
ytS = yt(survivor); % much shorter!!
gramS = gram(survivor, survivor);
%% the following loop finds the clusters of gramS, and fit PMLE on each.
beta = zeros(nSurvivor, 1); % preliminary only for surviors
remain = true(nSurvivor, 1); % logical
while (sum(remain) > 0)
    i = find(remain > .5, 1);
    cluster = false(nSurvivor, 1); % col vect
    cluster(i) = true;
    newCluster = logical(abs(gramS(:, i))); % col vect for all variables correlated with i
    while(~isequal(cluster, newCluster))
        cluster = newCluster;
        newCluster = logical(sum(abs(gramS(:, cluster)), 2)); % expand cluster
    end
    
    
    if sum(cluster)>15
        sum(cluster)
        error('Cluster too long')%'MATLAB:time', 'Large cluster larger for exhaustive search. May take some time.')
    end
    %%% run Penalized MLE for this cluster
    beta(cluster) = PMLE(gramS(cluster, cluster), ytS(cluster), lambda, uu, options);
    remain = logical(remain - cluster);
end

%% expand to a p-dim estimate by adding zeros
betaGS = zeros(p,1);
betaGS(survivor) = beta;

end

%% The PMLE perform the penalized MLE in the cleaning step of the graphlet
%%screening. We do NOT assume a maximal signal strength a*tau here.
function b= PMLE(gram, y, lambda, uu, options)

% fit Penalized MLE
if nargin < 4
    error('Insufficient number of arguments');%
end
n = length(y);
b = zeros(n,1);
L = 0;

for k=0:(3^n-1)
    idx = VectorizeBase(k, 3, n); % express k on base 3 as a row vector
    idx = idx' - 1; % convert to -1, 0, 1 to indicate sign
    cluster = logical(idx ~= 0);
    card = sum(cluster);
    if card == 0
        Lt = 0;
        bt = zeros(n, 1);
    else
        signs = idx(cluster);
        yc = y(cluster);
        gramcluster = gram(cluster, cluster);
        uucluster = uu.*ones(card, 1);    
        [bt, Lq] = quadprog(full(gramcluster), -yc, -diag(signs), -uucluster,[],[],[],[],1.1*uu.*signs, options);

        %[bt, Lq] = quadprog(gramcluster, -yc, -diag(signs), -uucluster);
        Lt = Lq + lambda^2 * card /2;
        if (Lt < L)
            b = zeros(n,1);
            b(cluster) = bt;
            L = Lt;
        end
    end
end

end


%% a private function, transform a non-negative integer to a vector on a given base.
function vector = VectorizeBase(i, base, vectorLength)

if (i < 0)||(base < 2)||(vectorLength < 1)
    error('Unreasonable arguments')
end

if(i > base^vectorLength -1)
    warning('CleeningStep:vectorizeBase:IncompatibleArguments',' the vector is not long enough')
end

vector = zeros(1, vectorLength);
for j = vectorLength:-1:1
    vector(j) = i - base *floor(i/base); % remainder
    i = (i - vector(j))/base;
end

end
% end private function vectorizeBase
