function c = dot(a,b,dim)
%DOT Vector dot product of codistributed array
%   C = DOT(A,B)
%   C = DOT(A,B,DIM)
%   
%   Example:
%   spmd
%       N = 1000;
%       d1 = codistributed.colon(1,N);
%       d2 = codistributed.ones(N,1);
%       d = dot(d1,d2)
%   end
%   
%   returns d = N*(N+1)/2.
%   
%   See also DOT, CODISTRIBUTED, CODISTRIBUTED/COLON, CODISTRIBUTED/ONES.


%   Copyright 2006-2022 The MathWorks, Inc.

narginchk(2, 3);

if nargin == 2
    distributedutil.CodistParser.verifyNonCodistributedInputs({a, b});
elseif nargin == 3
    distributedutil.CodistParser.verifyNonCodistributedInputs({a, b, dim});
    dim = distributedutil.CodistParser.gatherIfCodistributed(dim);
    if ~isa(a, 'codistributed') && ~isa(b, 'codistributed')
        c = dot(a, b, dim);
        return;
    end
end

% Special case: A and B are vectors and dim not supplied
if isvector(a) && isvector(b) && nargin < 3
    needRedistr = iIsRedistributionNeeded(a, b);
    if ~needRedistr
        % If no redistribution is needed then call spmdReduce directly and return
        localC = spmdReduce(@(x,y) sum([x y]), dot(getLocalPart(a),getLocalPart(b)));
        % Set the local part empty, with correct class, on all the workers
        % but 1
        if spmdIndex ~= 1
            localC = zeros(1, 0, 'like', localC);
        end
        codistrC = codistributor1d(2, [1, zeros(1, spmdSize - 1)], [1, 1]);
        c = codistributed.pDoBuildFromLocalPart(localC, codistrC);
    else
        % The binary operation .* does a redistribution if necessary. Note that
        % the result is scalar but still distributed.
        isColVec = @(v) size(v, 1) > 1;
        if isColVec(a) == isColVec(b)
            c = sum(conj(a).*b);
        else
            % Ensure that a and b are either both column vectors or both row vectors.
            c = sum(conj(a).*b.');
        end
        
    end
    return;
end

% Check dimensions
if ~isequal(size(a), size(b))
    error(message('parallel:distributed:DotInputSizeMismatch'));
end

if nargin == 2
    c = sum(conj(a).*b);
else
    c = sum(conj(a).*b, dim);
end

end

function needRedistr = iIsRedistributionNeeded(a, b)

% Basics assumption is that we need to redistribute
needRedistr = true;

if isa(a, 'codistributed') && isa(b, 'codistributed')
    codistrA = getCodistributor(a);
    codistrB = getCodistributor(b);
    % If a and b are distributed using the same scheme
    if strcmp(class(codistrA), class(codistrB))
        if isa(codistrA, 'codistributor1d')
            % If the partitions and the local sizes are exactly the same we do not need to redistribute
            if isequal(codistrA.Partition, codistrB.Partition) && ...
                    isequal(size(getLocalPart(a)), size(getLocalPart(b)))
                needRedistr = false;
            end
        elseif isa(codistrA, 'codistributor2dbc')
            % We assume we always need to redistribute when using 2dbc distribution
            % scheme
        end
    end
end
end
