function varargout = linsolve(A, B, varargin)
%LINSOLVE Solve linear system A*X=B for codistributed arrays
%   
%   X = linsolve(A,B)
%   X = linsolve(A,B,OPTS)
%   [X,R] = linsolve(...)
%   
%   See also CODISTRIBUTED, MLDIVIDE, DECOMPOSITION.


%   Copyright 2021-2022 The MathWorks, Inc.

narginchk(2, 3);
nargoutchk(1, 2);

if issparse(A) || issparse(B)
    error(message("MATLAB:linsolve:sparseNotSupported"));
end

if iscodistributed(A)
    [varargout{1:max(1, nargout)}] = iLinsolveACodistributed(A, B, varargin{:});
else % ~iscodistributed(A)
    [varargout{1:max(1, nargout)}] = iLinsolveAReplicated(A, B, varargin{:});
end

end

function varargout = iLinsolveAReplicated(A, B, varargin)
% LINSOLVE when A is not codistributed but B might be.

% If B is not codistributed, call in-memory LINSOLVE.
if ~iscodistributed(B)
    [varargout{1:max(1, nargout)}] = linsolve(A, B, varargin{:});
    return
end

% B is codistributed but A is not. We redistribute B by columns such that
% all the labs have different linear systems to solve.
codistB = getCodistributor(B);
if ~( isa(codistB, "codistributor1d") && (codistB.Dimension == 2) )
    B = redistribute(B, codistributor1d(2));
    codistB = getCodistributor(B);
end

% Then, we call LINSOLVE on each of the column blocks in parallel.
[varargout{1:max(1, nargout)}] = linsolve(A, getLocalPart(B), varargin{:});
% Finally, we build the codistributed solution X from its local parts. X is
% distributed by columns as B is, but it may have a different number of
% rows.
localX = varargout{1};
codistX = codistributor1d(2, codistB.Partition, [size(localX, 1) size(B, 2)]);
varargout{1} = codistributed.pDoBuildFromLocalPart(localX, codistX);
% 2nd output is a scalar representing rcond, for square, or rank otherwise.
if nargout == 2
    varargout{2} = codistributed.pDoBuildFromReplicatedScalar(varargout{2}, 2);
end
end

function varargout = iLinsolveACodistributed(A, B, opts)
% LINSOLVE when A is codistributed.

% Default is no conjugate transpose.
needCtranspose = false;
% If A is not square we must use QR, otherwise default is LU ; and
% "upper"/"lower" is not used.
[m, n] = size(A);
isARectangular = m ~= n;
% Default solver is "lu" for square and "qr" for rectangular.
decompTypeAndUploIfNeeded = "lu";
if isARectangular
    decompTypeAndUploIfNeeded = "qr";
end

if nargin > 2
    % Gather opts if needed.
    if iscodistributed(opts)
        opts = gather(opts);
    end
    % First round of parsing to filter out inconsistent opts.
    try
        iParseInconsistentOpts(opts, isARectangular);
    catch ME
        throwAsCaller(ME);
    end
    % Find out which decomposition to use using opts provided, but if A is
    % not square we must use QR.
    [decompTypeAndUploIfNeeded, needCtranspose] = iGetDecompositionTypeFromOpts(opts);
end

% If A is rectangular and the triangular solve is picked, we pre-process A
% and B to make use of a triangular solver.
useRectangularTriangularSolver = nargin > 2 && isARectangular ...
    && decompTypeAndUploIfNeeded(1) == "triangular";
if useRectangularTriangularSolver
    % Extract the square diagonal part of A, i.e., A = A(1:min(size(A), 1:min(size(A)))
    k = min(m, n);
    A = subsref(A, substruct('()',{1:k, 1:k}));
    % Shorten B if needed
    B = subsref(B, substruct('()',{1:k, ':'}));
end

% Decompose and solve.
decA = decomposition(A, decompTypeAndUploIfNeeded{:});
if needCtranspose
    varargout{1} = decA'\B;
else
    varargout{1} = decA\B;
end

% If the triangular solver has been used for a rectangular matrix we might
% need to expand B with zero rows to make its dimension consistent, its number
% of rows is the number of columns of A, i.e., X(n, end) = 0.
expansionIsNeeded = (m < n && ~needCtranspose) || (m > n && needCtranspose);
if useRectangularTriangularSolver && expansionIsNeeded    
    nExpansion = n;
    if needCtranspose
        % If ctranspose make sure to expand using the right dimension.
        nExpansion = m;
    end
    S = substruct('()',{nExpansion, size(varargout{1}, 2)});
    varargout{1} = subsasgn(varargout{1}, S, 0);
end

% 2nd output is rank for QR, and rcond otherwise.
if nargout == 2
    if decompTypeAndUploIfNeeded(1) ~= "qr"
        varargout{2} = rcond(decA);
        % rcond is single if the solution is
        if isUnderlyingType(varargout{1}, "single")
            varargout{2} = single(varargout{2});
        end
    else
        varargout{2} = rank(decA);
    end
end
end

function [decompTypeAndUploIfNeeded, needTranspose] = iGetDecompositionTypeFromOpts(opts)
% Parses opts to extract the decomposition we should use to solve the
% system. For triangular we should also specify "lower"/"upper" so we
% return a cell array containing both if needed, or only the decomposition
% type. Note we assume opts is consistent and there is no error checking
% here. Default values are different if A is rectangular or not.

% Initialise all the possible fields to false. Note we do not have any
% particular solver for upper Hessenberg so we don't use opts.UHESS. If it
% is specified then we use 'lu'.
isLowerTriangular  = false; % LT
isUpperTriangular  = false; % UT
isSymmetric        = false; % SYM - real symm or complex Hermitian
isPositiveDefinite = false; % POSDEF
isRectangular      = false; % RECT - force QR factorisation
needTranspose      = false; % TRANSA - conjugate transpose, or not

% Go through all the fields and retrieve the values.
if isfield(opts, "LT")
    isLowerTriangular = opts.LT;
end
if isfield(opts, "UT")
    isUpperTriangular = opts.UT;
end
if isfield(opts, "SYM")
    isSymmetric = opts.SYM;
end
if isfield(opts, "POSDEF")
    isPositiveDefinite = opts.POSDEF;
end
if isfield(opts, "RECT")
    isRectangular = opts.RECT;
end
if isfield(opts, "TRANSA")
    needTranspose = opts.TRANSA;
end

% From this information we pick the best solver we can. We use the criteria
% in the following order: rectangular, triangular, symmetry. Indeed,
% triangular can be used for rectangular matrices as well.
if isRectangular
    % Rectangular means 'qr', even if A is in fact square.
    decompTypeAndUploIfNeeded = "qr";
elseif isLowerTriangular
    decompTypeAndUploIfNeeded = ["triangular", "lower"];
elseif isUpperTriangular
    decompTypeAndUploIfNeeded = ["triangular", "upper"];
elseif isSymmetric
    if isPositiveDefinite
        % If A positive definite too, we can use 'chol'.
        decompTypeAndUploIfNeeded = "chol";
    else
        % Otherwise, we use 'lu' as 'ldl' is not supported for dense.
        decompTypeAndUploIfNeeded = "lu";
    end
else
    % Use "lu" because A is set to be non-rectangular.
    decompTypeAndUploIfNeeded = "lu";
end
end

function iParseInconsistentOpts(opts, isARectangular)
% Call in-memory LINSOLVE with opts and a template for A to get free
% parsing.

if isARectangular
    % Pass on a rectangular template for A if A is rectangular,
    % making sure to use consistent dimensions if TRANSA is used.
    needCtranspose = false;
    if isfield(opts, "TRANSA")
        needCtranspose = opts.TRANSA;
    end
    if needCtranspose
        linsolve(ones(1, 2), ones(2, 1), opts);
    else
        linsolve(ones(2, 1), ones(2, 1), opts);
    end
else
    linsolve(1, 1, opts);
end
end
