function [est_delta, Res_CI, Res_others, Res_diagnostic]  ...
             = Moreau_Approx_lm(y, X, gamma_0, Niter, sigmasq)
%Moreau_Approx_lm(y, X, gamma, Niter, sigmasq, theta_star)
% y, X are data
% gamma_0 controls the approx.
% Niter is the number of MCMC iterations
%sigma is noise level. Estimated from data if empty
% theta_star and delta_star are true values when available (mostly for
% testing)

%%Output:
%Res_beta: simulated beta
%Res_delta: simulated delta
% Res_o1 additional stats: Relative error, sensitivity, precision,
% sparsity, lambda_1, lambda_2, q, and timing.
%Note if beta_true and delta_true are empty, then first 3 columns of Res_o1
%are zero.
%Res_o2 contains acceptance probab. of the MCMC.



%data is n by p
n=size(X,1);
p=size(X,2);

BurnIn = ceil((1/5)*Niter);

%arguments for Moreau approx.
alpha=0.1; %value of alpha
display = 0; %verbose on not

if isempty(sigmasq) %then use empirical Bayes
    %Uses function lasso built-in Matlab
    if display ~= 0
        fprintf(1, '+------------------------------+\n');
        fprintf(1, '|      Estimating sigmasq      |\n');
        fprintf(1, '+------------------------------+\n');
    end
    [B,FitInfo] = lasso(X,y, 'CV', 10);
    betahat = B(:,FitInfo.IndexMinMSE); %or use FitInfo.IndexMinMSE
    nz = find(abs(betahat)>0);
    sigmasq = (1/(n- min(length(nz),n-1) ))*sum( (y-X(:,nz)*betahat(nz)).^2 );
end
S= (X')*X; lambda_max = sum(eig(S)); % max(eig(S)) or sum(eig(S))
M_lambda = (0.25/(1-alpha))*lambda_max;
gamma = gamma_0* (sigmasq/lambda_max);

if display ~= 0
    fprintf(1, '+---------------+\n');
    fprintf(1, '| Starting MCMC |\n');
    fprintf(1, '+---------------+\n');
end

tic;

%MCMC parameter needed
lsig=-2;
lsig_lbda = -3;
const_c = 5; %for Mala
mAcc1=0; mAcc2=0; mAcc_lbda = 0;


%prior q
u_q=1.5;

%prior lambda
a_lambda = 1e-5;

%Output
mean_theta=zeros(1,p);
mean_delta=zeros(1,p);
Res_theta = zeros(Niter,p); 
Res_others=zeros(Niter,3);
Res_ell = zeros(Niter,1);

%%Initialization
ypy=sum(y.^2); Xpy=(X')*y; XpX = (X')*X;
prior_q = 0.1; lambda_1 = 0.1; lambda_2 = 0.1;
theta_init=betahat;
theta=theta_init;
XpXt=XpX*theta;
ell_n= (0.5/sigmasq)*(ypy -2*(theta')*Xpy +(theta')*XpXt);
Grad = -(1/sigmasq)*(Xpy -XpXt);
logZ = logZ_fun_lm(lambda_1,lambda_2,sigmasq,alpha);

for kk = 1:Niter
    %Update delta
    tmp_J = theta - gamma * Grad;
    J = sign(tmp_J).*(abs(tmp_J)-alpha*lambda_1*gamma/sigmasq).*(abs(tmp_J) >= alpha*lambda_1*gamma/sigmasq)...
        ./( 1 + (1-alpha)*gamma*lambda_2/sigmasq );
    v = -(1/gamma)*tmp_J.*J +logZ +alpha*(lambda_1/sigmasq)*abs(J) ...
        + 0.5*((1-alpha)*(lambda_2/sigmasq) + (1/gamma) )*(J.^2);
    r = log(prior_q) - log(1-prior_q) +0.5*log(2*pi*gamma) - v;
    prob = 1./(1 + exp(-r));
    delta = (rand(p,1)<=prob);
    %COmpute h_n and J
    h_n = ell_n + (0.5/gamma)*sum(theta.^2) -sum(theta.*Grad) + sum(v(delta));
    J = zeros(p,1);
    ind = find(delta==1);
    v = tmp_J(ind);
    J(ind) = sign(v).*(abs(v)-alpha*lambda_1*gamma/sigmasq).*(abs(v) >= alpha*lambda_1*gamma/sigmasq)...
        ./( 1 + (1-alpha)*gamma*lambda_2/sigmasq );
%   %Update the selected components of theta Using a MaLa-style sampler
    if ~isempty(ind)
        Q=length(ind);
            for jj = 1:Q
                G=(1/gamma)*(theta(ind(jj))-J(ind(jj)));
                G=(const_c/(max(const_c,abs(G))))*G;
                theta_prop=theta;
                var_th = - 0.5*exp(2*lsig)*G + exp(lsig)*randn(1,1);
                theta_prop(ind(jj)) = theta(ind(jj)) + var_th;
                Grad_prop = Grad +(var_th/sigmasq)*XpX(:,ind(jj));
                J_prop=zeros(p,1); tmpJ_prop = theta_prop(ind)-gamma*Grad_prop(ind);
                J_prop(ind) = sign(tmpJ_prop).*(abs(tmpJ_prop)- alpha*gamma*(lambda_1/sigmasq))...
                    .*(abs(tmpJ_prop) >= alpha*gamma*(lambda_1/sigmasq))...
                        /(1+(1-alpha)*gamma*(lambda_2/sigmasq) );
                ell_nprop = ell_n + var_th*Grad(ind(jj)) + (0.5*var_th^2/sigmasq)*XpX(ind(jj),ind(jj));
                h_nprop = ell_nprop + sum(Grad_prop.*(J_prop-theta_prop)) + Q*logZ...
                    + alpha*(lambda_1/sigmasq)*sum(abs(J_prop(ind))) ...
                    +0.5*(1-alpha)*(lambda_2/sigmasq)*sum(J_prop(ind).^2) ...
                    +(0.5/gamma)*sum((J_prop-theta_prop).^2);
                G_prop = (1/gamma)*(theta_prop(ind(jj))- J_prop(ind(jj)));
                G_prop = (const_c/(max(const_c,abs(G_prop))))*G_prop;
                r = -0.5*exp(-2*lsig)*( (-var_th + 0.5*exp(2*lsig)*G_prop)^2 -(var_th + 0.5*exp(2*lsig)*G)^2 );
                r = r - (h_nprop-h_n);
                Acc = min(1,exp(r));
                if rand(1)<= Acc
                    theta=theta_prop; 
                    Grad=Grad_prop; ell_n = ell_nprop;
                    J=J_prop; h_n = h_nprop;
                end
                mAcc1 = mAcc1 + (1/kk)*(Acc - mAcc1);
                lsig = lsig +(1/kk^0.55)*(Acc - 0.4);
            end
    end
    
    %%Update the non-selected components
    ind = find(delta==0);
    if ~isempty(ind)
        ind_sel = find(delta==1);
        Prox_theta = zeros(p,1);
        Prox_theta(ind_sel) = sign(theta(ind_sel)).*(abs(theta(ind_sel)) - alpha*gamma*(lambda_1/sigmasq))...
                    .*(abs(theta(ind_sel)) >= alpha*gamma*(lambda_1/sigmasq))/(1+(1-alpha)*gamma*(lambda_2/sigmasq) );
        omega = -sum(Grad(ind_sel).*(J(ind_sel)-Prox_theta(ind_sel))) ...
            - alpha*(lambda_1/sigmasq)*sum(abs(J(ind_sel))-abs(Prox_theta(ind_sel)))...
            - 0.5*(1-alpha)*(lambda_2/sigmasq)*sum(J(ind_sel).^2-Prox_theta(ind_sel).^2)...
            -(0.5/gamma)*sum(J(ind_sel).^2-Prox_theta(ind_sel).^2 ...
            - 2*(J(ind_sel)-Prox_theta(ind_sel)).*theta(ind_sel));
        M=eye(length(ind)) - (gamma/sigmasq)*XpX(ind,ind);
        chol_M = chol(M,'upper');
        mu = chol_M\(chol_M'\( (1/sigmasq)*(XpX(ind,ind_sel))*(Prox_theta(ind_sel)-theta(ind_sel))));
        mu=gamma*mu;
        theta_prop = theta;
        theta_prop(ind) = mu + sqrt(gamma)*(chol_M\randn(length(ind),1));
        [ell_nprop, h_nprop, Grad_prop, J_prop] = grad_fun_lm(theta_prop, delta, lambda_1, lambda_2, sigmasq, gamma, alpha, ypy, Xpy, XpX);
        Prox_theta_prop = Prox_theta;
        Prox_theta_prop(ind_sel) = sign(theta_prop(ind_sel)).*(abs(theta_prop(ind_sel)) - alpha*gamma*(lambda_1/sigmasq))...
                    .*(abs(theta_prop(ind_sel)) >= alpha*gamma*(lambda_1/sigmasq))/(1+(1-alpha)*gamma*(lambda_2/sigmasq) );
        omega_prop = -sum(Grad_prop(ind_sel).*(J_prop(ind_sel) - Prox_theta_prop(ind_sel)))...
            - alpha*(lambda_1/sigmasq)*sum(abs(J_prop(ind_sel))-abs(Prox_theta_prop(ind_sel)))...
            - 0.5*(1-alpha)*(lambda_2/sigmasq)*sum(J_prop(ind_sel).^2-Prox_theta_prop(ind_sel).^2)...
            -(0.5/gamma)*sum(J_prop(ind_sel).^2 - Prox_theta_prop(ind_sel).^2 ...
            -2*(J_prop(ind_sel)-Prox_theta_prop(ind_sel)).*theta_prop(ind_sel));
        
        Acc = min(1,exp(omega_prop - omega));
        if rand(1) <=Acc
            theta=theta_prop; ell_n = ell_nprop;
            h_n=h_nprop;
            Grad= Grad_prop;
            J=J_prop;
        end
        mAcc2 = mAcc2 +(1/kk)*(Acc - mAcc2);
    end

    %Update hyper param
    
    %Update lambda1 and lambda2
     lambda_prop = [lambda_1,lambda_2]  + exp(lsig_lbda)*randn(1,2);
     if lambda_prop(1)<a_lambda || lambda_prop(1) >M_lambda || lambda_prop(2)<a_lambda || lambda_prop(2)>M_lambda
         Acc =0 ;
     else

        logZ_prop = logZ_fun_lm(lambda_prop(1),lambda_prop(2),sigmasq,alpha);
        ind = find(delta==1);
        v = theta(ind) -gamma*Grad(ind);
        J_prop = zeros(p,1);
        J_prop(ind) = sign(v).*(abs(v)-alpha*lambda_prop(1)*gamma/sigmasq)...
            .*(abs(v) >= alpha*lambda_prop(1)*gamma/sigmasq)...
            ./( 1 + (1-alpha)*gamma*lambda_prop(2)/sigmasq );
        r1 = -(1/gamma)*sum(v.*(J(ind)-J_prop(ind))) - sum(delta)*(logZ_prop-logZ) ...
            -0.5*((1-alpha)/sigmasq)*(lambda_prop(2)*sum(J_prop(ind).^2) - lambda_2*sum(J(ind).^2))...
            -(0.5/gamma)*(sum(J_prop(ind).^2) - sum(J(ind).^2))...
            -(alpha/sigmasq)*(lambda_prop(1)*sum(abs(J_prop(ind))) - lambda_1*sum(abs(J(ind))));
        Acc = min(1,exp(r1 ));
        if rand(1)<= Acc
            lambda_1 = lambda_prop(1); 
            lambda_2 = lambda_prop(2);
            logZ=logZ_prop;
        end
     end
    mAcc_lbda = mAcc_lbda + (1/kk)*(Acc - mAcc_lbda);
    lsig_lbda = lsig_lbda +(1/kk^0.55)*(Acc - 0.4);

    %Update prior_q

    prior_q = betarnd(sum(delta)+1,p+p^u_q-sum(delta));

    
    %%Collect information
    
    Res_theta(kk,:) = theta';
    Res_delta(kk,:) = delta';
    if kk>BurnIn
        mean_theta = mean_theta + (1/(kk-BurnIn))*(theta' - mean_theta);
        mean_delta = mean_delta + (1/(kk-BurnIn))*(delta' - mean_delta);
    end
    Res_others(kk,1) = lambda_1;
    Res_others(kk,2) = lambda_2;
    Res_others(kk,3) = prior_q;
    Res_ell(kk) = ell_n;
    %Timing
    Res_others(kk,7) = toc;
    
    
    if display ~= 0
      [kk*1e-3, Res_others(kk,:), mAcc1, mAcc2, mAcc_lbda]
    end
    
end

Conv =  Geweke_Conv(Res_ell((BurnIn+1):Niter));
Res_CI = zeros(p,3);
est_delta = [(mean_delta  > 0.5)',mean_delta'];
ind  = find( mean_delta > 0.5);
len = length(ind);
if len >0
    for jj = 1:len
        dt = Res_theta((BurnIn+1):Niter,ind(jj));
        CI = [quantile(dt,0.025),mean(dt),quantile(dt,0.975)];
        Res_CI(ind(jj),:) = CI;
    end
end

Res_diagnostic = [Conv(1), mAcc1, mAcc2, mAcc_lbda, sigmasq];


end



function [val0, val, grad, J] = grad_fun_lm(theta, delta, lambda_1, lambda_2, sigmasq, gamma, alpha, ypy, Xpy, XpX)
%Compute \ell_n and nabla \ell_n
% however val does not include the log-norm. constant

logZ = logZ_fun_lm(lambda_1,lambda_2,sigmasq,alpha);
p = size(XpX,1);
XpXt=XpX*theta;
val0= (0.5/sigmasq)*(ypy -2*(theta')*Xpy +(theta')*XpXt);
grad = -(1/sigmasq)*(Xpy -XpXt);

ind_sel = find(delta==1);
tmp_J=theta - gamma*grad;
J=zeros(p,1);
J(ind_sel) = sign(tmp_J(ind_sel)).*(abs(tmp_J(ind_sel)) - alpha*gamma*(lambda_1/sigmasq))...
    .*(abs(tmp_J(ind_sel)) >= alpha*gamma*(lambda_1/sigmasq))...
            /(1+(1-alpha)*gamma*(lambda_2/sigmasq) );
val = val0  + sum(grad.*(J-theta)) + length(ind_sel)*logZ + alpha*(lambda_1/sigmasq)*sum(abs(J(ind_sel))) ...
            +0.5*(1-alpha)*(lambda_2/sigmasq)*sum(J(ind_sel).^2) ...
            +(0.5/gamma)*sum((J-theta).^2);

end


function logZ = logZ_fun_lm(lambda_1, lambda_2, sigmasq, alpha)

x=alpha*(lambda_1/sigmasq)/sqrt(2*(1-alpha)*(lambda_2/sigmasq));
logZ = 0.5*log(2*pi/((1-alpha)*(lambda_2/sigmasq))) + log(erfcx(x));

end

