function [meanImg, reconEach] = reconMean(frames, ps, b, validframes, X, Y)
% Created by Yuandong Tian, Oct. 5, 2009
% Email: yuandong@andrew.cmu.edu
% reconstruct the image given the distortion. 

[m, n, nFrames] = size(frames);
if ~exist('validframes', 'var')
    % all frames are valid
    validframes = ones(nFrames, 1);
end;

if ~exist('X', 'var') && ~exist('Y', 'var')
    [X, Y] = meshgrid(1:n, 1:m);
    X = reshape(X, m*n, 1);
    Y = reshape(Y, m*n, 1);
    mMean = m;
    nMean = n;
else
    mMean = ceil(max(Y));
    nMean = ceil(max(X));
end;

if nargout == 2
    reconEach = zeros(mMean*nMean, nFrames);
    w = zeros(mMean*nMean, nFrames);
end;

meanImg = zeros(mMean*nMean, 1);
count = zeros(mMean*nMean, 1);

[Xdst, Ydst] = meshgrid(1:nMean, 1:mMean);
Xdst = Xdst(:);
Ydst = Ydst(:);

for i = 1:nFrames
    if ~validframes(i)
        continue;
    end;
    if isstruct(ps) 
        newx = reshape(ps.Xs(:, :, i), m*n, 1) + X;
        newy = reshape(ps.Ys(:, :, i), m*n, 1) + Y;
    elseif isstruct(b)
        nC = size(ps, 1);
        newx = b.bx(:, 1:nC) * ps(:, i) + X;
        newy = b.by(:, 1:nC) * ps(:, i) + Y;
    else
        newx = reshape(ps(:, :, i), m*n, 1) + X;
        newy = reshape(b(:, :, i), m*n, 1) + Y;
    end;
 
    % for mean, no need to do interpolation. simplest method works
    newxInt = round(newx);
    newyInt = round(newy);
    
    valid = (newxInt >= 1 & newxInt <= nMean & newyInt >= 1 & newyInt <= mMean);
    % note indices are repeatable.
    indices = sub2ind([mMean, nMean], newyInt(valid), newxInt(valid));
    thisFrame = reshape(frames(:, :, i), m, n);
    meanImg(indices) = meanImg(indices) + thisFrame(valid);
    count(indices) = count(indices) + 1;
    
    if nargout == 2
        % do spline interpolation. 
        %
       reconEach(indices, i) = thisFrame(valid);
       w(indices, i) = 1;
%        if mod(i, 10) == 0
%        fprintf(1, 'recon image = %d...\n', i);
%        end;
%        reconEach(:, i) = myinterp(newx, newy, thisFrame(:), Xdst, Ydst);
    end;
end;

eps = 1e-2;
meanImg = reshape(meanImg ./ (count + eps), mMean, nMean);
if nargout == 2
    reconEach = reshape(reconEach, [mMean, nMean, nFrames]);
    w = reshape(w, [mMean, nMean, nFrames]);
    % then do the weighted smoothing. 
    gf = fspecial('gaussian', [3, 3], 1);
    reconEach = imfilter(reconEach .* w, gf, 'same');
    w = imfilter(w, gf, 'same');
    reconEach = reconEach ./ w;
end;

function zi = myinterp(x, y, z, xi, yi)
% use 2D kernel regression, assume all vectors are column vectors 
% very slow so don't use it.
sigma = 1;
ni = length(xi);

zi = zeros(ni, 1);

for i = 1:ni
    if mod(i, 100) == 0
        fprintf(1, 'Interp point #%d...\n', i);
    end;
    dists = (xi(i) - x).^2 + (yi(i) - y).^2;
    w = exp(-dists / 2 / sigma^2);
    zi(i) = (z'*w) ./ sum(w);
end;