function s = norm(A,p)
%NORM Vector and matrix norms for codistributed array
%   
%   Supported matrix X syntaxes:
%   N = NORM(X)
%   N = NORM(X,2)
%   N = NORM(X,1)
%   N = NORM(X,Inf)
%   N = NORM(X,'fro')
%   N = NORM(X,P)
%   
%   Supported vector X syntaxes:
%   C = NORM(X)
%   C = NORM(X,P)
%   C = NORM(X,Inf)
%   C = NORM(X,-Inf) 
%   N = NORM(X,'fro')
%   
%   Supported N-D array X syntax:
%   C = NORM(X,'fro')
%   
%   See also NORM, CODISTRIBUTED, CODISTRIBUTED/NORMEST, CODISTRIBUTED/VECNORM.
%   


%   Copyright 2006-2021 The MathWorks, Inc.
    if nargin == 2
        distributedutil.CodistParser.verifyNonCodistributedInputs({A, p});
    end
    
    if nargin < 2
        p = 2;
    else 
        p = distributedutil.CodistParser.gatherIfCodistributed(p);
        if ~isa(A, 'codistributed')
            % Only the p was codistributed.
            s = norm(A, p);
            return;
        end 
    end 
    
    if ~isfloat(A)
        error(message('MATLAB:norm:inputType'));
    end  
    
    isVecA = isvector(A);
    
    % Check if p is a supported norm for input A, otherwise throw
    p = iIsValidNorm(isVecA, p);
    
    % Check for special cases that can exit early
    if hasNaN(A) || hasNaN(p) % Check for NaNs first!
        % NORM returns full NaN of correct class
        s = codistributed.NaN(1, underlyingType(A));
        return;
    elseif isempty(A) || strcmp(getFillPattern(A),'ZeroMatrix')
        % Ensure norm of an empty or zero matrix is full and the same type as A
        s = full(zeros(1, 1, 'like', A));
        return;
    elseif isscalar(A)
        % Ensure norm of a scalar is full independent of A
        s = full(abs(A));
        % Return early
        return;
    elseif hasInf(A)
        % NORM returns full Inf of correct class
        s = codistributed.Inf(1, underlyingType(A));
        return
    end
    
    % Check if Frobenius norm is selected
    isFrobenius = false;
    if strcmp(p, 'fro')
        isFrobenius = true;
    elseif isVecA && isnumeric(p)
        % For vectors 2-norm and Frobenius norm is the same
        isFrobenius = p==2;
    end
    
    % N-D arrays only support Frobenius norm
    if ~ismatrix(A) && ~isFrobenius
        error(message('MATLAB:norm:inputMustBe2D'));
    end

    % Switch between different methods
    if isFrobenius
        % Frobenius norm computation is the same for vectors, matrices,
        % and N-D arrays
        s = iFrobeniusNorm(A);
    elseif isVecA
        s = iVectorNorm(A, p);
    else
        s = iMatrixNorm(A, p);
    end
    
    % Ensure output is full
    s = full(s);
end % End norm

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function s = iFrobeniusNorm(A)
        LP = getLocalPart(A);
        codistr = getCodistributor(A);
        s = codistr.hFrobeniusNormImpl(LP);
        % Create codistributed output from replicated scalar local parts
        s = codistributed.pDoBuildFromReplicatedScalar(s, 2);
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function s = iVectorNorm(A, p)
% Compute non-trivial vector norms only for supported p.
    switch p
      case {inf, 'inf'} 
        s = max(abs(A));
      case -inf
        s = min(abs(A));
      otherwise 
          s = nthroot(sum(abs(A).^p), p);
    end
end 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function s = iMatrixNorm(A, p)
% Compute non-trivial matrix norms only for supported p.
    switch p
      case 1 
        s = max(sum(abs(A),1));
      case 2
        if issparse(A)
            ME = MException(message( ...
            'parallel:distributed:NormSparseMat2NormUnsupported'));
            throwAsCaller(ME)
        else
            s = max( svd(A) );
        end 
      case {inf, 'inf'}
        s = max(sum(abs(A),2));
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function p = iIsValidNorm(isVecA, p)
% Check if p triggers a supported case

    if (isscalar(p) && isnumeric(p)) || isScalarString(p)
        % Check if p for vectors A is supported
        if isVecA
            if isScalarString(p)
                [tf, p] = iCleanupP(p, {'fro', 'inf'});
            else 
                % Vectors accept all numeric p
                tf = true;
                % Ensure p is floating point
                if ~isfloat(p)
                    p = double(p);
                end
            end

            % Error if p is not supported for vectors throw
            if ~tf
                ME=MException(message(...
                    'parallel:distributed:NormVectorNormNotSupport'));
                throwAsCaller(ME)
            end
        else
        % Check if p for vectors A is supported
            if isScalarString(p)
                [tf, p] = iCleanupP(p, {'fro', 'inf'});
            else
                tf = any(p==[1 2 inf]);
            end

            % Error if p is not supported for vectors throw
            if ~tf
                ME = MException(message(...
                    'MATLAB:norm:unknownNorm'));
                throwAsCaller(ME)
            end
        end
    else
        ME = MException(message('parallel:distributed:NormInvalidP'));
        throwAsCaller(ME)
    end
end

function [tf, p] = iCleanupP(p, validSyntaxes)
% Check if p matches (even if only partially) any of the valid syntaxes and
% clean-up p in order to be recognised later on. Ensure strings of length 0
% are not accepted.
cmp = strncmpi(p, validSyntaxes, strlength(p)) & strlength(p);
tf = any(cmp);
if tf
    p = validSyntaxes{cmp};
end
end
