function fit = uBaL(X,y,family,groups,interact,intercept,r,burnIn,nSample)
%--------------------------------------------------------------------------
% uBaL.m: implement the unified Bayesian adaptive Lasso
%--------------------------------------------------------------------------
%
% DESCRIPTION:
%    Implemet the uBaL for various models (GLMs, Cox's model, ...). Can
%    deal with grouped variable and two-way interaction effects when accounting 
%    for the natural order between main and intertaction effects (main 
%    effects must be included into models before their interactions).
%    Note: When using the uBaL for selecting main and interaction effects, the
%    output fit.beta arranges the estimated coefficient vectors of effects 
%    in the order of the "adjacency" vector constructed by arranging out 
%    the "adjacency" matrix - an upper triangular matrix. That is, the
%    interaction between main effects i and j>=i will be stored at the
%    position J(i-1)+j-i(i-1)/2 in the adjacency vector, where J is the 
%    number of main effects.
%
% USAGE:                                                                 
%    fit = uBaL(X, y)
%    fit = uBaL(X, y, 'binomial', groups)
%
% INPUT ARGUMENTS:
% X           Input matrix of dimension nobs x nvars; each row is an
%             observation vector. Columns can be dummy variables.
% y           Response variable. Quantitative for family = 'normal'. For
%             family = 'binomial' should be a two-column matrix of counts
%             and frequencies.
% family      Reponse type, such as 'normal','binomial','poisson',... 
%             Default is 'normal'. The link function is assumed to be canonical.
%             Only GLMs are supported in this version.
% groups      a J-dimensional cell array indicating grouped variables, each
%             cell contains indices of the variables grouped. Default is 
%             not grouped. 
% interact    whether or not accounting for two-way interactions. Default is
%             'FALSE'.
% intercept   whether or not including the intercept into the model.
%             Default is 'TRUE'.
% r           the hyper-parameter in the Bayesian formulations. Default is .05
% burnIn      burn-in number used to do Gibbs sampling. Default is 1000.
% nSample     number of Gibbs samples to be drawn. Default is 10000.
%
% OUTPUT ARGUMENTS:
% fit         A structure.
% fit.beta    Estimated coefficients.
% fit.lambda  The vector of adaptive shrinkage parameters estimated 
%
% DATE: Jul 2010. Version 1
%
% AUTHOR: TRAN Minh Ngoc, Department of Statistics and Applied Probability,
% National University of Singapore, ngocmath@gmail.com
%
% REFERENCES:
%    Leng, Tran and Nott (2010). Bayesian Adaptive Lasso. 
%
% EXAMPLES:
% For variable selection in linear regression:
% n = 200; d = 8;
% X = randn(n,d);
% tbeta = [3;1.5;0;0;2;0;0;0];
% y = 5 + X*tbeta + randn(n,1);
% fit = uBaL(X,y);
% fit.beta
% For group selection in logistic regression:
% ng = 4; % number of groups;
% for i=1:ng groups{i}=[2*i-1,2*i]; end
% tbeta = [2,2,0,0,4,4,0,0]';
% mu=exp(5+X*tbeta)./(1+exp(5+X*tbeta));
% fre=10.*ones(n,1);
% y=[binornd(fre,mu),fre];
% fit = uBaL(X,y,'binomial',groups);
% fit.beta
%--------------------------------------------------------------------------

if (nargin<2)
   error('stats:glmfit:TooFewInputs','At least two arguments are required');
end
[nobs,nvar] = size(X);
if (nargin<3 | isempty(family)) family = 'normal'; end
if (nargin<4 | isempty(groups)) isgroup = 0; else isgroup = 1; end
if isgroup ind = [];
    for j = 1:length(groups)
        ind = [ind,groups{j}];
    end
    ind = sort(ind);
    if ((length(ind)~=nvar)||sum(1:nvar==ind)~=nvar) error('Grouping is not correct!'); end
end     
if (nargin<5 | isempty(interact)) interact = 0; end
if (nargin<6 | isempty(intercept)) intercept = 1; end
if (nargin<7 | isempty(r)) r = .01; end
if (nargin<8 | isempty(burnIn)) burnIn = 1000; end
if (nargin<9 | isempty(nSample)) nSample = 10000; end

if (~interact)
    if intercept [betaMLE,out1,out2] = glmfit(X,y,family,[],[],[],[],'on');
    else [betaMLE,out1,out2] = glmfit(X,y,family,[],[],[],[],'off'); end
    S=diag(out2.se)*out2.coeffcorr*diag(out2.se);
    IF=inv(S);
    Xtilde = sqrtm(IF);
    ytilde = Xtilde*betaMLE;

    if (~isgroup) lambda = gibbs_lasso(IF,betaMLE,r,burnIn,nSample);
        fit.lambda = lambda;
        fit.beta = BoostLasso(Xtilde,ytilde,lambda');
    end

    if (isgroup&intercept) newgroups{1} = 1;
        for j = 1:length(groups)
            newgroups{j+1} = 1+groups{j};
        end
        groups = newgroups;
        lambda = gibbs_glasso(Xtilde,ytilde,groups,r,burnIn,nSample);
        fit.lambda = lambda;
        lamb = lambda(1)*ones(1,length(groups{1}));
        for a = 2:length(groups)
            lamb = [lamb, lambda(a)*ones(1,length(groups{a}))];
        end
        Lamb = diag(lamb);
        Sigma = inv(Lamb)*Xtilde'*Xtilde*inv(Lamb);
        Q = inv(Lamb)*Xtilde'*ytilde;
        W = newton_grouplasso([],Sigma,Q,groups,1,1e-4);
        betahat = inv(Lamb)*W;
        betahat(betahat<1e-10) = 0;
        fit.beta = betahat;
    end

    if (isgroup&(~intercept)) lambda = gibbs_glasso(Xtilde,ytilde,groups,r,burnIn,nSample);
        fit.lambda = lambda;
        lamb = lambda(1)*ones(1,length(groups{1}));
        for a = 2:length(groups)
            lamb = [lamb, lambda(a)*ones(1,length(groups{1}))];
        end
        Lamb = diag(lamb);
        Sigma = inv(Lamb)*Xtilde'*Xtilde*inv(Lamb);
        Q = inv(Lamb)*Xtilde'*ytilde;
        W = newton_grouplasso([],Sigma,Q,groups,1,1e-4);
        betahat = inv(Lamb)*W;
        betahat(betahat<1e-10) = 0;
        fit.beta = betahat;
    end
end

if interact
    if (~isgroup) for j = 1:nvar 
            groups{j} = j;
        end
    end
    nme = length(groups); % number of main effects or grouped
    oldind = 1;
    XX = [];
    for i = 1:nme
        for j = i:nme
            if j==i newind = oldind+length(groups{i});
                newgroups{nme*(i-1)+j-i*(i-1)/2} = oldind:(newind-1);
                XX = [XX,X(:,groups{i})];
            else newind = oldind +length(groups{i})*length(groups{j});    
                newgroups{nme*(i-1)+j-i*(i-1)/2} = oldind:(newind-1);
                a = X(:,groups{i});
                b = X(:,groups{j});
                for ii = 1:length(groups{i})
                    for jj = 1:length(groups{j})
                        XX = [XX,a(:,ii).*b(:,jj)];
                    end
                end
            end
            oldind = newind;
        end
    end
    groups = newgroups; 
    ng = length(groups); % number of groups (with all main and interaction effects)
    for ig = 1:ng
        sg(ig) = length(groups{ig}); %size of groups
    end      
    for ig = 1:ng order{ig} = ig; end
    for ime = 1:nme
        MainEffectPosition = nme*(ime-1)+ime-ime*(ime-1)/2;
        index = order{MainEffectPosition};
        for j = (ime+1):nme
            index = [index,nme*(ime-1)+j-ime*(ime-1)/2];
        end
        for i = 1:(ime-1)
            index = [index,nme*(i-1)+ime-i*(i-1)/2];
        end
        order{MainEffectPosition} = index;
    end
    for i = 1:nme
        for j = i:nme
            parents{nme*(i-1)+j-i*(i-1)/2} = unique([nme*(i-1)+j-i*(i-1)/2,nme*(i-1)+i-i*(i-1)/2,nme*(j-1)+j-j*(j-1)/2]);
        end
    end
    
    if intercept newgroups{1} = 1;
        neworder{1} = 1;
        newparents{1} = 1;
        for j = 1:length(groups)
            newgroups{j+1} = 1+groups{j};
            neworder{j+1} = 1+order{j};
            newparents{j+1} = 1+parents{j};
        end
        groups = newgroups;
        order = neworder;
        parents = newparents;
    end
    if intercept [betaMLE,out1,out2] = glmfit(XX,y,family,[],[],[],[],'on');
    else [betaMLE,out1,out2] = glmfit(XX,y,family,[],[],[],[],'off'); end
    S=diag(out2.se)*out2.coeffcorr*diag(out2.se);
    IF=inv(S);
    Xtilde = sqrtm(IF);
    ytilde = Xtilde*betaMLE;

    lambda = gibbs_oglasso(Xtilde,ytilde,groups,order,parents,r,burnIn,nSample);
    fit.lambda = lambda;
    for ig = 1:length(lambda)
        if lambda(ig)>1e+10 lambda(ig) = 1e+10; end
    end
    fit.beta = BoostLasso(Xtilde,ytilde,lambda',groups,order);
end

function f = gibbs_lasso(IF,betaMLE,r,burnIn,nSample)   
%   helper function for the uBaL function, does Gibbs sampling for ungrouped Bayesian lasso
    delta = 1;
    mn = burnIn+nSample;
    [n,p] = size(IF);
    lambda2 = zeros(mn, p);
    beta = zeros(mn, p);
    tau2 = zeros(mn, p);
    lambda2(1,:) = ones(1,p);
    tau2(1,:) = ones(1,p);
    beta(1,:) = ones(1,p);
    xty = IF*betaMLE;
    for ii = 2:mn
        %% update beta
        A = IF+diag(1./tau2(ii-1,:));
        beta(ii,:) = mvnrnd(A\xty, inv(A));
        %% update tau2
        tmp = rinvnorm(p, sqrt(lambda2(ii-1,:)./beta(ii,:)./beta(ii,:)),lambda2(ii-1,:));
        tau2(ii,:) = 1./tmp;
        %% update lambda
        lambda2(ii,:) = gamrnd(1+r, 1./(tau2(ii,:)/2+delta));
        %% update delta by EB
        if floor(ii/30) ==  ii/30
            jj = ii-29;
            delta = p*r/sum(mean(lambda2(jj:ii,:),1));
        end
    end
    inc = (burnIn+1):mn;
    lambda2 = lambda2(inc,:);
    f = sqrt(mean(lambda2));
end

function f = gibbs_glasso(x,y,groups,r,burnIn,nSample)   
%   helper function for the uBaL function, does Gibbs sampling for
%   grouped Bayesian lasso
delta = 1;
mn = burnIn+nSample;
ng = length(groups);
for ig = 1:ng
    Xg{ig} = x(:,groups{ig});
    sg(ig) = length(groups{ig});
end
    lambda2 = zeros(mn, ng);
    beta = cell(ng,1);
    for a = 1:ng
        beta{a} = zeros(mn, sg(a));
    end
    tau2 = zeros(mn, ng);
    lambda2(1,:) = ones(1,ng);  
    tau2(1,:) = ones(1,ng);
    for a = 1:ng
        beta{a}(1,:) = ones(1,sg(a));
    end

    for ii = 2:mn
        %% update beta
        for a = 1:ng
            A = Xg{a}'*Xg{a}+(1/tau2(ii-1,a)).*eye(sg(a));
            b = y;
            for aa = 1:(a-1)
                b = b - Xg{aa}*(beta{aa}(ii,:))';
            end
            for aa = (a+1):ng
                b = b - Xg{aa}*(beta{aa}(ii-1,:))';
            end
            beta{a}(ii,:) = mvnrnd(A\(Xg{a}'*b), inv(A));
        end
        %% update tau2
        for a = 1:ng
            tmp = rinvnorm(1, sqrt(lambda2(ii-1,a))./norm(beta{a}(ii,:)),lambda2(ii-1,a));
            tau2(ii,a) = 1./tmp;
            tau2(ii,tau2(ii,:)<1e-50)=1e-50;
        end
        %% update lambda
        lambda2(ii,:) = gamrnd((r+(1+sg(:))./2)', 1./(tau2(ii,:)/2+delta));
        %% update delta by EB
        if floor(ii/30) ==  ii/30
            jj = ii-29;
            delta = ng*r/sum(mean(lambda2(jj:ii,:),1));
        end
    end
    inc = (burnIn+1):mn;
    lambda2 = lambda2(inc,:);
    f = sqrt(mean(lambda2));
end

function f = gibbs_oglasso(x,y,groups,order,parents,r,burnIn,nSample)   
%   helper function for the uBaL function, does Gibbs sampling for
%   ordered grouped Bayesian lasso
if (nargin < 6) r = .05; end
if (nargin < 7) burnIn = 1000; end
if (nargin < 8) nSample = 10000; end
delta = 1;
mn = burnIn+nSample;
ng = length(groups);

    for ig = 1:ng
        Xg{ig} = x(:,groups{ig});
        sg(ig) = length(groups{ig});
    end
  
    lambda2 = zeros(mn, ng);
    beta = cell(ng,1);
    for ig = 1:ng
        beta{ig} = zeros(mn, sg(ig));
    end
    tau2 = zeros(mn, ng);
    lambda2(1,:) = ones(1,ng);  
    tau2(1,:) = ones(1,ng);
    for ig = 1:ng
        beta{ig}(1,:) = ones(1,sg(ig));
    end

    for ii = 2:mn
        %% update beta
        for a = 1:ng
            sigma2j = sum(1./tau2(ii-1,parents{a})); 
            A = Xg{a}'*Xg{a}+sigma2j.*eye(sg(a));
            b = y;
            for aa = 1:(a-1)
                b = b - Xg{aa}*(beta{aa}(ii,:))';
            end
            for aa = (a+1):ng
                b = b - Xg{aa}*(beta{aa}(ii-1,:))';
            end
            beta{a}(ii,:) = mvnrnd(A\(Xg{a}'*b), inv(A));
        end
        %% update tau2
        for a = 1:ng
            b = [];
            for jj = 1:length(order{a})
                b = [b,beta{order{a}(jj)}(ii,:)];
            end
            tmp = rinvnorm(1, sqrt(lambda2(ii-1,a))./norm(b),lambda2(ii-1,a));
            tau2(ii,a) = 1./tmp;
            tau2(ii,tau2(ii,:)<1e-50)=1e-50;
        end
        %% update lambda
        for ig = 1:ng
            ss(ig) = sum(sg(order{ig}));
        end
        ss = ss';
        lambda2(ii,:) = gamrnd((r+(1+ss(:))./2)', 1./(tau2(ii,:)/2+delta));
        %% update delta by EB
        if floor(ii/30) ==  ii/30
            jj = ii-29;
            delta = ng*r/sum(mean(lambda2(jj:ii,:),1));
        end
    end
    inc = (burnIn+1):mn;
    lambda2 = lambda2(inc,:);
    lambda = sqrt(mean(lambda2)); 
f = lambda;
end

function [betahat,ite,excess] = BoostLasso(X,y,lambda,groups,order);
% using the boosting algorithm to search for the (ordered group) Lasso solution
% at the shrinkage vector lambda. The boosting loss function is the sum of 
% a (approximate) squared loss (induced by X and y) and a Lasso penlaty.
% groups: is a cell indicating grouped variables.
% order: is a cell indicating an ordering/hierarchy
if nargin < 4 grouped = 0;
    ordered = 0;
else grouped = 1;
    if nargin < 5 ordered = 0;
    else ordered = 1;
    end
end
[n,d] = size(X);
if grouped J = length(groups); else J = size(X,2); end
epsilon = 1e-10;
T = 500;
global currentbeta;
global Ij;
    function f = target(beta);
        if (~grouped) penalty = lambda'*abs(beta); end
        if (grouped&&(~ordered))
            penalty = 0;
            for j = 1:J
                penalty = penalty + lambda(j)*norm(beta(groups{j}));
            end
        end
        if (grouped&&(ordered))
            penalty = 0;
            for j = 1:J
                b = [];
                for i = 1:length(order{j})
                    b = [b,groups{order{j}(i)}];
                end
                penalty = penalty + lambda(j)*norm(beta(b));
            end
        end
        f = (y-X*beta)'*(y-X*beta)+penalty;
    end
    function f = fun(g)
        f = target(currentbeta+g.*Ij);
    end

currentbeta = zeros(d,1);
oldloss = target(currentbeta);
for jj = 1:d
    Ij = zeros(d,1); Ij(jj) = 1;
    [a,b] = fminsearch(@fun,0);
    ghat(jj) = a;
    loss(jj) = b;
end
[newloss,minindex] = min(loss);
Ij = zeros(d,1); Ij(minindex) = 1;
currentbeta = currentbeta+ghat(minindex).*Ij;
ite = 1;
excess = oldloss - newloss;
while (excess>epsilon)&&(ite<T)
    oldloss = newloss;
    for jj = 1:d
        Ij = zeros(d,1); Ij(jj) = 1;
        [ghat(jj),loss(jj)] = fminsearch(@fun,0);
    end
    [newloss,minindex] = min(loss);
    Ij = zeros(d,1); Ij(minindex) = 1;
    currentbeta = currentbeta+ghat(minindex).*Ij;
    excess = oldloss - newloss;
    ite = ite+1;
end
betahat = currentbeta;
end

end
        
    
