function [stateInput, data] = extractState(stateInput, nObservationsX, errorIDFcn)
% Extract the state input's data, and returns the dlarray shell (if
% any) along with the extracted data
%
%   Extract the data, and get the channel index and observation index.
%
%   Copyright 2019 The MathWorks, Inc.

channelDim = 1; % default for unformatted
observationDim = 2; % default for unformatted
if isa(stateInput, 'dlarray')
    if ~isempty(dims(stateInput)) % formatted dlarray
        % Permute (if needed) to ensure labels are ordered CB
        stateInput = permuteExplicitly(stateInput);
        % Ensure a C dimension was specified
        channelDim = finddatadim(stateInput, 'C');
        if isempty(channelDim)
            error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled'));
        end
        % Find the B dimension
        observationDim = finddatadim(stateInput, 'B');
    end
    % Extract the data
    fd = stateInput.FormattedData;
    [~,data] = extractData(fd);
else % Numeric array
    deep.internal.dlarray.validateNonDlarray(stateInput);
    data = stateInput;
    stateInput = [];
end


% Validate the size of the state input
sizeIn = size(data);
% nbObservationInput = sizeIn(observationIdx);
% observationSizesMismatch = nbObservationInput ~= nObservationsX;
if ~isvector(data) % Input is matrix or N-array
    if isempty(observationDim) % Was specified as a dlarray with no 'B' dimension
        error(iGetErrorMessage(errorIDFcn, 'InvalidLabeled'));
    elseif sizeIn(observationDim) ~= nObservationsX % Has non-singleton 'B' but with an incorrect size
        error(iGetErrorMessage(errorIDFcn, 'InvalidStateObservation'));
    elseif length(sizeIn) > 2 % Too many dimensions
        error(iGetErrorMessage(errorIDFcn, 'TooManyDimensions'));
    end
else
    if sizeIn(channelDim) ~= 1 || isscalar(data) 
        % The vector length is the number of hidden units, need to expand
        % to obtain a CxB matrix
        data = data(:);
        data = data(:, ones(1, nObservationsX));
    else % Only one hidden unit
        % The vector length should be the number of observations
        if length(data) ~= nObservationsX
            error(iGetErrorMessage(errorIDFcn, 'InvalidStateObservation'));
        end
        % Ensure it is a line vector
        data = data(:)';
    end
end
end

function msg = iGetErrorMessage(errorIDFcn, IDName)
s = errorIDFcn();
msg =  message(s.(IDName));
end

