function varargout = npufun(kernelFun, kernelSize, input_data, options)
%HDL.NPUFUN applies the kernelFun function to each sliding window of
% kernelSize of the input data. The size of the sliding window is
% determined by kernelSize. Function kernelFun is called for each
% kernelSize window of the input and computes an element of the output.
% output_data is same size as input_data.
%
%     output_data = hdl.npufun(kernelFun, kernelSize, input_data);
%
% The size of the output data matches with the size of the input data.
%
% Example: Call npufun to apply image blurring on image input A.
%
%     output_data = hdl.npufun(@blurringKernel, kernelSize, input_data);
%
% Example: Call npufun to blur image with custom boundary constant.
%
%     output_data = hdl.npufun(@blurringKernel, kernelSize, input_data, ...
%         'BoundaryConstant', 5);
%
% Example: Call npufun to blur image with boundary pixels replicated.
%
%     output_data = hdl.npufun(@blurringKernel, kernelSize, input_data, ...
%         'BoundaryMethod', 'replicate');
%
% Example: Call npufun to apply custom coefficient to the input data.
%
%     output_data = hdl.npufun(@blurringKernel, kernelSize, input_data, ...
%         'KernelArg', 3);
%
% Example: Call npufun with kernel returning multiple outputs
%
%     [output_data1, output_data2] = hdl.npufun(@multiOutKernel, ...
%         kernelSize, input_data, 'KernelArg', 3);
%
% See also hdl.iteratorfun.

%   Copyright 2021-2023 The MathWorks, Inc.

%#codegen

arguments
    kernelFun (1,1) function_handle {mustBeValidKernelFunction}
    kernelSize {mustBeInteger, mustBePositive, mustBeVector, mustBe2D, mustBeConst(kernelSize,'kernelSize')}
end

arguments (Repeating)
    input_data
end

arguments
    options.BoundaryConstant (1,1) {mustBeInteger, mustBeInt32} = 0
    options.BoundaryMethod (1,:) char ...
        {mustBeValidBoundaryMethod, mustBeConst(options.BoundaryMethod,'BoundaryMethod')} = 'constant'
end


firstImgIdx = 0;

%% Read in streamed and non-streamed inputs

% All streamed or non-streamed inputs to the kernel will have the
% associated value in this array set to true
isKernelInput = false(1, numel(input_data));

% All streamed inputs to the kernel will have the associated value in this
% array set to true
isStreamedInput = false(1, numel(input_data));

for i=coder.unroll(1:numel(input_data))

    if ischar(input_data{i})
        option = lower(input_data{i});

        % Check that this isn't a trailing Name/Value pair with missing
        % Value field
        coder.internal.assert(~strcmpi(option, 'boundaryconstant') && ...
            ~strcmpi(option, 'boundarymethod'), ...
            'hdlmllib:hdlmllib:IncompleteNameValuePair');

        % Check if there are any arguments coming after and 
        % that this is a valid option flag
        coder.internal.assert(i+1 <= numel(input_data) && ...
            (strcmpi(option,'kernelarg') || strcmpi(option,'nonsampleinput')), ...
            'hdlmllib:hdlmllib:NpufunInvalidArgument', input_data{i})

        % The next index is a value for the kernel function
        isKernelInput(i+1) = true;

    elseif i==1 || ~ischar(input_data{i-1})
        % If this is the first input or an input not preceeded by a char
        % array option, then it is a streamed input
        img = input_data{i};

        if firstImgIdx == 0
            % This is the first input image we have found. Make sure
            % that it is no more than 2D and that it is non-empty.
            coder.internal.assert(numel(img) > 0, ...
                'hdlmllib:hdlmllib:NpufunImageMustNotBeEmpty');
            coder.internal.assert(ndims(img) <= 2, ...
                'hdlmllib:hdlmllib:NpufunImageMustBe2D');

            firstImgIdx = i;
        else
            % For later images, make sure that their size matches the
            % first image we found
            firstImg = input_data{firstImgIdx};
            coder.internal.assert(isequal(size(firstImg), size(img)), ...
                'hdlmllib:hdlmllib:NpufunImagesMustBeSameSize');
        end

        % Mark this index as a streamed kernel input
        isKernelInput(i) = true;
        isStreamedInput(i) = true;
    end
end

inputIdxs = find(isKernelInput);
streamedIdxs = find(isStreamedInput(isKernelInput));

allInputs = cell(1, numel(inputIdxs));

for i=1:numel(inputIdxs)
    allInputs{i} = input_data{inputIdxs(i)};
end

coder.internal.assert(numel(streamedIdxs) > 0, 'hdlmllib:hdlmllib:NpufunNeedsInputImage');

argOptions = options;

%% Apply kernel function

% get idx offsets for x and y
xIdxOffs = getIdxOffsets(kernelSize(1));
yIdxOffs = getIdxOffsets(kernelSize(2));

sampleOutput = allInputs{streamedIdxs(1)};
[outSizeX, outSizeY] = size(sampleOutput);

% call the function once to initialize the output(s)
numOut = nargout(kernelFun);
outs = cell(1, numOut);

% Boundary method doesn't matter for initialization
[outs{:}] = callKernelFcn(kernelFun, xIdxOffs, yIdxOffs, outSizeX, outSizeY, ...
    streamedIdxs, allInputs, ...
    coder.internal.ignoreRange(1), coder.internal.ignoreRange(1), argOptions);

% Initialize the output image(s)
imagesOut = cell(1, numOut);
for i=1:numOut
    coder.internal.assert(isscalar(outs{i}), ...
        'hdlmllib:hdlmllib:NpufunKernelMustHaveScalarOutputs');

    imagesOut{i} = repmat(outs{i}, outSizeX, outSizeY);
end

for i=1:outSizeX
    for j=1:outSizeY
        if i == 1 && j == 1
            continue;
        end

        [outs{:}] = callKernelFcn(kernelFun, xIdxOffs, yIdxOffs, outSizeX, outSizeY, ...
            streamedIdxs, allInputs, i, j, argOptions);

        for k=1:numOut
            imagesOut{k}(i, j) = outs{k};
        end
    end
end

varargout = imagesOut;
end

%% Helper functions

function idxOffs = getIdxOffsets(kernelSize)
if mod(kernelSize, 2) == 0
    halfKernelSize = kernelSize / 2;
    idxOffs = -(halfKernelSize-1) : halfKernelSize;
else
    halfKernelSize = (kernelSize - 1) / 2;
    idxOffs = -halfKernelSize : halfKernelSize;
end
end


%% Kernel Function Caller
function varargout = callKernelFcn(kernelFun, xIdxOffs, yIdxOffs, outSizeX, outSizeY,...
    imagesIdx, allInputs, xIdx, yIdx, options)

coder.inline('always');

% determine indices for the current window
xIdxRange = xIdx + xIdxOffs;
yIdxRange = yIdx + yIdxOffs;

xIsOOB_pre  = xIdxRange < 1;
xIsOOB_post = xIdxRange > outSizeX;
yIsOOB_pre  = yIdxRange < 1;
yIsOOB_post = yIdxRange > outSizeY;

% replicate borders by default to make all indices in range
xIdxRange(xIsOOB_pre) = 1;
yIdxRange(yIsOOB_pre) = 1;
xIdxRange(xIsOOB_post) = outSizeX;
yIdxRange(yIsOOB_post) = outSizeY;

% If boundary method is 'constant', overwrite the relevant cells in the loop
useBoundaryConstant = strcmpi(options.BoundaryMethod,'constant');

% Calculate windows
allInputsWithWindows = cell(1, numel(allInputs));
for ii=coder.unroll(1:numel(allInputs))
    if ismember(ii, imagesIdx)
        img = allInputs{ii};
        wdw = img(xIdxRange, yIdxRange);

        if useBoundaryConstant
            wdw(xIsOOB_pre | xIsOOB_post, :) = options.BoundaryConstant;
            wdw(:, yIsOOB_pre | yIsOOB_post) = options.BoundaryConstant;
        end

        allInputsWithWindows{ii} = wdw;
    else
        allInputsWithWindows{ii} = allInputs{ii};
    end
end

% Apply kernel function to window
pixelsOut = cell(1, nargout(kernelFun));
[pixelsOut{:}] = kernelFun(allInputsWithWindows{:});

varargout = pixelsOut;
end


%% Argument Validation Functions

function mustBeValidKernelFunction(kernelFun)
name = func2str(kernelFun);
coder.internal.assert(name(1) ~= '@', ...
    'hdlmllib:hdlmllib:KernelMustNotBeAnonymous', 'hdl.npufun');
coder.internal.assert(nargout(kernelFun) ~= 0, ...
    'hdlmllib:hdlmllib:KernelMustHaveAnOutput', 'hdl.npufun');
coder.internal.assert(nargout(kernelFun) > 0, ...
    'hdlmllib:hdlmllib:KernelMustNotUseVarargout', 'hdl.npufun');
coder.internal.assert(nargin(kernelFun) >= 0, ...
    'hdlmllib:hdlmllib:KernelMustNotUseVarargin', 'hdl.npufun');
end

function mustBe2D(kernelSize)
coder.internal.assert(numel(kernelSize) == 2, ...
    'hdlmllib:hdlmllib:NpufunKernelSizeIncorrect');
end

function mustBeConst(arg, name)
coder.internal.prefer_const(arg);
coder.internal.assert(coder.internal.isConst(arg), ...
    'hdlmllib:hdlmllib:ArgumentMustBeConstant', name, 'hdl.npufun');
end

function mustBeInt32(arg)
% check that this value can be represented as int32
if arg ~= cast(int32(arg), 'like', arg)
    coder.internal.error('hdlmllib:hdlmllib:NpufunInvalidBoundaryConstant');
end
end

function mustBeValidBoundaryMethod(method)
coder.internal.assert(strcmpi(method,'constant') || ...
    strcmpi(method,'replicate'), ...
    'hdlmllib:hdlmllib:NpufunInvalidBoundaryMethod');
end

% LocalWords:  iteratorfun boundaryconstant boundarymethod kernelarg nonsampleinput OOB
