%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% clean up the workspace
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clear all ; close all ; clc
randn('seed' , 0);
rand('seed' , 0);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% generate excess density profile
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
N1 = 25;
N2 = 5;
[z1,z2] = meshgrid(1:N1,1:N2);
x = double((abs(z1-15) > 5) | (abs(z2-3.5) > 1));

figure();
imagesc(x);
colormap gray;
title('density profile, x');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% make noisy gravity measurements
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[L1,L2] = meshgrid(1:0.1:N1,(0:0.2:N2)/N2);
L = [L1(:)' ; L2(:)'];
m = size(L,2);
n = N1*N2;
perr = 0.0001;

y = nan(m,1);
theta = nan(m,n);
d = nan(m,n);
for i = 1:m
    y(i) = 0;
    for j = 1:n
        theta(i,j) = atan2(z2(j) + L(2,i) , z1(j) - L(1,i));
        d(i,j) = (z1(j) - L(1,i))^2 + (z2(j) + L(2,i))^2;
        y(i) = y(i) + cos(theta(i,j)) / d(i,j)^2 * x(j);
    end
end
y_true = y;
y = (1 - perr + 2 * perr * rand(m,1)) .* y;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% estimate the excess density profile
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
A = nan(m,n);
for i = 1:m
    for j = 1:n
        A(i,j) = cos(theta(i,j)) / d(i,j)^2;
    end
end
xhat_true = reshape(A\y_true , N2 , N1);
xhat = reshape(A\y , N2 , N1);

lambda = 5e-10;
Areg = [A ; sqrt(lambda) * eye(n)];
yreg = [y ; zeros(n,1)];
xhat_reg = reshape(Areg \ yreg , N2 , N1);

figure();
imagesc(xhat_true);
colormap gray;
title('xhat with exact measurements');

figure();
imagesc(xhat);
colormap gray;
title('xhat with noisy measurements');

figure();
imagesc(xhat_reg);
colormap gray;
title('regularized xhat with noisy measurements');