function varargout = eig(varargin)
%EIG Eigenvalues and eigenvectors of codistributed array
%   D = EIG(A)
%   [V,D] = EIG(A)
%   [V,D,W] = EIG(A)
%   [...] = EIG(A,'balance')
%   [...] = EIG(A,'nobalance')
%   D = EIG(A,B)
%   [V,D] = EIG(A,B)
%   [...] = EIG(A,B,'chol')
%   [...] = EIG(...,'vector')
%   [...] = EIG(...,'matrix')
%   
%   Limitations:
%   
%   If you specify A and B, both must be Hermitian matrices and
%   B must be positive definite.
%   
%   See also EIG, CODISTRIBUTED, CODISTRIBUTED/RAND.


%   Copyright 2006-2022 The MathWorks, Inc.

narginchk(1, 5);
nargoutchk(0, 3);

distributedutil.CodistParser.verifyNonCodistributedInputs(varargin);

[numericArgs, flags] = iParseArgs(varargin);
% Check flags and number of numeric inputs
if ~isempty(flags)
    C = cell(1, numel(numericArgs));
    eig(C{:}, flags{:});
end
% Check numeric inputs are non-sparse float matrices
if any(cellfun(@iIsUnsupported, numericArgs))
    error(message('parallel:distributed:EigNotReal'));
end

% No numeric input is codistributed, hence, call core MATLAB
if ~any(cellfun(@iscodistributed, numericArgs))
    [varargout{1:max(1, nargout)}] = eig(numericArgs{:}, flags{:});
    return
end

% Extract the data inputs
A = numericArgs{1};
generalizedSystem = numel(numericArgs)==2;
if generalizedSystem
    B = numericArgs{2};
    % Check that A & B are the same size
    if ~isequal(size(A), size(B))
        error(message('MATLAB:dimagree'));
    end
end

% Extract flags
[isvec, balance] = iCheckFlags(nargout, flags);

% Special cases not handled by ScaLAPACK
[type, codistr] = iGetOutputTypeAndCodistr(numericArgs);
if isempty(A)
    % Empty input
    [varargout{1:max(1, nargout)}] = iEigEmpty(isvec, type, codistr);
    return;
elseif hasNaN(A) || hasInf(A)
    % Input contains Infs or NaNs
    [varargout{1:max(1, nargout)}] = eigSvdInfOrNaN(isvec, type, codistr);
    return
end

isHermitianA = ishermitian(A);

if generalizedSystem
    [varargout{1:max(1, nargout)}] = iGeneralizedEig(A, B, isHermitianA, isvec);
else
    [varargout{1:max(1, nargout)}] = iStandardEig(A, isHermitianA, isvec, balance);
end
end

function varargout = iStandardEig(A, isHermitianA, isvec, balance)
% For standard systems we support non-symmetric and symmetric case. Notice
% that the later one is always balanced, hence, the balance flag is
% ignored.
if isHermitianA
    [varargout{1:nargout}] = scalaEig(A, isvec);
elseif ishermitian(A, 'skew') && ~isreal(A)
    % If A is complex skew-Hermitian, multitply by 1i, pass it to the
    % Hermitian solver and scale eigenvalues back.
    [varargout{1:nargout}] = scalaEig(A.*1i, isvec);
    % Get output index of the eigenvalues, can be 1st or 2nd output.
    idxValues = ~isvec + 1;
    varargout{idxValues} = varargout{idxValues}.*(-1i);
else
    [varargout{1:nargout}] = scalaNonSymmetricEig(A, isvec, balance);
end
end

function varargout = iGeneralizedEig(A, B, isHermitianA, isvec)
% For generalized systems we only support the symmetric case, which is
% always balanced.

% Special case not handled by ScaLAPACK, B contains NaNs or Infs
if hasNaN(B) || hasInf(B)
    [type, codistr] = iGetOutputTypeAndCodistr({A, B});
    [varargout{1:nargout}] = eigSvdInfOrNaN(isvec, type, codistr);
    return
end

% Non-symmetric, generalized systems are not yet supported
if ~isHermitianA || ~ishermitian(B)
    error(message('parallel:distributed:EigNoGeneralizedNonSymEig'))
end

if nargout == 3
    error(message('parallel:distributed:EigUnsupportedLeftVectors'));
end

try
    R = chol(B, "upper"); % B = R'*R
catch ME
    error(message('parallel:distributed:EigBMustBePosDef'));
end
% A*U == B*U*S <=> inv(R')*A*inv(R) * U == U*S
if isvec
    [varargout{1:nargout}] = eig(R'\A/R, "vector");
else
    [varargout{1:nargout}] = eig(R'\A/R, "matrix");
end
if nargout == 2
    % Post-process eigenvectors to make them B-orthogonal, U'*B*U == eye
    varargout{1} = R\varargout{1};
end
end

function [numericArgs, flags] = iParseArgs(argList)

trailingArgIdx = find(cellfun(@matlab.internal.datatypes.isScalarText, argList), 1, 'first');

if isempty(trailingArgIdx)
    % All data
    numericArgs = argList;
    flags = {};
elseif trailingArgIdx == 1
    % All flags
    numericArgs = {};
    flags = argList;
else
    numericArgs = argList(1:trailingArgIdx-1);
    flags = argList(trailingArgIdx:end);
end

% No data inputs is an error
if isempty(numericArgs)
    ME = MException(message('parallel:distributed:EigNotReal'));
    throwAsCaller(ME);
end

% More than two data inputs is an error
if numel(numericArgs)>2
    ME = MException(message('MATLAB:eig:unknownArgumentGeneralized'));
    throwAsCaller(ME);
end

% Gather all flags and convert to char
flags = cellfun(@(x)char(gather(x)), flags, 'UniformOutput', false);
end

function [isvec, balance] = iCheckFlags(numOut, flags)

% Default for eigenvalues is vector for less than 2 outputs, matrix
% othwerwise.
isvec = (numOut < 2);
% Default for balance is true
balance = true;

% Find if any of {'vector', 'matrix'} is specified
matvecFlag = ismember(lower(flags), {'vector', 'matrix'});
if any(matvecFlag)
    isvec = strcmpi(flags{matvecFlag}, 'vector');
end

% Find if any of {'balance', 'nobalance'} is specified
balanceFlag = ismember(lower(flags), {'balance', 'nobalance'});
if any(balanceFlag)
    balance = strcmpi(flags{balanceFlag}, 'balance');
end

% Find if 'qz' is defined and error since we don't support this
if any(ismember(lower(flags), 'qz'))
    ME = MException(message('parallel:distributed:EigUnsupportedQz'));
    throwAsCaller(ME);
end
% Notice that the 'chol' flag is ignored since we can only solve
% generalized problems with spd B, which always uses chol.
end

function tf = iIsUnsupported(A)
tf = ~isfloat(A) || ~ismatrix(A) || issparse(A);
end

function varargout = iEigEmpty(isvec, type, codistr)
% we got here because all of numericArgs are empty

if nargout==0
    return
end

if isvec
    codistr = codistr.hGetNewForSize([0 1]);
    eigenVals = zeros(0, 1, type, codistr);
else
    eigenVals = zeros(0, 0, type, codistr);
end

if nargout == 1
    % Just the eigenvalues
    varargout{1} = eigenVals;
elseif nargout >= 2
    % Left Eigenvectors, Eigenvalues
    codistr = codistr.hGetNewForSize([0 0]);
    varargout{1} = zeros(0, 0, type, codistr);
    varargout{2} = eigenVals;
    if nargout == 3
        % Right Eigenvectors
        varargout{3} = zeros(0, 0, type, codistr);
    end
end
end

function [typeArg, codistr] = iGetOutputTypeAndCodistr(numericArgs)
% Get type and codistributor of outputs.

isArgCodistributed = cellfun(@(x) isa(x, 'codistributed'), numericArgs);
firstCodistributedArg = find(isArgCodistributed, 1, 'first');
firstCodistributedMatrix = numericArgs{firstCodistributedArg};
codistr = getCodistributor(firstCodistributedMatrix);
if any(cellfun(@(x) isUnderlyingType(x, 'single'), numericArgs))
    typeArg = 'single';
else
    typeArg = 'double';
end
end
