function [retRes,retStat] = POCRE(inY,inX,inTP,inOption)
%function [retRes,retStat] = POCRE(inY,inX,inTP,inOption)
% Fitting a linear regression model using the Penalized Orthogonal-
% Components REgression (POCRE).
%
% 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:
%   inX: X
%   inY: Y --- allow for multiple response variables
%   inTP: the tuning parameter (lambda)
%   inOption:
%       tol -- tolerance;
%       maxit -- maximum number of iterations;
%       maxcmp -- maximum number of components;
%       maxvar -- maximum number of selected variables;
%       ptype -- type of penalization:
%                'ebt' ~ EBT w/o z-transformation,
%                'ebtz' ~ EBT w/ z-transformation,
%                'l1' ~ L1,
%                'scad' ~ SCAD (gamma=3.7 by default),
%                'mcp' ~ MCP (gamma=3.7 by default);
%       gamma -- parameter for SCAD or MCP (=3.7 by default);
%       covidx -- indices of covariates of no interest from components.
%
% Output:
%   retRes: nCmp, mu, beta, varpi, vartheta, omega, theta, sigmae2,
%           bSparse(is beta sparse?), n, p;
%   retStat: some statistical/inference results based on the model with all
%            non-zero beta, e.g., pvalue, seqpv, indpv, loglik, effp
%            (effective number of predictors),sigmae2.
%
% Note: POCRE stops the construction of new components when number of
%       nonzero-coefficient variables exceeds maxvar (=n/2 by default).
%
% Dabao Zhang     July 1, 2013
%
% Revised by:
%   Dabao Zhang, 9/16/2013 -- Add inOption.covidx, retRes.pvalue, 
%                             retRes.seqpv, retRes.indpv;
%   Dabao Zhang, 9/23/2013 -- Add retStat and calculate effp.
%   Dabao Zhang, 11/6/2013 -- p-values for multiple Y (k>1) need
%                             to be implemented using mvregress.
%   Dabao Zhang, 6/16/2014 -- Output sigmae2.
%   Dabao Zhang, 8/13/2014 -- Calc of sigmae2 for multivariate Y.
%   Dabao Zhang, 7/25/2016 -- Output n,p,id to retRes.
%   Dabao Zhang, 7/20/2017 -- Construct all components for the
%                             first non-sparse tuning parameter.
%   Dabao Zhang, 12/15/2017 -- default maxvar.
%
% Copyright (c) 2013-2017 by Dabao Zhang.

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

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 = 10; end
if( ~isfield(inOption,'maxvar') ) inOption.maxvar = min(size(inX))/2; end
if( ~isfield(inOption,'ptype') ) inOption.ptype = 'ebtz'; end
if( ~isfield(inOption,'covidx') ) inOption.covidx = []; end

[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

% 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;
retRes.n = n;
retRes.p = p;
retRes.id = 'pocre';

% 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 = [];
omega = ones(p,1);      omeganrm = 1;
fprintf('\b  ');
%while( (omeganrm>eps)&&(idxCmp<inOption.maxcmp)&&(retRes.bSparse) )
while( (omeganrm>eps)&&(idxCmp<inOption.maxcmp) )
    idxCmp = idxCmp+1;

    tmpRes = POCRE1POC(Y,X,retRes,inTP,inOption);
    idxNZBeta = union(idxNZBeta,find(abs(tmpRes.varpi)>eps));
    if( length(idxNZBeta)>inOption.maxvar )
        retRes.bSparse = 0;
    end
    if( ~isempty(tmpRes.varpi) )
        % 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;

        X = tmpRes.X;
        fprintf('.');
    else
        omeganrm = 0;
    end
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);

% Calculate p-values and log-likelihood value if requested
if( nargout>1 )
    if( k==1 )
        retStat = EvalComponentsSig(inY,inX,retRes,inOption);
    else
         % See mvregress() to implement likelihood ratio tests
         warning('  Calculation of p-values: to be implemented!');
    end
    
    % Calculate the log-likelihood function value
    retStat.loglik = 0;
    retStat.effp = zeros(1,k);
    retStat.sigmae2 = ones(1,k);
    for j=1:k
        %effp = effp + sum(abs(retRes.beta(:,j))>eps));
        tmpX = [ones(n,1) inX(:,abs(retRes.beta(:,j))>eps)];
        [tmpR,tmpJ] = rref(tmpX);
        retStat.effp(j) = length(tmpJ);
        tmpX = tmpX(:,tmpJ);
        yEst = tmpX*regress(inY(:,j),tmpX);
        retStat.loglik = retStat.loglik-(log(2*pi)...
                         +1+log(sum((inY(:,j)-yEst).^2,1)/n))*n/2;
        retStat.sigmae2(j) = sum((inY(:,j)-yEst).^2,1)/(n-retStat.effp(j));
    end
end



%--------------------------------------------------------------------------
%
function retRes = POCRE1POC(inY,inX,inRes,inTP,inOption)
%function retRes = POCRE1POC(inY,inX,inRes,inTP,inOption)
% 
% Construct one orthogonal component in POCRE.
%
% Model: Y = mu + X*beta + Epsilon, Epsilon ~ (0, sigmae);
%
% Input:
%   inX: X
%   inY: Y, which can be multivariate
%   inRes: results of previous construction w/ (omega,theta,varpi,vartheta)
%   inTP: tuning parameter (lambda)
%   inOption:
%       tol -- tolerance;
%       maxit -- maximum number of iterations;
%       maxcmp -- maximum number of components;
%       maxvar -- maximum number of selected variables;
%       ptype -- type of penalization:
%                'ebt' ~ EBT w/o z-transformation,
%                'ebtz' ~ EBT w/ z-transformation,
%                'l1' ~ L1,
%                'scad' ~ SCAD (gamma=3.7 by default),
%                'mcp' ~ MCP (gamma=3.7 by default);
%       gamma -- parameter for SCAD or MCP (=3.7 by default);
%
% Output:
%   retRes: X_j (updated X), omega, varpi, vartheta, theta
%
% Dabao Zhang     May 30, 2013
%
% Revised by:
%       Dabao Zhang, July 22, 2016 -- Inherited from Construct1POC(...).
%
% Copyright (c) 2013 by Dabao Zhang.

error(nargchk(5, 5, 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);

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);
    
    switch( inOption.ptype )
        case 'ebt'  % POCRE0 via Empirical Bayes Thresholding
            tmpSD = median(abs(omega))*inTP/norminv(0.75,0,1);
            tmpRes = GEBCauchy(omega,tmpSD,[1,1,1]);
            omega = tmpRes.mu;
            
        case 'ebtz'  % POCRE via Empirical Bayes Thresholding
            omega = omega/norm(omega);
            omega = .5*log((1+omega)./(1-omega));
            tmpSD = inTP/sqrt(length(omega)-3);
            tmpRes = GEBCauchy(omega,tmpSD,[1,1,1]);
            omega = 1-2./(exp(2*tmpRes.mu)+1);
            
        case 'l1'
            tmpSD = median(abs(omega))*inTP/norminv(0.75,0,1);
            omega = max(abs(omega)-tmpSD*sqrt(2*log(p)),0)...
                    .*sign(omega);
            
        case 'scad'
            if( ~isfield(inOption,'gamma') ) inOption.gamma = 3.7; end
            tmpSD = median(abs(omega))*inTP/norminv(0.75,0,1);
            tmpT = tmpSD*sqrt(2*log(length(omega))); % Threshold of L1

            % |z|<=2*lambda
            tIdx = find(abs(omega)<=2*tmpT);
            omega(tIdx) = max(abs(omega(tIdx))-tmpT,0).*sign(omega(tIdx));

            % 2*lambda<|z|<=gamma*lambda
            tIdx = setdiff(find(abs(omega)<=inOption.gamma*tmpT),tIdx);
            tmpSO = max(abs(omega(tIdx))-tmpT/(1-1/inOption.gamma),0);
            omega(tIdx) = tmpSO.*sign(omega(tIdx))/(1-1/(inOption.gamma-1));
            
        case 'mcp'
            if( ~isfield(inOption,'gamma') ) inOption.gamma = 3.7; end
            tmpSD = median(abs(omega))*inTP/norminv(0.75,0,1);
            tmpT = tmpSD*sqrt(2*log(length(omega))); % Threshold of L1
            tIdx = find(abs(omega)<=inOption.gamma*tmpT);
            tmpSO = max(abs(omega(tIdx))-tmpT,0);
            omega(tIdx) = tmpSO.*sign(omega(tIdx))/(1-1/inOption.gamma);
            
        otherwise
            error('POCRE: The penalization method is not supported!');
    end

    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

if( omeganrm>eps )
    % 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';
end



%--------------------------------------------------------------------------
%
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



%--------------------------------------------------------------------------
%
function retStat = EvalComponentsSig(inY,inX,inRes,inOption)
%function retStat = EvalComponentsSig(inY,inX,inRes,inOption)
%
% Evaluate the significance of constructed components.
%
% Note: This function is implemented only for univariate Y.
%
% Input:
%   inX: X
%   inY: Y --- allow for multiple response variables
%   inRes: first output of POCRE;
%   inOption:
%       covidx -- indices of covariates of no interest from components.
%
% Dabao Zhang     November 6, 2013
%
% Revised by:
%       Dabao Zhang, September 23, 2016 --  Replace regstats() by fitlm().
%
% Copyright (c) 2013 by Dabao Zhang.

[n,p] = size(inX);
X = bsxfun(@minus,inX,inRes.xShift);

% Covariates of no interest
covx = [];
for j = 1:inRes.nCmp
    if( sum(abs(inRes.varpi(inOption.covidx,j)))>eps )
        covx = [covx ...
                X(:,inOption.covidx)*inRes.varpi(inOption.covidx,j)];
    end
end
ncovx = size(covx,2);

tidx = setdiff(1:p,inOption.covidx);
X = [covx X(:,tidx)*inRes.varpi(tidx,:)];
ts = fitlm(X,inY);
retStat.pvalue = ts.Coefficients.pValue((ncovx+2):end);

retStat.seqpv = retStat.pvalue;
for j=1:(inRes.nCmp-1)
    ts = fitlm(X(:,1:(j+ncovx)),inY);
    retStat.seqpv(j) = ts.Coefficients.pValue(j+ncovx+1);
end

retStat.indpv = retStat.seqpv;
for j=2:inRes.nCmp
    ts = fitlm(X(:,[1:ncovx j+ncovx]),inY);
    retStat.indpv(j) = ts.Coefficients.pValue(ncovx+2);
end
