classdef DifferentiableFunction <  matlab.mixin.internal.FunctionObject
    %DifferentiableFunction   Interface for custom differentiable function
    %
    %  Example:
    %    % Define a DifferentiableFunction to round the input for the forward
    %    % pass and be the identity function in the backward pass.
    %    classdef roundFunction < deep.DifferentiableFunction
    %        methods
    %            function fun = roundFunction()
    %                fun@deep.DifferentiableFunction(1);
    %            end
    %            function y = forward(fun,x)
    %                y = round(x);
    %            end
    %            function dLdX = backward(fun,dLdY,computeGradients)
    %                dLdX = dLdY;
    %            end
    %        end
    %    end
    %    newround = roundFunction;
    %    dlfeval(@(x)dlgradient(newround(x),x), dlarray(rand))
    %    % displays 1
    %    dlfeval(@(x)dlgradient(round(x),x), dlarray(rand))
    %    % displays 0

    %   Copyright 2023 The MathWorks, Inc.

    properties(Access = private)
        NumOutputs
        NumMemoryValues
        SaveInputsForBackward
        SaveOutputsForBackward
    end

    methods(Abstract)
        % [Y1,...,YM,Mem1,...,MemQ] = forward(fun,X1,...,XN)
        % Forward input data through the function and output the results.
        %
        % Inputs:
        %   fun           - DifferentiableFunction to forward propagate
        %                   through.
        %   X1,...XN      - Input data.
        % Outputs:
        %   Y1,...YM      - Outputs of forward function. Outputs must be
        %                   numeric.
        %   Mem1,...,MemQ - (Optional) Memory values to pass to backward
        %                   method.
        % M is the number of outputs. Q is the number of memory values.
        % M and Q are specified in the DifferentiableFunction constructor.
        varargout = forward(fun,varargin)

        % [dLdX1,...,dLdXN] = backward(fun,dLdY1,...,dLdYM,computeGradients,X1,...,XN,Y1,...YM,Mem1,...,MemQ)
        % Backward propagate the derivative of the loss function through the DifferentiableFunction.
        %
        % Inputs:
        %   fun              - DifferentiableFunction to backward propagate
        %                      through.
        %   dLdY1,...dLdYM   - Derivative of loss with respect to outputs.
        %   computeGradients - Logical vector of length N indicating which
        %                      gradients to compute. If computeGradients(i)
        %                      is false, then dLdXi is unused for subsequent
        %                      computations and can be any value. The
        %                      values of computeGradients depend on the
        %                      computation graph and the inputs to dlgradient.
        %   X1,...,XN        - (Optional) Forward inputs
        %   Y1,...,YM        - (Optional) Forward outputs
        %   Mem1,...,MemQ    - (Optional) Memory values to pass to backward
        %                      method.
        % Outputs:
        %   dLdX1,...dLdXN   - Derivative of loss with respect to inputs
        % M is the number of outputs. Q is the number of memory values.
        % M and Q are specified in the DifferentiableFunction constructor.
        varargout = backward(fun,varargin)
    end

    methods
        function fun = DifferentiableFunction(numOutputs, options)
            % DifferentiableFunction  Construct a differentiable function
            %   fun = DifferentiableFunction(numOutputs) creates a
            %   differentiable function. Evaluate fun using the syntax 
            %     [Y1,...,YM] = fun(X1,...,XN)
            %   M is numOutputs.
            %
            %   [___] = DifferentiableFunction(___,Name=Value) specifies
            %   additional options using one or more name-value arguments
            %
            %    NumMemoryValues        - Number of memory values to pass
            %                             from forward to backward.
            %                             The default is 0.
            %    SaveInputsForBackward  - Save inputs to pass to backward.
            %                             The default is false.
            %    SaveOutputsForBackward - Save outputs to pass to backward.
            %                             The default is false.

            arguments
                numOutputs(1,1) double {mustBePositive,mustBeInteger}
                options.NumMemoryValues(1,1) double {mustBeNonnegative,mustBeInteger} = 0
                options.SaveInputsForBackward(1,1) logical = false
                options.SaveOutputsForBackward(1,1) logical = false
            end
            fun.NumOutputs = numOutputs;
            fun.NumMemoryValues = options.NumMemoryValues;
            fun.SaveInputsForBackward = options.SaveInputsForBackward;
            fun.SaveOutputsForBackward = options.SaveOutputsForBackward;
        end
    end

    methods (Hidden)
        function varargout = parenReference(fun,X)
            arguments
                fun deep.DifferentiableFunction
            end
            arguments(Repeating)
                X
            end
            numInputs = length(X);
            numOutputs = fun.NumOutputs;

            nargoutchk(0,numOutputs);
            if isempty(X)
                error(message('deep:DifferentiableFunction:NoDataInputs'));
            end

            % Check inputs for dlarrays.
            hasDlarrayInput = false;
            for i = 1:length(X)
                if isdlarray(X{i})
                    if ~isempty(dims(X{i}))
                        X{i} = stripdims(X{i});
                    end
                    hasDlarrayInput = true;
                else
                    try
                        X{i} = dlarray(X{i});
                    catch
                        % Throw a custom error explaining what the inputs must be. The dlarray errors
                        % assume the user themselves are calling the dlarray function.
                        error(message('deep:DifferentiableFunction:NonDlarrayInput'));
                    end
                end
            end

            % The "+1" is to include the Sizes memory output which we need internally but don't return to users.
            numInternalOutputs = numOutputs + 1;

            numMemory = fun.NumMemoryValues;
            numInternalOutputs = numInternalOutputs + numMemory;

            fcn = deep.internal.CustomBackwardMethod(fun,numInputs,numOutputs,numMemory,...
                fun.SaveInputsForBackward,fun.SaveOutputsForBackward);
            try
                % The struct below is to get in-place optimizations of the arrays in X.
                [varargout{1:numInternalOutputs}] = deep.internal.dlarray.extension.applyExtensionMethod(fcn, struct('v',matlab.lang.internal.move(X)).v);
            catch err
                id = 'deep:DifferentiableFunction:EvaluationError';
                ex = MException(id,message(id));
                ex = addCause(ex,err);
                throw(ex);
            end

            % Extract data from outputs if none of the inputs were dlarrays.
            if ~hasDlarrayInput
                for i = 1:max(1,nargout)
                    varargout{i} = extractdata(varargout{i});
                end
            end
        end

        function varargout = feval(fun,X)
            arguments
                fun deep.DifferentiableFunction
            end
            arguments(Repeating)
                X
            end
            [varargout{1:nargout}] = fun(X{:});
        end

        function n = nargin(~)
            % Nargin is -1 meaning it is like "varargin". There is no
            % reliable way to automatically detect nargin from the subclass
            % definition.
            n = -1;
        end

        function n = nargout(fun)
            n = fun.NumOutputs;
        end

        function c = num2cell(fun)
            c = {fun};
        end

    end
end
