function [Z,xdot,ydot,zdot] = focalflow(I1,I2,I3,params)
% This function computes the focal flow given three images and parameters.
% INPUT:
% I1, I2, I3: HxW double images
% params: a structure specifying the following parameters
% Sigma: the filter standard deviation
% mu_s: sensor distance
% mu_f: focal distance
% xp, yp: principal point (default to the center)
% szxW, szyW: half size of window function (default to 50)
% mask: a HxW 0,1 mask to remove some areas like background, shadow, etc. (default to all 1)

% OUTPUT:
% Z: depth
% xdot, ydot, zdot: velocity
%
%% Initialization
Sigma = params.Sigma;mu_s = params.mu_s;mu_f = params.mu_f;
if isfield(params,'szxW')
    szxW = params.szxW;
    szyW = params.szyW;
else
    szxW = 35;
    szyW = 35;
end
diffKernels = [-1/2,0,1/2];
im = cell(3,1);im{1} = I1;im{2} = I2;im{3} = I3;
H = size(I1,1);W = size(I1,2);
if isfield(params,'xp')
    xp = params.xp;
    yp = params.yp;
else
    xp = floor(W/2);
    yp = floor(H/2);
end
% Remove the edge
szx = xp - 20; szy = yp - 20;
if isfield(params,'mask')
    mask = params.mask;
else
    mask = ones(H,W);
end
%% Temporal difference
imT = (im{3} - im{1})/2;
imT = imT(yp-szy:yp+szy, xp-szx:xp+szx);

%% Spatial difference
[XX,YY] = meshgrid(-szx+xp-xp:szx+xp-xp,-szy+yp-yp:szy+yp-yp);
ratio = 10^4; % normalization factor to avoid numerical instability
XX = XX ./ ratio;
YY = YY ./ ratio;

imX = conv2(im{2},-diffKernels,'same');
imY = conv2(im{2},-diffKernels','same');
imXX = conv2(imX,-diffKernels,'same');
imYY = conv2(imY,-diffKernels','same');
% Crop only the necessary part
imX = imX(yp-szy:yp+szy, xp-szx:xp+szx) .* mask(yp-(szy):yp+(szy), xp-(szx):xp+(szx));
imY = imY(yp-szy:yp+szy, xp-szx:xp+szx) .* mask(yp-(szy):yp+(szy), xp-(szx):xp+(szx));
imXX = imXX(yp-szy:yp+szy, xp-szx:xp+szx) .* mask(yp-(szy):yp+(szy), xp-(szx):xp+(szx));
imYY = imYY(yp-szy:yp+szy, xp-szx:xp+szx) .* mask(yp-(szy):yp+(szy), xp-(szx):xp+(szx));
XimX = imX .* XX;
YimY = imY .* YY;

% Precompute entries of the normal equation A'Ax=A'b (Ix^2, IxIy, Ix(xIx + yIy), etc.)
% at every pixel. These will be windowed momentarily.
A11 = imX .* imX;
A12 = imX .* imY;
A13 = imX .* (XimX + YimY);
A14 = imX .* (imXX + imYY);

A22 = imY .* imY;
A23 = imY .* (XimX + YimY);
A24 = imY .* (imXX + imYY);

A33 = (XimX + YimY) .* (XimX + YimY);
A34 = (XimX + YimY) .* (imXX + imYY);

A44 = (imXX + imYY) .* (imXX + imYY);

b1 = imX .* imT;
b2 = imY .* imT;
b3 = (XimX + YimY) .* imT;
b4 = (imXX + imYY) .* imT;


mask = mask(yp-(szy-szyW):yp+(szy-szyW), xp-(szx - szxW):xp+(szx - szxW));
mask = reshape(mask,(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1);
W = ones(szyW * 2 + 1, szxW * 2 + 1);

% Window the entries by convolution with a windowing function W.
% ATA and ATb are big matrices storing all the windowed entries of A'A and A'b for each pixel

ATA = zeros(4,4,(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1));
ATb = zeros(1,4,(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1));
ATA(1,1,:) = reshape(conv2(A11,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(1,2,:) = reshape(conv2(A12,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(1,3,:) = reshape(conv2(A13,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(1,4,:) = reshape(conv2(A14,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(2,1,:) = ATA(1,2,:);
ATA(2,2,:) = reshape(conv2(A22,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(2,3,:) = reshape(conv2(A23,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(2,4,:) = reshape(conv2(A24,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(3,1,:) = ATA(1,3,:);
ATA(3,2,:) = ATA(2,3,:);
ATA(3,3,:) = reshape(conv2(A33,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(3,4,:) = reshape(conv2(A34,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATA(4,1,:) = ATA(1,4,:);
ATA(4,2,:) = ATA(2,4,:);
ATA(4,3,:) = ATA(3,4,:);
ATA(4,4,:) = reshape(conv2(A44,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;

ATb(1,1,:) = reshape(conv2(b1,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATb(1,2,:) = reshape(conv2(b2,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATb(1,3,:) = reshape(conv2(b3,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;
ATb(1,4,:) = reshape(conv2(b4,W,'valid'),(2*(szx - szxW) + 1) * (2*(szy - szyW) + 1),1) .* mask;

%% Computing the u vector, the u = [u_1,u_2,u_3,v]'
% Invert all per­pixel 4x4 matrices in parallel using cofactors
detATA = ATA(1,1,:) .* ATA(2,2,:) .* ATA(3,3,:) .* ATA(4,4,:) + ...
    ATA(1,1,:) .* ATA(2,3,:) .* ATA(3,4,:) .* ATA(4,2,:) + ...
    ATA(1,1,:) .* ATA(2,4,:) .* ATA(3,2,:) .* ATA(4,3,:) + ...
    ...
    ATA(1,2,:) .* ATA(2,1,:) .* ATA(3,4,:) .* ATA(4,3,:) + ...
    ATA(1,2,:) .* ATA(2,3,:) .* ATA(3,1,:) .* ATA(4,4,:) + ...
    ATA(1,2,:) .* ATA(2,4,:) .* ATA(3,3,:) .* ATA(4,1,:) + ...
    ...
    ATA(1,3,:) .* ATA(2,1,:) .* ATA(3,2,:) .* ATA(4,4,:) + ...
    ATA(1,3,:) .* ATA(2,2,:) .* ATA(3,4,:) .* ATA(4,1,:) + ...
    ATA(1,3,:) .* ATA(2,4,:) .* ATA(3,1,:) .* ATA(4,2,:) + ...
    ...
    ATA(1,4,:) .* ATA(2,1,:) .* ATA(3,3,:) .* ATA(4,2,:) + ...
    ATA(1,4,:) .* ATA(2,2,:) .* ATA(3,1,:) .* ATA(4,3,:) + ...
    ATA(1,4,:) .* ATA(2,3,:) .* ATA(3,2,:) .* ATA(4,1,:) - ...
    ...
    ATA(1,1,:) .* ATA(2,2,:) .* ATA(3,4,:) .* ATA(4,3,:) - ...
    ATA(1,1,:) .* ATA(2,3,:) .* ATA(3,2,:) .* ATA(4,4,:) - ...
    ATA(1,1,:) .* ATA(2,4,:) .* ATA(3,3,:) .* ATA(4,2,:) - ...
    ...
    ATA(1,2,:) .* ATA(2,1,:) .* ATA(3,3,:) .* ATA(4,4,:) - ...
    ATA(1,2,:) .* ATA(2,3,:) .* ATA(3,4,:) .* ATA(4,1,:) - ...
    ATA(1,2,:) .* ATA(2,4,:) .* ATA(3,1,:) .* ATA(4,3,:) - ...
    ...
    ATA(1,3,:) .* ATA(2,1,:) .* ATA(3,4,:) .* ATA(4,2,:) - ...
    ATA(1,3,:) .* ATA(2,2,:) .* ATA(3,1,:) .* ATA(4,4,:) - ...
    ATA(1,3,:) .* ATA(2,4,:) .* ATA(3,2,:) .* ATA(4,1,:) - ...
    ...
    ATA(1,4,:) .* ATA(2,1,:) .* ATA(3,2,:) .* ATA(4,3,:) - ...
    ATA(1,4,:) .* ATA(2,2,:) .* ATA(3,3,:) .* ATA(4,1,:) - ...
    ATA(1,4,:) .* ATA(2,3,:) .* ATA(3,1,:) .* ATA(4,2,:);

B = zeros(4,4,size(ATA,3));
B(1,1,:) = ATA(2,2,:) .* ATA(3,3,:) .* ATA(4,4,:) ...
    + ATA(2,3,:) .* ATA(3,4,:) .* ATA(4,2,:) ...
    + ATA(2,4,:) .* ATA(3,2,:) .* ATA(4,3,:) ...
    - ATA(2,2,:) .* ATA(3,4,:) .* ATA(4,3,:) ...
    - ATA(2,3,:) .* ATA(3,2,:) .* ATA(4,4,:) ...
    - ATA(2,4,:) .* ATA(3,3,:) .* ATA(4,2,:);

B(1,2,:) = ...
      ATA(1,2,:) .* ATA(3,4,:) .* ATA(4,3,:) ...
    + ATA(1,3,:) .* ATA(3,2,:) .* ATA(4,4,:) ...
    + ATA(1,4,:) .* ATA(3,3,:) .* ATA(4,2,:) ...
    - ATA(1,2,:) .* ATA(3,3,:) .* ATA(4,4,:) ...
    - ATA(1,3,:) .* ATA(3,4,:) .* ATA(4,2,:) ...
    - ATA(1,4,:) .* ATA(3,2,:) .* ATA(4,3,:);

B(1,3,:) = ...
      ATA(1,2,:) .* ATA(2,3,:) .* ATA(4,4,:) ...
    + ATA(1,3,:) .* ATA(2,4,:) .* ATA(4,2,:) ...
    + ATA(1,4,:) .* ATA(2,2,:) .* ATA(4,3,:) ...
    - ATA(1,2,:) .* ATA(2,4,:) .* ATA(4,3,:) ...
    - ATA(1,3,:) .* ATA(2,2,:) .* ATA(4,4,:) ...
    - ATA(1,4,:) .* ATA(2,3,:) .* ATA(4,2,:);

B(1,4,:) = ...
      ATA(1,2,:) .* ATA(2,4,:) .* ATA(3,3,:) ...
    + ATA(1,3,:) .* ATA(2,2,:) .* ATA(3,4,:) ...
    + ATA(1,4,:) .* ATA(2,3,:) .* ATA(3,2,:) ...
    - ATA(1,2,:) .* ATA(2,3,:) .* ATA(3,4,:) ...
    - ATA(1,3,:) .* ATA(2,4,:) .* ATA(3,2,:) ...
    - ATA(1,4,:) .* ATA(2,2,:) .* ATA(3,3,:);

B(2,2,:) = ...
      ATA(1,1,:) .* ATA(3,3,:) .* ATA(4,4,:) ...
    + ATA(1,3,:) .* ATA(3,4,:) .* ATA(4,1,:) ...
    + ATA(1,4,:) .* ATA(3,1,:) .* ATA(4,3,:) ...
    - ATA(1,1,:) .* ATA(3,4,:) .* ATA(4,3,:) ...
    - ATA(1,3,:) .* ATA(3,1,:) .* ATA(4,4,:) ...
    - ATA(1,4,:) .* ATA(3,3,:) .* ATA(4,1,:);

B(2,3,:) = ...
      ATA(1,1,:) .* ATA(2,4,:) .* ATA(4,3,:) ...
    + ATA(1,3,:) .* ATA(2,1,:) .* ATA(4,4,:) ...
    + ATA(1,4,:) .* ATA(2,3,:) .* ATA(4,1,:) ...
    - ATA(1,1,:) .* ATA(2,3,:) .* ATA(4,4,:) ...
    - ATA(1,3,:) .* ATA(2,4,:) .* ATA(4,1,:) ...
    - ATA(1,4,:) .* ATA(2,1,:) .* ATA(4,3,:);

B(2,4,:) = ...
      ATA(1,1,:) .* ATA(2,3,:) .* ATA(3,4,:) ...
    + ATA(1,3,:) .* ATA(2,4,:) .* ATA(3,1,:) ...
    + ATA(1,4,:) .* ATA(2,1,:) .* ATA(3,3,:) ...
    - ATA(1,1,:) .* ATA(2,4,:) .* ATA(3,3,:) ...
    - ATA(1,3,:) .* ATA(2,1,:) .* ATA(3,4,:) ...
    - ATA(1,4,:) .* ATA(2,3,:) .* ATA(3,1,:);

B(3,3,:) = ...
      ATA(1,1,:) .* ATA(2,2,:) .* ATA(4,4,:) ...
    + ATA(1,2,:) .* ATA(2,4,:) .* ATA(4,1,:) ...
    + ATA(1,4,:) .* ATA(2,1,:) .* ATA(4,2,:) ...
    - ATA(1,1,:) .* ATA(2,4,:) .* ATA(4,2,:) ...
    - ATA(1,2,:) .* ATA(2,1,:) .* ATA(4,4,:) ...
    - ATA(1,4,:) .* ATA(2,2,:) .* ATA(4,1,:);

B(3,4,:) = ...
      ATA(1,1,:) .* ATA(2,4,:) .* ATA(3,2,:) ...
    + ATA(1,2,:) .* ATA(2,1,:) .* ATA(3,4,:) ...
    + ATA(1,4,:) .* ATA(2,2,:) .* ATA(3,1,:) ...
    - ATA(1,1,:) .* ATA(2,2,:) .* ATA(3,4,:) ...
    - ATA(1,2,:) .* ATA(2,4,:) .* ATA(3,1,:) ...
    - ATA(1,4,:) .* ATA(2,1,:) .* ATA(3,2,:);

B(4,4,:) = ...
      ATA(1,1,:) .* ATA(2,2,:) .* ATA(3,3,:) ...
    + ATA(1,2,:) .* ATA(2,3,:) .* ATA(3,1,:) ...
    + ATA(1,3,:) .* ATA(2,1,:) .* ATA(3,2,:) ...
    - ATA(1,1,:) .* ATA(2,3,:) .* ATA(3,2,:) ...
    - ATA(1,2,:) .* ATA(2,1,:) .* ATA(3,3,:) ...
    - ATA(1,3,:) .* ATA(2,2,:) .* ATA(3,1,:);

B(2,1,:) = B(1,2,:);
B(3,1,:) = B(1,3,:);
B(3,2,:) = B(2,3,:);
B(4,1,:) = B(1,4,:);
B(4,2,:) = B(2,4,:);
B(4,3,:) = B(3,4,:);

% Clean the space
ATA = [];
IdetATA = 1./ detATA;
flag = reshape(IdetATA > 0.000001,size(B,3),1);
IdetATA = repmat(reshape(IdetATA,1,1,size(B,3)),4,4,1);
B = IdetATA .* B;
IdetATA = [];
B = B .* repmat(ATb,4,1,1);
u = reshape(sum(B,2),4,size(B,3))';
u(flag == 0,:) = 0;
u(:,3) = u(:,3) ./ ratio; %Recover from the normalization

% Recover Z, xdot, ydot, zdot
Z = mu_s.^2 .* Sigma.^2 .* u(:,3) ./ (mu_s.^2 .* Sigma.^2 .* u(:,3)./ mu_f - mu_f .* u(:,4));

xdot = -u(:,1) .* Z ./ mu_s;
ydot = -u(:,2) .* Z ./ mu_s;
zdot = -u(:,3) .* Z;
Z = reshape(Z,(szy - szyW) * 2 + 1,(szx - szxW) * 2 + 1);
xdot = reshape(xdot,(szy - szyW) * 2 + 1,(szx - szxW) * 2 + 1);
ydot = reshape(ydot,(szy - szyW) * 2 + 1,(szx - szxW) * 2 + 1);
zdot = reshape(zdot,(szy - szyW) * 2 + 1,(szx - szxW) * 2 + 1);
end