function [p1, p2] = estKernelReg(frame1, frame2, b, sigma, penalty, nC, p1, p2)
% Created by Yuandong Tian, Oct. 5, 2009
% Email: yuandong@andrew.cmu.edu
% water tracking. minimize Eqn. 9 in the paper

%--------------------
%   best parameters for simulated data:
%  span = 1 & xsel = 102:135; ysel = 42:85 ->  
%           alpha = 20; sigma = 2, tau = 2; firstNC = 20; maxIter = 20;
%           no n divided on the gradient.
%  span = 2 & xsel = 102:135; ysel = 42:85 ->  
%           alpha = 50; sigma = 2-4-6, tau = 2; firstNC = 20; maxIter = 20-40;
%           perform better: sigma = 2, firstNC = 20 
%           even better:    sigma = 2, firstNC = 10;
%                           sigma = 2, firstNC = 5 is also ok, but slightly worse than first = 10;
%                           sigma = 2, firstNC = 3 the estimation blows up (much higher than true value, since they try to estimate everything using only the first several bases)
%
%           no penalty term (the key is not the penalty, but the number of basis vectors)
%           alpha = 100; sigma = 2, tau = 2, firstNC = 5-15 (i guess the best is 10)
%           (if firstNC is too large, it only hurts you by estimating those
%           coefficients in small bases, which cannot be compensate by more
%           iterations, too small then the coefficients are also not
%           trustible, since they tend to explain everything by themselves)
%  it seems like too many iterations with too many components will get you
%  screwed up. 
%  The best thing you can do is to pick bases one by one, if
%  the first n basis cannot intepret the data (or they can still decrease the
%  function, but has to do it in a "screwed"/somehow expensive way), then use the n+1, and so on.)
%  sounds like l1-norm. But i don't have time to go for it.
%  finally best:
%    span = 2, alpha = 100; no penalty. firstNC = 5-10. iter = 20
% 
% for real data, since there are some blurring happening, it is possible we
% could have higher alpha
%--------------------
[h, w] = size(frame1);

span = 2;

[X, Y] = meshgrid(1:w, 1:h);
X = X(:);
Y = Y(:);
subValid = (mod(X(:), span) == 0 & mod(Y(:), span) == 0);

if ~exist('nC', 'var')
    nC = b.nC;
end;

%n = h*w;
n = sum(subValid);

y1 = reshape(frame1, h*w, 1);
y1 = y1(subValid);
y2 = reshape(frame2, 1, h*w);
y2 = y2(subValid);

% repbxs = zeros(n, n, nC);
% repbys = zeros(n, n, nC);
% for k = 1:nC
%     repbxs(:, :, k) = repmat(b.bx(:, k), 1, n) / sigma^2;
%     repbys(:, :, k) = repmat(b.by(:, k), 1, n) / sigma^2;
% end;

% minimize the following function:
% J(p1, p2) = sum_i1 (y_i1 - \hat y_i using 2)^2 + symmetry

% 22 * 17 is the size of 102:135; ysel = 42:85, divided by span = 2. 
% it serves as the normalized term here. for patch with different size,
% alpha should change accordingly.
alpha = 3000 * (22 * 17) / n; 
% try smaller alpha;
% alpha = 500 * (22 * 17) / n; 
eps = 1e-10;
epsSparse = 1e-4;

if ~exist('p1', 'var')
    p1 = zeros(nC, 1);
end;
if ~exist('p2', 'var')
    p2 = zeros(nC, 1);
end;

if ~isscalar(penalty)
    penalty = penalty(1:nC);
else 
    penalty = repmat(penalty, nC, 1);
end;

cc1p1 = zeros(n, nC);
cc1p2 = zeros(n, nC);

cc2p1 = zeros(nC, n);
cc2p2 = zeros(nC, n);

prevFuncVal = Inf;

nIter = 1;
maxIter = 20;
doubleCount = 0;
doubleMaxCount = 5;

% figure;
while nIter < maxIter;
    % estimate w
    %
    pos1 = [b.bx(subValid, 1:nC) * p1 + X(subValid), b.by(subValid, 1:nC) * p1 + Y(subValid)];
    pos2 = [b.bx(subValid, 1:nC) * p2 + X(subValid), b.by(subValid, 1:nC) * p2 + Y(subValid)];
    
%     positive1 = (y1 > 0.5);
%     positive2 = (y2 > 0.5)';
%     plot(pos1(positive1, 1), pos1(positive1, 2), 'r+', pos1(~positive1, 1), pos1(~positive1, 2), 'b+');
%     hold on;
%     plot(pos2(positive2, 1), pos2(positive2, 2), 'ro', pos2(~positive2, 1), pos2(~positive2, 2), 'bo');
%     hold off;
%     drawnow;
    
    % estimate pairwise L2 distance
    [dx, dy] = dist2(pos1, pos2);
    distMat = dx.^2 + dy.^2;
    w = exp(-distMat / 2 / sigma^2);
%    w(w < epsSparse) = 0;
%    w = sparse(w);
    
    sumFor1 = sum(w, 2);
    sumFor2 = sum(w, 1);
    
    % objective function
    y1est = (w*y2') ./ sumFor1;
    y1err = y1est - y1;
    y2est = (y1'*w) ./ sumFor2;
    y2err = y2est - y2;
    
    funcVal = sum(y1err.^2) + sum(y2err.^2); % + p1' * (penalty .* p1) + p2' * (penalty .* p2);
    fprintf(1, '#Iter = %d, funcVal = %d', nIter, funcVal);
    if (prevFuncVal < funcVal)
        alpha = alpha / 2;
        fprintf(1, 'error! func ascending! alpha = %d', alpha);
        nIter = nIter - 1;
        doubleCount = 0;
    else
%         doubleCount = doubleCount + 1;
%         if (doubleCount > doubleMaxCount) 
%             alpha = alpha * 2;
%             doubleCount = 0;
%             fprintf(1, 'alpha = %d', alpha);
%         end;
    end;
    prevFuncVal = funcVal;
    
    % gradient for w_ij
    for k = 1:nC
%         wdx = w .* dx;
%         wdy = w .* dy;
%         
%         gp1 = - wdx .* repbxs(:, :, k) - wdy .* repbys(:, :, k);
%         gp2 = wdx .* repbxs(:, :, k)' + wdy .* repbys(:, :, k)';
        
        rect1 = repmat(b.bx(subValid, k), 1, n) .* dx + repmat(b.by(subValid, k), 1, n) .* dy;
        rect2 = dx .* repmat(b.bx(subValid, k)', n, 1) + dy .* repmat(b.by(subValid, k)', n, 1);
         
        gp1 = -w .* rect1 / sigma^2;
        gp2 = w .* rect2 / sigma^2;
        
        cc1p1(:, k) = (gp1 * y2' - y1est .* sum(gp1, 2) ) ./ sumFor1;
        cc2p1(k, :) = (y1' * gp1 - y2est .* sum(gp1, 1) ) ./ sumFor2; 
        
        cc1p2(:, k) = (gp2 * y2' - y1est .* sum(gp2, 2) ) ./ sumFor1;
        cc2p2(k, :) = (y1' * gp2 - y2est .* sum(gp2, 1) ) ./ sumFor2; 
    end;
    
    gradp1 = (y1err' * cc1p1)' + (cc2p1 * y2err'); % + penalty .* p1;
    gradp2 = (y1err' * cc1p2)' + (cc2p2 * y2err'); % + penalty .* p2;
    
    normg = gradp1' * gradp1 + gradp2' * gradp2;
    fprintf(1, '  normg = %d\n', normg);
    
    if (normg < eps)
        break;
    end;
    
    % gradient descent.
    p1 = p1 - alpha * gradp1;
    p2 = p2 - alpha * gradp2;
    nIter = nIter + 1;    
end;
%fprintf(1, '#Iter = %d, norm(g) = %d\n', nIter, normg);
% finally print pos1 and pos2

