function R = rescale(A, varargin)
%RESCALE   Rescales the range of data
%   R = RESCALE(A)
%   RESCALE(A,b,c)
%   RESCALE(...,'InputMin',IMIN)
%   RESCALE(...,'InputMax',IMAX)
%
%   If the first input A is unlabeled, all additional inputs must also be
%   unlabeled.
%
%   If the first input A is labeled, all additional inputs must either be
%   unlabeled scalars, or have dimension labels that are a subset of the
%   dimension labels of A. In this case, each dimension must either be
%   singleton or match the corresponding dimension of A.

%   Copyright 2019 The MathWorks, Inc.

narginchk(1,inf);

% Process inputs
[A, a, b, inputMin, inputMax] = preprocessInputs(A, varargin{:});

% Quick return for empty inputs
if isempty(A)
    R = A;
    return;
end

% Regularize constant values to return lowerbound of output range
constReg = (inputMin == inputMax);

% Determine where to center the problem based on the input range 
sigma = max(min(0,inputMax),inputMin);
inputMin = inputMin - sigma;
inputMax = inputMax - sigma;

% Let dlarray do its thing with the same algorithm as is used in rescale
% This will have dimension labels, internal permutations, automatic
% differentiation work correctly, as they have already been checked and
% adjusted if necessary.

% Scale to prevent overflow/underflow. All variables here are not dlarrays,
% no matter what a, b, inputMax, inputMin are.
e1 = nextpow2(max(abs(inputMax), abs(inputMin)));
r1 = 2.^(e1-1);
e2 = nextpow2(max(abs(a), abs(b)));
r2 = 2.^(e2-1);
r3 = 2.^(fix((e1+e2)/2)-1);

inputMax_r1 = inputMax ./ r1;
inputMin_r1 = inputMin ./ r1;
constReg_r1 = constReg ./ r1;

z = (inputMax_r1.*(a./r3) - inputMin_r1.*(b./r3) + (a./r3).*constReg_r1) ...
    ./ (inputMax_r1 - inputMin_r1 + constReg_r1);
slope = (b./r2 - a./r2)./(inputMax./r3 - inputMin./r3 + constReg./r3);
R = r2 .* (slope./r3 .* (A - sigma) + (r3./r2).*z);

% Check to make sure the output is within the output range
R = max(R, a, 'includenan');
R = min(R, b, 'includenan');

end
%--------------------------------------------------------------------------
function [A, a, b, inputMin, inputMax] = preprocessInputs(A, varargin)
% Parse RESCALE inputs

% Process A
if ~(isnumeric(A) || islogical(A)) || ~isreal(A)
    error(message('MATLAB:rescale:InvalidA'));
end

hasLabelsA = isa(A, 'dlarray') && hasFormat(A);

% Set defaults
a = 0;
b = 1;
minFlag = false;
maxFlag = false;


ndimsA = ndims(A);
sizA = size(A);
if hasLabelsA
    dimsA = dims(A);
    dimCountA = iCountDims(dimsA, size(A));
else
    dimsA = '';
    dimCountA = [];
end

if ~isempty(varargin)
    indStart = 1;
    
    % Parse output range
    if isnumeric(varargin{1}) || islogical(varargin{1})
        if 2 > length(varargin)
            error(message('MATLAB:rescale:RequiredThirdInput'));
        end
        a = varargin{1};
        b = varargin{2};
        if ~isnumeric(a) || ~isreal(a)
            error(message('MATLAB:rescale:InvalidOutputRange'));
        end
        if ~isnumeric(b) || ~isreal(b)
            if rem(length(varargin),2) ~= 0
                error(message('MATLAB:rescale:RequiredThirdInput'));
            else
                error(message('MATLAB:rescale:InvalidOutputRange'));
            end
        end
        
        % Check sizes and dimension labels of a, b. Return them with
        % the same labels and compatible dimensions, or error.
        a = iProcessInput(a, A, ndimsA, sizA, dimsA, dimCountA, 'MATLAB:rescale:NumberDimsOut');
        b = iProcessInput(b, A, ndimsA, sizA, dimsA, dimCountA, 'MATLAB:rescale:NumberDimsOut');
        
        if any(a > b, 'all')
            error(message('MATLAB:rescale:OutMinGreaterThanOutMax'));
        end
        if ~isfloat(a)
            a = double(a);
        end
        if ~isfloat(b)
            b = double(b);
        end
        indStart = 3;
    end
    
    % Parse name-value pairs
    nvNames = ["InputMin", "InputMax"];
    opts = struct;
    for j = indStart:2:length(varargin)
        name = varargin{j};
        if j+1 > length(varargin)
            error(message('MATLAB:rescale:KeyWithoutValue'));
        elseif (~(ischar(name) && isrow(name)) && ~(isstring(name) && isscalar(name))) ...
                || (isstring(name) && strlength(name) == 0)
            error(message('MATLAB:rescale:ParseFlags'));
        end
        ind = startsWith(nvNames, name, 'IgnoreCase', true);
        if nnz(ind) ~= 1
            error(message('MATLAB:rescale:ParseFlags'));
        end
        opts.(nvNames{ind}) = varargin{j+1};
    end
    if isfield(opts, 'InputMin')
        inputMin = opts.InputMin;
        if ~(isnumeric(inputMin) || islogical(inputMin)) || ~isreal(inputMin)
            error(message('MATLAB:rescale:InvalidInputMin'));
        end
        inputMin = iProcessInput(inputMin, A, ndimsA, sizA, dimsA, dimCountA, 'MATLAB:rescale:NumberDimsInLower');
        minFlag = true;
    end
    if isfield(opts, 'InputMax')
        inputMax = opts.InputMax;
        if ~(isnumeric(inputMax) || islogical(inputMax)) || ~isreal(inputMax)
            error(message('MATLAB:rescale:InvalidInputMax'));
        end
        inputMax = iProcessInput(inputMax, A, ndimsA, sizA, dimsA, dimCountA, 'MATLAB:rescale:NumberDimsInUpper');
        maxFlag = true;
    end
end

% Set input range if not set above
if ~minFlag
    inputMin = min(A, [], 'all');
end
if ~maxFlag
    inputMax = max(A, [], 'all');
end

% Cast inputMin/inputMax to double if not double or single
if ~isfloat(inputMin)
    inputMin = double(inputMin);
end
if ~isfloat(inputMax)
    inputMax = double(inputMax);
end

% Check to make sure min < max
if any(inputMin > inputMax, 'all')
    error(message('MATLAB:rescale:MinGreaterThanMax'));
end

% Preprocess input range if needed
if minFlag
    A = max(A, inputMin, 'includenan');
end
if maxFlag
    A = min(A, inputMax, 'includenan');
end
if ~isfloat(A)
    A = double(A);
end

end


function X = iProcessInput(X, A, ndimsA, sizA, dimsA, dimCountA, errIDNumberDimsOut)
% X represents one of the optional inputs a, b, inputMin, inputMax, which
% must be matched with A using implicit expansion.
hasLabelsA = ~isempty(dimsA);
hasLabelsX = isa(X, 'dlarray') && hasFormat(X);

if hasLabelsA
    if ~hasLabelsX
        if isscalar(X)
            % Unlabeled scalars are fine to combine with labeled A
            return
        else
            error(message('deep:dlarray:RescaleLabeledFirstInput'));
        end
    end
    
    dimCountX = iCountDims(dims(X), size(X));
    
    if any(dimCountX > dimCountA)
        error(message('deep:dlarray:RescaleLabelsSubset'));
    end
    
    % For each dimension of A, set it to true if this dimension is
    % represented by a dimension of x.
    isInA = false(1, length(dimsA));
    cumsumDimCountA = [0 cumsum(dimCountA)];
    for ii=1:length(dimCountA)
        isInA(cumsumDimCountA(ii)+(1:dimCountX(ii))) = true;
    end
    
    if nnz(isInA) > 0 && ~all(size(X, 1:nnz(isInA)) == size(A, find(isInA)) | size(X, 1:nnz(isInA)) == 1)
        error(message(errIDNumberDimsOut));
    end
    
    % Relabel x to have the same labels as A
    X = dlarray(X, [dims(X) dimsA(~isInA)]);
else
    if hasLabelsX
        error(message('deep:dlarray:RescaleUnlabeledFirstInput'));
    end
    
    if (ndims(X) > ndimsA) || ~all(size(X) == sizA(1:ndims(X)) | size(X) == 1)
        error(message(errIDNumberDimsOut));
    end
end
end


function fmtCounts = iCountDims(fmt, sz)
% Count number of occurences of each label in dims:
% dimCounts = [nnz(dims=='S') nnz(dims=='C') nnz(dims=='B') nnz(dims=='T') nnz(dims=='U')]

dimTypes = dlarray.OrderOfDims;

fmtCounts = zeros(size(dimTypes));
if ~isempty(fmt)
    for n = 1:numel(dimTypes)
        thisLetter = dimTypes(n);
        isType = thisLetter==fmt;
        fmtCounts(n) = sum(isType);
    end
    
    if length(fmt) == 2 && fmt(2) == 'U' && sz(2) == 1
        if fmt(1) == 'U' && sz(1) == 1
            fmtCounts(end) = 0;
        else
            fmtCounts(end) = fmtCounts(end) - 1;
        end
    end
end
end
