% read the image and the mask from input files
f = double(imread('corrupted.png'));
[ny, nx, c] = size(f);
N=nx*ny*c;
% mask(i,j) = 1 iff pixel (i,j) is corrupted.
mask = imread('mask.png') == zeros(ny, nx, c);

% vectorize image and mask
f=f(:);
mask = mask(:);
%%
% count corrupted pixels
M = sum(mask);

% construct a sparse linear operator for the color gradient
e = ones(ny, 1);
A = spdiags([-e e], 0:1, ny, ny);
A(end, end) = 0;
Dy = kron(speye(nx), A);
e = ones(nx, 1);
B = spdiags([-e e], 0:1, nx, nx)';
B(:, end) = 0;
Dx = kron(B', speye(ny));
Nabla_col = cat(1, kron(speye(3), Dx), kron(speye(3), Dy));

%%
% Let u_tilde be the vector containing only the unknown pixels that we
% optimize for and let u be the whole image that we are looking for.
% Construct operators X, Y s.t. u=X*u_tilde + Y*f is the reconstructed image.
X = sparse(find(mask), 1:M, 1, N, M);
Y = speye(N);
Y(mask, :) = 0;

% Construct a least squares problem min_{u_tilde} F(u_tilde) 
% with F(u_tilde)=|| A * u_tilde - b ||^2
A = Nabla_col * X;
b = -Nabla_col * Y * f;

%% solve
%[u,flag] = lsqr(A, b, 1e-1, 100);
% Solve least squares problem via first order optimality condition:
% \nabla F(x) = 0
u_tilde = mldivide(A'*A, A'*b);
% Inpaint unknown pixels in f
u = X*u_tilde + Y*f;

imshow([uint8(reshape(f, ny, nx, 3)) uint8(reshape(u, ny, nx, 3))]);