function [retX,retSIdx,retRes,retStat] = POCREScreen(inY,inX,inEIdx,inOption)
%function [retX,retXIdx,retRes,retStat] = POCREScreen(inY,inX,inEIdx,inOption)
%
% Using the Penalized Orthogonal-Components REgression (POCRE) to screen
% the variables for fitting a linear regression model.
%
% Model: Y = mu + X*beta + Epsilon, Epsilon ~ (0, sigmae2),
%        X*beta = \sum_j (X*varpi_j)*vartheta_j'
%               = \sum_j (X_j*omega_j)*vartheta_j';
%
% Input:
%   inY: Y --- allow for multiple response variables
%   inX: X
%   inEIdx: Indices of exceptional variables which are always included in
%           the model (but not counted into maxvar);
%   inOption:
%       maxit -- maximum number of iterations;
%       maxcmp -- maximum number of components;
%       maxvar -- maximum number of selected variables;
%       covidx -- indices of covariates of no interest from components.
%       tol -- tolerance;
%
% Output:
%   retX: inX[,[retEIdx,retSIdx]];
%   retSIdx: indices of selected variables;
%   retRes: nCmp, mu, beta, varpi, vartheta, omega, theta, sigmae2;
%   retStat: updated mu & beta after constructing each component.
%
% Note: POCRE will select [maxvar/maxcmp] variables for each component, and
%       stop the construction of new components when total number of
%       selected variables reaches maxvar (=n by default).
%
% Dabao Zhang     July 23, 2016
%
% Revised by:
%       Dabao Zhang, November 7, 2016 -- Default value of inOption.maxvar.
%       Dabao Zhang, July 31, 2017 -- Output retStat.
%
% Copyright (c) 2013-2017 by Dabao Zhang.

error(nargchk(2, 4, nargin));

[n,p] = size(inX);  % X with n objects p variables;
k = size(inY,2);    % Y is k-variate;
if( n~=size(inY,1) )
    error('inX and inY should have the same number of rows!');
end

if( nargin<4 ) inOption = []; end
if( ~isfield(inOption,'tol') ) inOption.tol = 1e-6; end    % 1e-12?
if( ~isfield(inOption,'maxit') ) inOption.maxit = 100; end
if( ~isfield(inOption,'maxcmp') ) inOption.maxcmp = 5; end
if( ~isfield(inOption,'maxvar') ) inOption.maxvar = min(n/log(n),p/2); end
inOption.maxvar = round(inOption.maxvar);
if( ~isfield(inOption,'ptype') ) inOption.ptype = 'ebtz'; end
if( ~isfield(inOption,'covidx') ) inOption.covidx = []; end
if( nargin<3 ) inEIdx = []; end

% Initialization
retRes.nCmp = 0;
retRes.mu = zeros(1,k);
retRes.beta = zeros(p,k);
retRes.varpi = zeros(p,0);
retRes.vartheta = zeros(k,0);
retRes.omega = zeros(p,0);
retRes.theta = zeros(p,0);
%retRes.bSparse = 1;

if( nargout==4 )
    retStat.mu = zeros(1,k);
    retStat.beta = zeros(p,k);
end

% Define indicator vectors for inEIdx (idcE) & retSIdx (idcS)
idcE = zeros(p,1);
idcE(inEIdx) = 1;
idcS = zeros(p,1);

% Centralize both X and Y
retRes.xShift = mean(inX,1);
X = bsxfun(@minus,inX,retRes.xShift);
retRes.yShift = mean(inY);
Y = bsxfun(@minus,inY,retRes.yShift);

% Squentially construct each components
%idxCmp = 0;             idxNZBeta = [];
idxCmp = 0;      omega = ones(p,1);      omeganrm = 1;
nSele = round(inOption.maxvar/inOption.maxcmp);
fprintf(' Screening variables ');
while( (sum(idcS)<inOption.maxvar)&&(idxCmp<inOption.maxcmp) )
    idxCmp = idxCmp+1;

    if( idxCmp==inOption.maxcmp )
        nSele = inOption.maxvar-nSele*(inOption.maxcmp-1);
    end
    tmpRes = POCREScreen1POC(Y,X,retRes,idcE+idcS,nSele,inOption);
    %idxNZBeta = union(idxNZBeta,find(abs(tmpRes.varpi)>eps));
    idcS = idcS+tmpRes.idcS;

    % Update omega, theta, varpi, vartheta
    retRes.omega = [retRes.omega tmpRes.omega];
    retRes.theta = [retRes.theta tmpRes.theta];
    retRes.varpi = [retRes.varpi tmpRes.varpi];
    retRes.vartheta = [retRes.vartheta tmpRes.vartheta];
    retRes.nCmp = idxCmp;

    % Save mu & beta
    if( nargout==4 )
        retStat(idxCmp).beta = retRes.varpi*retRes.vartheta';
        retStat(idxCmp).mu = retRes.yShift-retRes.xShift*retStat(idxCmp).beta;
    end
    
    X = tmpRes.X;
    fprintf('.');
end
fprintf('\n');

% Calculate beta
retRes.beta = retRes.varpi*retRes.vartheta';

% Calculate mu
retRes.mu = retRes.yShift-retRes.xShift*retRes.beta;

% Calculate sigmae2
retRes.sigmae2 = sum(bsxfun(@minus,inY-inX*retRes.beta,retRes.mu).^2) ...
                 /(n-retRes.nCmp);

retSIdx = find(idcS);
retX = inX(:,[inEIdx,retSIdx]);



%--------------------------------------------------------------------------
%
function retRes = POCREScreen1POC(inY,inX,inRes,inIdc,inNX,inOption)
%function retRes = POCREScreen1POC(inY,inX,inRes,inIdc,inNX,inOption)
% 
% Construct one orthogonal component for POCREScreen(...)
%
% Model: Y = mu + X*beta + Epsilon, Epsilon ~ (0, sigmae);
%
% Input:
%   inY: Y, which can be multivariate
%   inX: X
%   inRes: results of previous construction w/ (omega,theta,varpi,vartheta)
%   inIdc: a vector of 0 & 1 indicating whether the variables have been
%          selected in prevously constructed components;
%   inNX: number of variables to be selected for current component;
%   inOption:
%       tol -- tolerance;
%       maxit -- maximum number of iterations;
%       maxcmp -- maximum number of components;
%       maxvar -- maximum number of selected variables;
%
% Output:
%   retRes: X_j (updated X), omega, varpi, vartheta, theta, idcS.
%
% Dabao Zhang     July 25, 2016
%
% Revised by:
%
% Copyright (c) 2013-2016 by Dabao Zhang.

error(nargchk(6, 6, nargin));

% Initialize retRes
p = size(inX,2);
retRes.omega = zeros(p,0);
retRes.theta = zeros(p,0);
retRes.varpi = zeros(p,0);
retRes.vartheta = zeros(size(inY,2),0);
retRes.idcS = zeros(p,1);

idxCmp = length(inRes.varpi)+1; % Index of current component
tmpMat = inY'*inX;

% Calculate the leading eigenvector of X_j'YY'X_j
[tmpU,tmpS,omega] = svds(tmpMat,1);

% Calculate the regularized component vector
idIter = 0;
omegapre = omega;   omegadf = 1;  omeganrm = norm(omega);
alpha = tmpMat'*(tmpMat*omega); alpha = alpha/norm(alpha);
while((omeganrm>eps)&&(omegadf>inOption.tol)&&(idIter<=inOption.maxit))
    idIter = idIter + 1;
    omegapre = omega;
    omega = tmpMat'*(tmpMat*alpha);
    
    % Q: Hard Threshold or Soft Threshold?
    tmpCutoff = kthlargestvalue(abs(omega(~inIdc)),inNX+1);
    omega = max(abs(omega)-tmpCutoff,0).*sign(omega);
    %omega = omega.*(abs(omega)>tmpCutoff);

    alpha = tmpMat'*(tmpMat*omega);
    omeganrm = norm(omega);
    if( omeganrm>eps )
        alpha = alpha/norm(alpha);  % norm(omega)>eps ==> norm(alpha)>0?
        omegadf = norm(omegapre-omega)/omeganrm;
    end
end

% Update omega, theta, varpi, vartheta
retRes.omega = omega/omeganrm;
if( idxCmp==1 )
    retRes.varpi = retRes.omega;
else
    retRes.varpi = varpifromomega(retRes.omega,inRes.omega,inRes.theta);
end
eta = inX*retRes.omega;
retRes.theta = inX'*eta/(eta'*eta);
retRes.vartheta = inY'*eta/(eta'*eta);

% Update X
retRes.X = inX-eta*retRes.theta';

% Update retRes.idcS
retRes.idcS = max((abs(omega)>eps)-inIdc,0);



%--------------------------------------------------------------------------
%
function retVarpi = varpifromomega(omega,omegapre,thetapre)
%function retVarpi = varpifromomega(omega,omegapre,thetapre)
% Update varpi from omega at a given iteration.
%
% Input:
%   omega: current value of omega;
%   omegapre: omega of previous components;
%   thetapre: theta of previous components.
%
% Output:
%   retVarpi: Updated varpi.
%
% Dabao Zhang     October 4, 2012
%
% Revised by:
%
% Copyright (c) 2012 by Dabao Zhang.

nCmp = size(omegapre,2);
retVarpi = omega;
for( j=1:nCmp )
    retVarpi = retVarpi - omegapre(:,j)*(thetapre(:,j)'*retVarpi);
end
