function calculate_collection_area2(Z,NX,NY)
%input/must be set beforehand: ZS pool drain drainage;
%output : A = collection area ~ amount of water
global A ZS pool drain drainage;
flowU = zeros(NX,NY); % there is some flow that goes up UP from (i,j) to (i,j+1)
flowD = zeros(NX,NY); % there is some flow that goes up DOWN from (i,j) to (i,j+1)
flowL = zeros(NX,NY); % there is some flow that goes up LEFT from (i,j) to (i,j+1)
flowR = zeros(NX,NY); % there is some flow that goes up RIGHT from (i,j) to (i,j+1)
A = zeros(NX,NY);
%loop over all cells, from highest to lowest using sorted array ZS
for i = NX*NY:-1:1
x = ZS(i,2);
y = ZS(i,3);
xU = mod(x-1-1,NX)+1; % normally x-1 but observe p.b.c.
xD = mod(x+1-1,NX)+1; % normally x+1 but observe p.b.c.
yL = mod(y-1-1,NY)+1; % normally y-1 but observe p.b.c.
yR = mod(y+1-1,NY)+1; % normally y+1 but observe p.b.c.
%compute fraction of flow to all lower cells
sum = 0;
if (Z(x,y)>Z(xU,y)) sum = sum+Z(x,y)-Z(xU,y); end
if (Z(x,y)>Z(xD,y)) sum = sum+Z(x,y)-Z(xD,y); end
if (Z(x,y)>Z(x,yR)) sum = sum+Z(x,y)-Z(x,yR); end
if (Z(x,y)>Z(x,yL)) sum = sum+Z(x,y)-Z(x,yL); end
if sum>0
if (Z(x,y)>Z(xU,y)) flowU(x,y) = (Z(x,y)-Z(xU,y))/sum; end
if (Z(x,y)>Z(xD,y)) flowD(x,y) = (Z(x,y)-Z(xD,y))/sum; end
if (Z(x,y)>Z(x,yR)) flowR(x,y) = (Z(x,y)-Z(x,yR))/sum; end
if (Z(x,y)>Z(x,yL)) flowL(x,y) = (Z(x,y)-Z(x,yL))/sum; end
end
% the first value is rain on the area (multiply by dx^2 outside)
% I can only receive water from higher cells, for the flow arrays have
% been computed.
if drain(x,y)
A(x,y) = 1 + flowL(x,yR)*A(x,yR) + flowR(x,yL)*A(x,yL) + flowU(xD,y)*A(xD,y) + flowD(xU,y)*A(xU,y);
elseif pool(x,y) % if this cell is part of a pool (boundary), add all water that flows into this cell to the drainage point
A(drainage == pool(x,y)) = A(drainage == pool(x,y)) + flowL(x,yR)*A(x,yR) + flowR(x,yL)*A(x,yL) + flowU(xD,y)*A(xD,y) + flowD(xU,y)*A(xU,y);
elseif drainage(x,y) % if this is drainage point of a pool,
A(x,y) = 1 + A(x,y) + flowL(x,yR)*A(x,yR) + flowR(x,yL)*A(x,yL) + flowU(xD,y)*A(xD,y) + flowD(xU,y)*A(xU,y);
end
end