function Z = mpower(X, Y)
%^   Matrix power of codistributed array
%   C = A^b
%   C = MPOWER(A,b)
%   Either A or b (but not both) may be a matrix.  The other input must be scalar.
%   
%   Example:
%   spmd
%       A = rand(1000, 'codistributed');
%       C = A^1.3;
%   end
%    
%   See also MPOWER, CODISTRIBUTED.


%   Copyright 2012-2022 The MathWorks, Inc.

narginchk(2, 2);

distributedutil.CodistParser.verifyNonCodistributedInputs({X, Y});

% Do the input checks in the same order as MATLAB.
if isscalar(X) && isscalar(Y)
    % Redirect to element-wise.
    Z = power(X, Y);
    return
end

if isinteger(X) || isinteger(Y)
    error(message('MATLAB:mpower:integerNotSupported'));
end

if ~ismatrix(X) || ~ismatrix(Y)
    error(message('MATLAB:mpower:inputsMustBe2D'));
end

if ~isfloat(X)
    X = double(X);
end
if ~isfloat(Y)
    Y = double(Y);
end
if superiorfloat(X, Y) == "single"
    if issparse(X) || issparse(Y)
        error(message('MATLAB:mpower:sparseSingleNotSupported'))
    end
    X = single(X);
    Y = single(Y);
end

if isscalar(Y) && isreal(Y)
    if isfinite(Y) && round(Y) == Y && size(X, 1) == size(X, 2)
        Z = integerMpower(X, Y);
        return
    end
end
Z = generalMpower(X, Y);
end

function Z = integerMpower(X, y)
% Calculate an integer matrix power (or its inverse) by repeated
% multiplication.
outType = parallel.internal.array.dominantType(X,y);
y = gather(y); % Scalar so save time by replicating.
if hasNaN(X)
    if y <= 0 || y >= 3
        % NaN to 0 is NaN; NaN to any power >= 3 or <= 0 is NaN
        Z = NaN(size(X), "like", outType);
    elseif y == 1
        Z = cast(X, "like", outType);
    else % y == 2
        assert(y == 2, "error, y should be equal to 2")
        Z = cast(X*X, "like", outType);
    end
elseif y == 0
    % Anything else to power 0 is identity (always real)
    Z = eye(size(X), 'like', outType);
else
    % X and y can be sparse
    % Z = X^y for integer y: use repeated squaring.
    % For example X^13 = X^1 + X^4 + X^8.
    p = abs(y);
    D = X;
    first = true;
    while p>0
        if rem(p, 2) == 1 % if odd
            if first
                Z = D;  % hit first time. D*I
                first = false;
            else
                Z = D*Z;
            end
        end
        p = fix(p/2);
        if p ~= 0
            D = D*D;
        end
    end
    if y < 0
        Z = inv(Z);
    end
    % Since we didn't combine y into the calculation itself, we must
    % explicitly force the output to the right type.
    Z = cast(Z, "like", outType);
end
if issparse(X)
    Z = sparse(Z);
end
end

function Z = generalMpower(X, Y)
if isscalar(Y) && size(X, 1) == size(X, 2)
    scalarX = false;
    Z = generalMpowerImpl(X, Y, scalarX);
elseif isscalar(X) && size(Y, 1) == size(Y, 2)
    scalarX = true;
    Z = generalMpowerImpl(Y, X, scalarX);
else
    Me = MException(message('MATLAB:mpower:notScalarAndSquareMatrix'));
    throwAsCaller(Me);
end
if issparse(X)
    Z = sparse(Z);
end
end

function Z = generalMpowerImpl(X, Y, scalarX)
% Calculate matrix power using eigen-decomposition.
% X and Y should never be sparse.
try
    if ishermitian(X)
        [V, d] = eig(full(X), 'vector');
    else
        [V, d] = eig(full(X), 'nobalance', 'vector');
    end
    if scalarX
        d = Y.^d.';
    else
        d = d.'.^Y;
    end
    if ishermitian(X)
        Z = V.*d*V';
    else
        Z = V.*d/V;
    end
catch err
    rethrow(err);
end
if issparse(X)
    Z = sparse(Z);
end
end
