function obj = parenReference(obj,varargin)
% Implement obj(indices)

%   Copyright 2019-2020 The MathWorks, Inc.

dimsobj = dims(obj);
if ~isempty(dimsobj) && numel(varargin) < ndims(obj)
    % Drop the format if the indexing does not address all of the
    % dimensions (result will be reduced in dimensionality compared to the
    % input).
    dimsobj = '';
end

[fd, objdata] = extractData(obj.FormattedData);
viewToData = obj.ViewToDataMap;

done = false;

% Remove dlarray subscripts
[objdata, varargin] = iRemoveDlarray(objdata,varargin);

if isempty(viewToData)
    % Apply indexing to underlying data directly
    objdata = objdata(varargin{:});
    done = true;
elseif isscalar(varargin)
    % Linear indexing case, translate that index if it's short enough
    ind = varargin{1};
    if isnumeric(ind) && numel(ind) < numel(objdata)/10 && ~iIsRecordingArray(ind)
        % Heuristic: Transforming the index is only faster if the index
        % array is smaller than obj by a significant number.
        try
            [objdata, viewToData] = iLinearNumericIndexRef(objdata, viewToData, obj.DataToViewMap, ind);
            done = true;
        catch
            % In case of error, let iExplicitlyPermuteRef throw the error.
        end
    end
elseif length(varargin) >= length(viewToData)
    % Possible to permute varargin to match the internal permutation
    try
        [objdata, viewToData] = iPermuteSubscriptsRef(objdata, viewToData, varargin);
        done = true;
    catch
        % In case of error, let iExplicitlyPermuteRef throw the error.
    end
end

if ~done
    % No special case applied, or an error was thrown. Fall back on
    % explicitly permuting the input dlarray.
    [objdata, viewToData] = iExplicitlyPermuteRef(objdata, viewToData, varargin);
end

% Adjust dimension labels (add 'U' for trailing dimensions that used to be
% singletons without a dimension label, drop 'U' for trailing dimensions
% that have become singleton). May call squeezedata if dimension labels are
% dropped.
[dimsobj, objdata, viewToData] = addOrDropDimensionLabels(dimsobj, objdata, viewToData);
fd.DimensionMetadata = dimsobj;

% Set data
obj.FormattedData = insertData(fd, objdata);

% Set implicit permutations
if isempty(viewToData)
    obj.ViewToDataMap = [];
    obj.DataToViewMap = [];
else
    dataToView(viewToData) = 1:length(viewToData);
    obj.ViewToDataMap = viewToData;
    obj.DataToViewMap = dataToView;
end
end


function [data, viewToData] = iExplicitlyPermuteRef(data, viewToData, args)
% Permute obj explicitly, then just flow through to indexing data. This
% assumes viewToData is not empty.

data = permute(data, viewToData);
    
% Remove all information related to permutation
viewToData = [];

% Apply indexing to underlying data
data = data(args{:});
end


function [data, viewToData] = iPermuteSubscriptsRef(data, viewToData, args)
% Permute the subcripts to match the underlying data.

% Adjust length of viewToData to be at least as long as args cell array
lenDtoV = length(viewToData);
lenArgs = length(args);
viewToData(lenDtoV+1:lenArgs) = lenDtoV+1:lenArgs;

% Permute arguments to internal order
args(viewToData) = args;

% Index into data
data = data(args{:});
end


function [data, viewToData] = iLinearNumericIndexRef(data, viewToData, dataToView, ind)
% Apply a linear index in obj's view to underlying data by translating the
% index (only efficient if numel(x) is significantly smaller than numel(obj)).

% Size of object seen from the outside (may have additional trailing
% singleton dimensions)
szobj = size(data, viewToData);

% Convert ind to a 1-based double array
if isobject(ind)
    ind = subsindex(ind) + 1;
end
ind = double(ind);

% Check that index is valid. It's necessary to check that ind is all
% integer, because it's possible that iTranslateIndex returns an integer
% translation for a non-integer input. For example,
% iTranslateIndex(2.5, [4 2], [2 1], [2 1]) returns 4.
if any(ind(:) > numel(data)) || any(ind(:) < 1) || any(floor(ind) ~= ind, 'all')
    error(message('MATLAB:badsubscript'));
end

% Translate outside linear index to linear index into underlying data
indData = iTranslateIndex(ind, szobj, viewToData, dataToView);
% Apply subsref to underlying data
data = data(indData);

% If both obj and ind are "ND-vectors", their direction depends on the size
% of the array being indexed into. Need to adjust for internal permutation
% here.
if nnz(szobj ~= 1) == 1 && nnz(size(ind) ~= 1) == 1
    szobj(szobj ~= 1) = numel(data);
    
    data = reshape(data, szobj);
end

viewToData = [];
end


function dataind = iTranslateIndex(ind, szObj, viewToData, dataToView)
% Translate linear indices into the view object into linear indexes into
% the underlying data.

% Size of obj in view coordinates
szObj(end+1:length(viewToData)) = 1;

% The following is a more efficient version of calling ind2sub and then
% sub2ind.
%   c = cell(1, length(szObj));
%   [c{:}] = ind2sub(szObj, ind);
%   c = c(dataToView);
%   dataind = sub2ind(szObj(dataToView), c{:});

% Increments for each dimension in the outside and inside representation:
incView = ones(size(szObj));
incData = ones(size(szObj));
for ii=1:length(incView)-1
    incView(ii+1) = incView(ii) * szObj(ii);
    incData(dataToView(ii+1)) = incData(dataToView(ii)) * szObj(dataToView(ii));
end

% Transform to 0-based indexing
ind = ind - 1;

% Initialize output dataind
dataind = ones(size(ind));


for i = length(szObj):-1:2
    % Remaining index without the part of ind used in dimension i.
    indRemainder = rem(ind, incView(i));
    
    % subI is the (all-integer, 0-based) subscript into dimension i.
    % subI = (ind - indRemainder) / incView(i)
    % Multiply subI with incData(i) to find its contribution to dataind
    % (based on the increment for dimension i in the data representation).
    % Note:
    % - not computing subI explicitly to avoid intermediate memory usage.
    % - using (ind - indRemainder) * (incData(i) / incView(i)) would result
    %   in round-off errors.
    
    % Apply subscript to the outside linear index
    dataind = dataind + ((ind - indRemainder) / incView(i)) * incData(i);
    
    % Upate the index
    ind = indRemainder;
end

% Apply last subscript and move back to 1-based indexing
dataind = dataind + ind*incData(1);
end

function [objdata, subs] = iRemoveDlarray(objdata, subs)
% Convert dlarray subscripts into RecordingArray or CPU/GPU array
% If any subscripts are recording then make objdata recording
isObjRecording = iIsRecordingArray(objdata);
for i=1:length(subs)
    sub = subs{i};
    if iIsDlarray(sub)
        recordingSub = getRecorded(sub);
        if iIsRecordingArray(recordingSub)
            subs{i} = iGetPermutedRecordingArray(sub,recordingSub);
            if ~isObjRecording
                objdata = deep.internal.recording.RecordingArray(objdata);
                isObjRecording = true;
            end
        else
            subs{i} = extractdata(sub);
        end
    end
end
end

function recordingSub = iGetPermutedRecordingArray(sub,recordingSub)
[~,externalPermutation] = getRecordingPermutations(sub);
if ~isempty(externalPermutation)
    recordingSub = permute(recordingSub,externalPermutation);
end
end

function tf = iIsRecordingArray(x)
tf = isa(x,'deep.internal.recording.RecordingArray');
end

function tf = iIsDlarray(x)
tf = isa(x,'dlarray');
end
