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
