function [dZ, dh0, dc0, dW, dR, db] = internal_lstmBackward(dZ, dh0, dc0, X, h0, c0, W, R, b, Z, H, C, Ws, stateActivation, gateActivation)
% This function is undocumented and will change in a future release.

% The help for this method is similar to the CPU method
% internal_lstmBackward
%
% This method needs the input Ws, which is the workspace generated by cuDNN
% during the forward pass.

%   Copyright 2019-2023 The MathWorks, Inc.

import deep.internal.recording.errorOnComplexGradientOutputs

dZ = errorOnComplexGradientOutputs(dZ, 'lstm');
dh0 = errorOnComplexGradientOutputs(dh0, 'lstm');
dc0 = errorOnComplexGradientOutputs(dc0, 'lstm');

% Ensure all inputs have the same data type
%   Regardless of where the forward path was computed, the GPU LSTM
%   implementation only expects the final cell state, so index here
[dZ, dh0, dc0, X, h0, c0, W, R, b, Z, H, C] = ...
    nnet.internal.cnngpu.castToSuperiorFloat(dZ, dh0, dc0, X, h0, c0, W, R, b, Z, H, C);

if iHasDefaultActivations(stateActivation, gateActivation)
    % Ensure only final element of cell state is used for cuDNN backward.
    C = C(:,:,end);

    % If Ws is empty it means the forward pass was computed on the CPU. We
    % need to create that workspace so we must call the GPU forward method
    % in order to generate the workspace.
    if isempty(Ws)
        % Make sure the GPU version is definitely called by moving one of the
        % inputs to the GPU.
        [~,~,~,Ws] = internal_lstm(gpuArray(X), h0, c0, W, R, b, stateActivation, gateActivation, true);
    end

    % Prepare data for builtin
    [Wc, nHiddenUnits, nInputChannels] = iPrepareDataForBackward(W, R, b);
    args = { Z, H, C, dZ, dh0, dc0, X, h0, c0, Wc, Ws };

    % Call cuDNN backward method
    [dZ, dh0, dc0, dWc] = nnet.internal.cnngpu.lstmBackward(args{:});

    % If other outputs requested, convert back the CuDNN weights vector
    if nargout > 1
        % Convert to NNT weights format
        [dW, dR, db] = nnet.internal.cnngpu.util.LSTMWeightsConverter.fromCudnnDerivative(dWc, nHiddenUnits, nInputChannels);
    end

else
    % Prepare the inputs for cnnhost implementation
    learnable.W = W;
    learnable.R = R;
    learnable.b = b;
    state.c0 = c0;
    state.y0 = h0;
    outputs.Y = Z;
    outputs.CS = C;
    outputs.HS = H;
    derivatives.dZ = dZ;
    derivatives.dHS = dh0;
    derivatives.dCS = dc0;
    options.ReturnLast = false;
    options = iSetActivationFcns(options, stateActivation, gateActivation);

    % Call cnnhost implementation
    [dZ, dW, dR, db, dh0, dc0] = ...
        nnet.internal.cnngpu.lstmBackwardGeneral(X, learnable, state, outputs, derivatives, options);

end
end

function [Wc, nHiddenUnits, nInputChannels] = iPrepareDataForBackward(W, R, b)
% Prepare data passed into the strategy for the internal
% backward method

% Convert to cuDNN weights format
Wc = nnet.internal.cnngpu.util.LSTMWeightsConverter.toCudnn(W, R, b);
% Determine problem dimensions
[fourH, nInputChannels] = size( W );
nHiddenUnits = fourH./4;
end

function tf = iHasDefaultActivations(stateActivation, gateActivation)
tf = (stateActivation == "tanh") && (gateActivation == "sigmoid");
end

function options = iSetActivationFcns(options, stateActivation, gateActivation)
if stateActivation == "tanh"
    options.Activation = @tanh;
    options.dActivation = @iTanhDiff;
elseif stateActivation == "relu"
    options.Activation = @iRelu;
    options.dActivation = @iReluDiff;
else
    options.Activation = @iSoftSign;
    options.dActivation = @iSoftSignDiff;
end

if gateActivation == "sigmoid"
    options.RecurrentActivation = @iSigmoid;
    options.dRecurrentActivation = @iSigmoidDiff;
else
    options.RecurrentActivation = @nnet.internal.cnnhost.hardSigmoidForward;
    options.dRecurrentActivation = @nnet.internal.cnnhost.hardSigmoidBackward;
end
end

%% Activation functions and derivatives
function y = iSigmoid(x)
y = 1./(1 + exp(-x));
end

function y = iSigmoidDiff(x)
y = ( 0.5.*sech(0.5.*x) ).^2;
end

function y = iTanhDiff(x)
y = sech(x).^2;
end

function y = iRelu(x)
y = max(0,x);
end

function y = iReluDiff(x)
y = 1 .* (x > 0);
end

function y = iSoftSign(x)
y = x./(1 + abs(x));
end

function y = iSoftSignDiff(x)
y = (1 + abs(x)).^-2;
end
