classdef deepSignalAnomalyDetectorLSTMForecaster < signal.internal.anomalyDetector.deepSignalAnomalyDetectorBase
%deepSignalAnomalyDetectorLSTMForecaster Detect signal anomalies using an LSTM forecaster
%
% deepSignalAnomalyDetectorLSTMForecaster properties:
%
%   ModelType               - Type of detector model
%
%   IsTrained               - True if the detector has been trained
%
%   NumChannels             - Number of channels expected for each input signal
%
%   NumHiddenUnits          - Number of hidden units of LSTM layers
%
%   ForecastSteps           - Number of steps forward for forecasting
%
%   WindowLength            - Length of window for window-based anomaly detection
%
%   OverlapLength           - Overlap length in samples between windows
%
%   WindowLossAggregation   - Method to aggregate sample losses within a
%                             window to get window loss
%
%   Threshold               - Reconstruction loss threshold used to detect anomalies
%
%   ThresholdMethod         - Method for threshold computation
%
%   ThresholdParameter      - Parameter for threshold computation
%
%   ThresholdFunction       - Custom function for threshold computation.
%                             Applies only when ThresholdMethod is
%                             "customFunction".
%
% deepSignalAnomalyDetectorLSTMForecasting methods:
%
%   trainDetector           - Train detector using a training data set
%
%   detect                  - Detect anomalies over an input data set
%
%   updateDetector          - Update settings of trained detector and
%                             recompute threshold based on updated parameters             
%
%   plotAnomalies           - Plot anomalies of an input signal
%
%   plotLossDistribution    - Plot CDF and histogram of window reconstruction
%                             loss
%                                                          
%   plotLoss                - Plot window reconstruction loss
%
%   getModel                - Get underlying neural network model of detector
%
%   resetState              - Reset the internal states of the model
%
%   saveModel               - Save the underlying neural network model and
%                             and detector parameters for loading into Deep 
%                             signal anomaly detector Simulink block

% Copyright 2023 The MathWorks, Inc.

    properties (SetAccess = protected)
        % LSTM Forecaster Model-specific parameters

        % Number of hidden units of LSTM layers
        NumHiddenUnits

        % Number of steps ahead for forecasting
        ForecastSteps
    end

    properties (Access = private)
        pDataMemory
        pLossMemory
        pAccumulatedStep;
    end

    methods (Hidden)
        function this = deepSignalAnomalyDetectorLSTMForecaster(numChannels,modelType,Nvargs,baseProps)
            %deepSignalAnomalyDetectorLSTM Detect signal anomalies using a deep learning LSTM auto encoder
            arguments
                numChannels
                modelType = 'lstmforecaster'
                Nvargs.NumHiddenUnits {mustBeNumeric,mustBeInteger,mustBePositive,mustBeVector} = 16;
                Nvargs.ForecastSteps (1,1) {mustBeNumeric,mustBeInteger,mustBePositive}  = 1;
                baseProps.?signal.internal.anomalyDetector.deepSignalAnomalyDetectorBase;
            end
            % Set properties in the deepSignalAnomalyDetectorBase
            basePropsCell = namedargs2cell(baseProps);
            this@signal.internal.anomalyDetector.deepSignalAnomalyDetectorBase(basePropsCell{:})

            % Validate channels
            validateattributes(numChannels,'numeric',{'positive','integer'},'deepSignalAnomalyDetector','number of channels');

            % Placeholder for modelType argument. To ensure error message has
            % consistent argument index with the wrapper function API.
            this.ModelType = validatestring(modelType,{'lstmforecaster'},'deepSignalAnomalyDetector','ModelType');

            this.ForecastSteps = Nvargs.ForecastSteps;

            this.NumChannels = numChannels;
            this.NumHiddenUnits = Nvargs.NumHiddenUnits(:)';
            this.buildModel()
            this.updateModelList()

            % Initialize Memory
            this.pDataMemory = [];
            this.pAccumulatedStep = 0;
            if isnumeric(this.WindowLength)
                this.pLossMemory = zeros(this.WindowLength-1,1);
            else
                this.pLossMemory = [];
            end
        end
    end

    methods (Access = protected)
        % Build the network
        function buildModel(this) % LSTM version
            this.pLayers = [sequenceInputLayer(this.NumChannels)];
            for ii = 1:length(this.NumHiddenUnits)
                this.pLayers = [this.pLayers,lstmLayer(this.NumHiddenUnits(ii))];
            end
            this.pLayers = [this.pLayers,fullyConnectedLayer(this.NumChannels)];
            this.pLayerGraph = layerGraph(this.pLayers);
            dlnet = dlnetwork(this.pLayerGraph);
            this.pDlnet= dlnet;
        end
    end

    %% Overridden methods (customized for lstmforecaster model)
    methods (Access=protected,Hidden=true)
        % Internal Training Method
        function [trainedNet,trainHistory] = iModelTrain(D,dataProcessed,trainingOpts)
            % Ensure that the size of the signals is greater than or equal
            % to D.ForecasterSteps+1. This guarantees valid data points
            % exist for comparison against forecasting.
            cellfun(@(x) assert(size(x,1)>D.ForecastSteps, ...
                message("signal_deep:deepSignalAnomalyDetector:InvalidSizeForecaster")), ...
                dataProcessed);

            % Index the data to obtain the expected input and the
            % corresponding forecasting ground truth.
            getDataX = @(X)cellfun(@(x) x(1:end-D.ForecastSteps, :), X, 'UniformOutput', false);
            getDataT = @(X)cellfun(@(x) x(1+D.ForecastSteps:end, :), X, 'UniformOutput', false);

            if isnumeric(trainingOpts.SequenceLength)
                trainingOpts.SequenceLength = "longest";
                [dataProcessedSeq,~] = signal.internal.anomalyDetector.segmentData(dataProcessed, ...
                    SequenceLength=D.pTrainingOptions.SequenceLength);

                [trainedNet,trainHistory] = trainnet( ...
                    getDataX(dataProcessedSeq), ...
                    getDataT(dataProcessedSeq),D.pDlnet, ...
                    "mean-squared-error",trainingOpts);
            else
                [trainedNet,trainHistory] = trainnet( ...
                    getDataX(dataProcessed), ...
                    getDataT(dataProcessed),D.pDlnet, ...
                    "mean-squared-error",trainingOpts);
            end
        end

        % Internal Prediction Method
        function reconstData = iModelPredict(D,dataProcessed,Nvargs)
            arguments
                D
                dataProcessed
                Nvargs.KeepState = false
                Nvargs.ExecutionEnvironment = "auto"
                Nvargs.MiniBatchSize = 128
            end

            % If update memory
            if Nvargs.KeepState
                % Not support batch prediction while keeping states
                if length(dataProcessed)>1
                    error(message("signal_deep:deepSignalAnomalyDetector:InvalidBatchWithState"))
                end
                X = dataProcessed{1};
                [Xpred,state] = predict(D.pDlnet,X);
                D.pDlnet.State = state;
                newDataStep = size(X,1);
                warmUpStep = min(max(0,D.ForecastSteps-D.pAccumulatedStep),newDataStep);
                D.pDataMemory = [D.pDataMemory;Xpred];
                Xpred(1:warmUpStep,:) = X(1:warmUpStep,:);
                Xpred(warmUpStep+1:newDataStep,:) =  D.pDataMemory(1:newDataStep-warmUpStep,:);
                D.pDataMemory = D.pDataMemory(newDataStep-warmUpStep+1:end,:);

                D.pAccumulatedStep = D.pAccumulatedStep+newDataStep;
                reconstData{1} = Xpred;
            else
                D.pAccumulatedStep = 0;
                if isnumeric(D.pTrainingOptions.SequenceLength)
                    nwarmup = floor(D.pTrainingOptions.SequenceLength*0.05);
                    % Segment data into target sequence length
                    [dataProcessedSeq,V] = signal.internal.anomalyDetector.segmentData(dataProcessed, ...
                        SequenceLength=D.pTrainingOptions.SequenceLength);
                    % Batch Prediction
                    % No zero padding when all sequences have a consistent size
                    Xpred = minibatchpredict(D.pDlnet,dataProcessedSeq, ...
                        ExecutionEnvironment=Nvargs.ExecutionEnvironment, ...
                        MiniBatchSize=Nvargs.MiniBatchSize, ...
                        UniformOutput=false);
                    % Replace the signals for the part not yet predicted
                    % and the warmup phase with the original signals
                    reconstDataSeq = cellfun(@(x,y) ...
                        [y(1:D.ForecastSteps+nwarmup,:);x(1+nwarmup:end-D.ForecastSteps,:)], ...
                        Xpred,dataProcessedSeq, ...
                        UniformOutput=false);
                    % Reassemble data into the original complete full-length sequence.
                    reconstData = signal.internal.anomalyDetector.reassembleData(reconstDataSeq,V, ...
                        SequenceLength=D.pTrainingOptions.SequenceLength);
                else
                    % Batch Prediction
                    Xpred = minibatchpredict(D.pDlnet,dataProcessed, ...
                        ExecutionEnvironment=Nvargs.ExecutionEnvironment, ...
                        MiniBatchSize=Nvargs.MiniBatchSize, ...
                        UniformOutput=false);
                    % Remove additional trailing zero padding
                    reconstData = cellfun(@(x,y)x(1:size(y,1),:),Xpred,dataProcessed,UniformOutput=false);
                    % Replace the signals for the part not yet predicted
                    % with the original signals
                    reconstData = cellfun(@(x,y) ...
                        [y(1:D.ForecastSteps,:);x(1:end-D.ForecastSteps,:)], ...
                        reconstData,dataProcessed, ...
                        UniformOutput=false);
                end
            end
        end

        % Internal Aggregation Method
        % Aggregate point losses to get window loss for one signal
        function winLoss = iAggregatePointLoss(D,loss,winLen)
            method = D.WindowLossAggregation;
            % When windowLength is Inf, i.e. "fullSignal"
            if ~isfinite(winLen)
                switch method
                    case "mean"
                        winLoss = mean(loss);
                    case "median"
                        winLoss = median(loss);
                    case "min"
                        winLoss = min(loss);
                    case "max"
                        winLoss = max(loss);
                end
            else
                loss = [D.pLossMemory;loss];
                % Directional window length
                dirWinLen = [winLen-1,0];
                switch method
                    case "mean"
                        winLoss = movmean(loss,dirWinLen,"Endpoints",0);
                    case "median"
                        % movmedian with gpuArray has limitation:
                        % winLen should <= 31
                        if winLen > 31 && isgpuarray(loss)
                            winLoss = movmedian(gather(loss),dirWinLen,"Endpoints",0);
                            winLoss = gpuArray(winLoss);
                        else
                            winLoss = movmedian(loss,dirWinLen,"Endpoints",0);
                        end
                    case "min"
                        winLoss = movmin(loss,dirWinLen,"Endpoints",0);
                    case "max"
                        winLoss = movmax(loss,dirWinLen,"Endpoints",0);
                end
                D.pLossMemory = loss(end-winLen+2:end,:);
                winLoss = winLoss(winLen:end);
            end
        end

        % Internal Reset Method
        function iResetState(D)
            D.pDlnet= resetState(D.pDlnet);
            D.pAccumulatedStep = 0;
            D.pDataMemory = [];
            D.pLossMemory = [];
            if isnumeric(D.WindowLength)
                D.pLossMemory = zeros(D.WindowLength-1,1);
            end
        end
    end
    %%
    methods (Access = private)
        % Update the list of model peroperties to display
        function updateModelList(this)
            this.pModelList = [
                this.pModelList,...
                "ForecastSteps",...
                "NumHiddenUnits"
                ];
        end
    end
end