function [B, rCond] = scalaTridiagSolve(transposeFlagA, A, B)
% SCALATRIDIAGSOLVE Sparse Banded LU solver using MKL ScaLAPACK routines
% Solves sparse linear system op(A)X=B, where op() is defined by
% transposeFlagA which is one of
% TransposeFlags.NO_TRANSPOSE/TRANSPOSE/CONJUGATE_TRANSPOSE.
% Right-hand side B is overwritten by solution X.

%   Copyright 2020-2022 The MathWorks, Inc.

import distributedutil.TransposeFlags

% This function handles both sparse and full A
assert(~isempty(A), "Empty case should be handled in pMldivideAndMrdivide.");

sizeB = size(B);

% Unfortunately we aren't able to calculate the reciprocal condition number
% for a tridiagonal matrix, so we will simply assume that the problem is
% well-conditioned and return rCond=[] except for 2 cases: 
% * If the input has any NaN, we set rCond to NaN and return
% codistributed(B).*NaN without calling the solver.
% * If the result of the solver has any NaN, we set rCond to NaN.
rCond = [];
if hasNaN(A)
    B = codistributed(B).*NaN;
    rCond = NaN;
    return;
end

% Output distribution of X should match distribution of A
distA = getCodistributor(A);
distX = codistributor1d(1, distA.Partition, sizeB);

% scaLAPACK needs bands of A 1D-distributed by columns in full format with
% the superdiagonal left justified and subdiagonal right justified.
if transposeFlagA == TransposeFlags.NO_TRANSPOSE
    bandA = spdiags(A.', -1:1); % m x 3 dense matrix distributed 1D by rows
else
    bandA = spdiags(A, -1:1);  % m x 3 dense matrix distributed 1D by rows
end

distBandA = getCodistributor(bandA);
assert(distBandA.Dimension == 1, "band matrix returned from spdiags in scalaTridiagSolve should be distributed along rows.") 
assert(isequal(distBandA.Cached.GlobalSize, [sizeB(1), 3]), "band matrix returned from spdiags in scalaTridiagSolve has wrong size.")

% Scalapack requires that partition be uniform except for last block
partA = distBandA.Partition;

if length(partA) > 1 && (any(diff(partA(1:end-1))) || partA(end) > partA(1)...
        || partA(1) <= 2)
    % Redistribution needed
    szA = size(bandA);
    newBlockSize = max(ceil(szA(1)/spmdSize), 3);
    newPartA = diff(min(szA(1), (0:spmdSize)*newBlockSize));
    newDistBandA = codistributor1d(1, newPartA, szA);
    bandA = redistribute(bandA, newDistBandA);
    distBandA = newDistBandA;
end

% New codistributor for B by rows with matching partition as A
distB = codistributor1d(1, distBandA.Partition, sizeB);

% Check distribution scheme of B and redistribute if necessary
B = distributedutil.redistributeIfNeeded(B, distB);

localA = hGetLocalPartWithCorrectComplexity(bandA);
localB = hGetLocalPartWithCorrectComplexity(B);

% if localA is empty, it can be sparse
localA = full(localA);

isSparseB = issparse(localB);
% localB needs to be full for use within scaLAPACK
localB = full(localB);

% Conjugate localA if required.
if transposeFlagA == TransposeFlags.CONJUGATE_TRANSPOSE
    localA = conj(localA);  % Conjugate the diagonals of A
end

nGlobal = distB.Cached.GlobalSize(1);

% Define codistributor of one row of transposed bandA though we don't
% actually transpose the vectors.
distBandA = codistributor1d(2, distBandA.Partition, [1, nGlobal]);

descA = bandarraydescriptor(distBandA);
descB = bandarraydescriptor(distB);

% localA contains the three diagonals of A as columns:
% 3 is the subdiagonal
% 2 is the main diagonal
% 1 is the superdiagonal

% Casts are needed to preserve complexity in case local part is empty.
superDiag = cast(localA(:, 1), 'like', localA);
mainDiag = cast(localA(:, 2), 'like', localA);
subDiag = cast(localA(:, 3), 'like', localA);

% Pad localB on last worker if it has fewer rows than on the other workers
rowsLocalB = size(localB, 1);
rowsPadB = distB.Partition(1) - rowsLocalB;
if rowsPadB > 0
    localB = [localB; zeros(rowsPadB, size(localB, 2), 'like', localB)];
end

localB = parallel.internal.codistextern.sparseMldivideTridiag( ...
    subDiag, mainDiag, superDiag, descA, localB, descB);

% Remove padding if it had been applied
if rowsPadB > 0
    localB = localB(1:rowsLocalB, :);
end

if isSparseB
    % The sparsity of the result should be the same as the sparsity of
    % the right hand side.
    localB = sparse(localB);
end

B = codistributed.pDoBuildFromLocalPart(localB, distB);

B = distributedutil.redistributeIfNeeded(B, distX);

% Check if result has any NaN and if so, set rCond to NaN.
if hasNaN(B)
    rCond = NaN;
end
end
