% gibbs sampling for hierarchical model
% Gaussian + smoothness + lam 
clear, close all
set(0,'DefaultAxesFontSize',14)

%%  ill-posed problem: first-order derivative
n = 100;
m = n;
A = ones(n);
A = tril(A);
A = A/norm(A);
tt= linspace(0,1,n)';
xe= zeros(n,1);
xe([10 20 50]) = 1;
be = A*xe;
sig = 1e-2*max(abs(be));
randn('seed',0)
b  = be + sig*randn(m,1);
figure(10), plot(1:n,xe,'r')

%% Gibbs sampling with hierarchical model
N = 1e4;
Ensx = zeros(m,N);
x   = zeros(m,1);      % initial guess for x
lam2 = 1e-1;
al0 = 1; bt0 = 1e-4;     % parameter for Gamma distribution
tau = 1/sig^2;
w = ones(m,1);
for i = 2:N
    % update x componentwise, given w
    for j = 1:m
        a0 = tau*sum(A(:,j).^2) + 1./w(j);
        mus = (b-A*x + x(j)*A(:,j));
        b0 = 2*tau*sum(A(:,j).*mus);
        mu = b0/(2*a0);
        sigx = sqrt(1/a0);
        x(j) = randn(1)*sigx+mu;
    end
    Ensx(:,i) = x;
    for j = 1:m 
        w(j) = gigrnd(0.5,lam2,x(j)^2,1);
    end
    figure(1),  plot(1:n,xe,'k',1:n,x,'r','linewidth',2), axis square, title(['iter =' num2str(i)]), drawnow
    if rem(i,1000) == 1, i, end
end

%% post processing
x_mean = mean(Ensx(:,100:N),2);
figure(2), plot(1:n,xe,1:n,x_mean,'linewidth',2), axis square
legend('exact','mean')
x_var  = var(Ensx(:,100:N),0,2);
figure(3), errorbar(1:n,x_mean,1.5*sqrt(x_var),'linewidth',2), axis([0 100 -0.5 1.5]), axis square
print(3,'-depsc','mcmc_gibbs_sparsity_mean')
hold on, plot(1:n,xe,'r','linewidth',2), hold off
figure(4), plot(1:N,Ensx(1,:),1:N,Ensx(n/2,:)), axis([0 N -4 4]), axis square
print(4,'-depsc','mcmc_gibbs_sparsity_trace')