classdef stftLayer < nnet.layer.Layer & nnet.layer.Formattable & ...
        nnet.internal.cnn.layer.CPUFusableLayer
    % stftLayer Short-time Fourier transform layer
    %   layer = stftLayer creates a short-time Fourier transform (STFT) layer
    %   with a 128-sample Hann window. There are 96 overlapped samples between
    %   successive windows. Equivalently, the stride between consecutive
    %   windows is 32 samples. 
    %
    %   Tensor inputs to stftLayer must be CBT arrays. The size of the T
    %   dimension of the tensor input must have at least as many samples as
    %   the window. 
    %
    %   By default, the output of stftLayer is an SCBT array. The spatial
    %   dimension corresponds to frequency. Use the 'OutputMode' parameter to
    %   control the format of the layer output.
    %
    %   layer = stftLayer('PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
    %   parameters for the layer using name-value arguments:
    %
    %   'Window'                Window used to compute the STFT, specified as
    %                           a vector with two or more elements. The default
    %                           is a Hann window of length 128.
    %
    %   'OverlapLength'         Number of overlapped samples, specified as a
    %                           a positive integer strictly smaller than the
    %                           length of Window. The default is 96 samples.
    %                           The stride between consecutive windows is the
    %                           difference between the window length and the
    %                           number of overlapped samples.
    %
    %   'FFTLength'             Number of frequency points used to calculate
    %                           the discrete Fourier transform, specified as an
    %                           integer greater than or equal to the window
    %                           length. If the length of the input data along
    %                           the time dimension is less than FFTLength, the
    %                           layer right-pads the data and the window with
    %                           zeros to have length FFTLength. The default is
    %                           the length of Window.
    %
    %   'TransformMode'         Layer transform mode, specified as one of 
    %                           these:
    %
    %                           'mag' - STFT magnitude (default)
    %
    %                           'squaremag' - STFT squared magnitude
    %
    %                           'logmag' - natural logarithm of STFT 
    %                           magnitude
    %
    %                           'logsquaremag' - natural logarithm of STFT
    %                           squared magnitude
    %
    %                           'realimag' - STFT real and imaginary parts
    %                           concatenated along the channel dimension
    %
    %   'OutputMode'            Layer output mode, specified as one of 
    %                           these:
    %
    %                           'spatiotemporal' - format output as a
    %                           sequence of 1-D images where the image
    %                           height corresponds to frequency, the second
    %                           dimension corresponds to channel, the third 
    %                           dimension corresponds to the batch,
    %                           and the fourth dimension corresponds to
    %                           time. Use this format in a 1-D
    %                           convolutional network when you want to 
    %                           convolve, or operate, along frequency. 
    %                           (default)                      
    %                           
    %                           'spatial' - format output as a 2-D image
    %                           where the image height corresponds to 
    %                           frequency and the image width corresponds to 
    %                           time. The third and fourth dimensions of
    %                           the output are channel and batch, 
    %                           respectively. Use this format in a 2-D
    %                           convolutional network when you want to
    %                           convolve along the two spatial dimensions,
    %                           which here represent frequency and time.
    %
    %                           'temporal' - format output as a 1-D sequence. 
    %                           This format takes the 'spatiotemporal'
    %                           output format and flattens the image height
    %                           into the channel dimension. The second
    %                           dimension of the output corresponds to the
    %                           batch and the third dimension corresponds
    %                           to time. Use this format in a 1-D
    %                           convolutional network if you wish to
    %                           convolve along time, or in a recurrent
    %                           network.
    %
    %   'WeightLearnRateFactor' Multiplier for the learning rate of the
    %                           weights, specified as a nonnegative scalar.
    %                           The default is 0, which results in weights
    %                           that do not update with training.
    %
    %   'Name'                  Layer name, specified as a character vector
    %                           or string scalar. The default is ''.
    %
    %   %Example 1:
    %   %   Create a short-time Fourier transform layer using the default
    %   %   values.
    %
    %   layer = stftLayer;
    %
    %   %Example 2:
    %   %   Create a short-time Fourier transform layer using a Kaiser 
    %   %   window. Specify the Kaiser window length to be 250 samples with 
    %   %   a shape factor of 2.5. Set the number of overlapped samples to
    %   %   100. Output the logarithm of the squared magnitude of the STFT. 
    %   %   Set the output mode to 'temporal'.
    %
    %   layer = stftLayer('Window',kaiser(250,2.5),'OverlapLength',100, ...
    %       'TransformMode','logsquaremag','OutputMode','temporal');
    %
    %   %Example 3:
    %   %   Create a two-layer dlnetwork containing a sequence input layer 
    %   %   and an STFT layer. There is one feature in the sequence input. 
    %   %   Set the minimum signal length in the sequence input layer to 
    %   %   1000 samples. Use a Hamming window of length 100. Set the
    %   %   weight learn rate factor to 1 and the output mode of the layer
    %   %   to 'spatial'. Run a batch of 10 random single-channel signals
    %   %   through the dlnetwork.
    %
    %   layers = [sequenceInputLayer(1,'Name','input','MinLength',1e3)
    %             stftLayer('Window',hamming(100),'OutputMode','spatial',...
    %                    'WeightLearnRateFactor',1)];
    %   dlnet = dlnetwork(layers);
    %   dataout = dlnet.forward(dlarray(randn(1,10,1000,'single'),'CBT'));
    %   size(dataout)
    %   dims(dataout)
    %
    %   See also dlstft
    %   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>
    
    
    %   Copyright 2021 The Mathworks, Inc.
    
    properties (Learnable)
        Weights
    end
    
    properties(Dependent)
        WeightLearnRateFactor
    end
    
    properties
        % Layer output type
        TransformMode
        % Format of output dimensions
        OutputMode
    end
    
    properties(SetAccess = private)
        % Layer properties.
        % Spectral window
        Window
        % Number of overlapped samples
        OverlapLength
        % Number of DFT points
        FFTLength
        
    end
    
    properties(Hidden,Access= private)
        PaddingSize = 0
    end
    
    methods
        function this = stftLayer(NVargs)
            arguments
                NVargs.Window {mustBeFloat,mustBeVector,mustBeNumeric,...
                    mustBeReal,mustBeFinite,mustBeNonempty} = hann(128,'periodic')
                NVargs.OverlapLength {mustBeScalarOrEmpty,mustBeReal}  = []
                NVargs.FFTLength  = []
                NVargs.TransformMode = 'mag'
                NVargs.Name = ''
                NVargs.OutputMode = 'spatiotemporal'
                NVargs.WeightLearnRateFactor = 0
            end
            this.Name = NVargs.Name;
            this.Type = 'Custom STFT Layer';
            this.Window = NVargs.Window(:);
            this.OverlapLength = NVargs.OverlapLength;
            if isempty(NVargs.FFTLength)
                this.FFTLength = length(this.Window);
            else
                this.FFTLength = NVargs.FFTLength;
            end
            if isempty(NVargs.OverlapLength)
                this.OverlapLength = floor(0.75*length(this.Window));
            end
            this.TransformMode = NVargs.TransformMode;
            % Validate the layer properties
            this.OutputMode = NVargs.OutputMode;
            this = validateProperties(this);
            this.PaddingSize = this.FFTLength-length(this.Window);
            this = initializeWeights(this);
            this.WeightLearnRateFactor = NVargs.WeightLearnRateFactor;
            
        end
        
        function Z = predict(this, X)
            
            % Forward input data through the layer at prediction time and
            % output the result.
            if ~isequal(dims(X),'CBT')
                error(message('signal_deep:stftLayer:NotCBT'));
            end
            
            T = size(X,3);
            nwin = length(this.Window);
            if T < nwin
                error(message('signal_deep:stftLayer:LessThanWindow'));
            end
            
            % Add padding if FFTLength > length(Window)
            if this.PaddingSize
                X(:,:,T:T+this.PaddingSize) = 0;
            end
            
            [YR,YI] = dlstft(X,'OverlapLength',this.OverlapLength,...
                'Weights',{this.Weights,nwin});
            
            Z = iFormatOutput(this,YR,YI);
            if ~strcmp(this.OutputMode,'spatiotemporal')
                Z = iFormatTensor(this,Z);
            end
        end
        
        function this = set.WeightLearnRateFactor(this,lrate)
            this = setLearnRateFactor(this,"Weights",lrate);
        end
        
        function lrate = get.WeightLearnRateFactor(this)
            lrate = getLearnRateFactor(this,"Weights");
        end
        
        function this = set.OutputMode(this,outputMode)
            this.OutputMode = outputMode;
        end
        
        function this = set.TransformMode(this,transformMode)
            this.TransformMode = transformMode;
        end        
        
    end
    
    methods (Hidden)
        
        function this = validateProperties(this)
            nwin = length(this.Window);
            % Validate window length
            try
                validateattributes(nwin,{'numeric'},...
                    {'integer','>',1},'stftLayer');
            catch e
                error(message('signal_deep:stftLayer:windowlength'));
            end
            
            % OverlapLength
            validateattributes(this.OverlapLength,{'numeric'},...
                {'scalar','integer','nonnegative','real','nonnan',...
                'finite','nonempty','>=',0,'<',nwin},...
                'stftLayer','OverlapLength');
            
            % FFTLength
            validateattributes(this.FFTLength,{'numeric'},...
                {'scalar','integer','nonnegative','real','nonnan',...
                'finite','nonempty','>=',nwin},'stftLayer','FFTLength');
            
            % TransformMode
            validTransformMode = {'mag','squaremag','logmag',...
                'logsquaremag','realimag'};
            this.TransformMode = validatestring(this.TransformMode,...
                validTransformMode,'stftLayer','TransformMode');
            
            % Number of spatial dimensions
            validDims = {'spatiotemporal','spatial','temporal'};
            this.OutputMode = ...
                validatestring(this.OutputMode,validDims,'stftLayer',...
                'OutputMode');
        end
        
        
        function Z = iFormatOutput(this,YR,YI)
            YR = stftLayer.iboundAwayFromZero(YR);
            YI = stftLayer.iboundAwayFromZero(YI);
            switch this.TransformMode
                case 'mag'
                    Z = sqrt(YR.^2+YI.^2);
                case 'logmag'
                    Z = log(sqrt(YR.^2+YI.^2));
                case 'squaremag'
                    Z = YR.^2+YI.^2;
                case 'logsquaremag'
                    Z = log(YR.^2+YI.^2);
                case 'realimag'
                    dimC = finddim(YR,'C');
                    Z = cat(dimC,YR,YI);
            end
        end
        
        function Z = iFormatTensor(this,Z)
            switch this.OutputMode
                case 'temporal'
                    sz = size(Z);
                    tensorOutputSize = [prod(sz(1:2)) sz(3:end)];
                    Z = reshape(Z,tensorOutputSize);
                    Z = dlarray(Z,'CBT');
                case 'spatial'
                    Z = dlarray(Z,'SCBS');
            end
        end        
            
        function this = initializeWeights(this)
            % Initialize windowed DFT basis
            nwin = length(this.Window);
            win = this.Window;
            if nwin < this.FFTLength
                win(end+this.FFTLength-nwin,1) = 0;
            end
            
            nfreq = fix(this.FFTLength/2+1);
            N = length(win);
            n = 0:N-1;
            idx = 1:nfreq;
            
            wR = win.*cos(2*pi*(idx-1).*n'/N);
            wI = win.*-sin(2*pi*(idx-1).*n'/N);
            this.Weights = cat(3,wR,wI);
            this.Weights = cast(this.Weights,'single');
            
        end
        
        function Z = cpuInferenceEnginePredict(this, X)
            % Prediction method for the CPU inference engine.
            X = dlarray(X,'CBT');
            Z = this.predict(X);
            Z = extractdata(Z);
        end
        
        function layerArgs = getFusedArguments(this)
            % getFusedArguments  Returned the arguments needed to call the
            % layer in a fused network.
            % 1,1 refers to 1 input, 1 output argument
            layerArgs = { 'external', @this.cpuInferenceEnginePredict, ...
                1, 1, @this.propagateSize};
        end
        
        function tf = isFusable(~)
            % isFusable  Indicates if the layer is fusable in a given network.
            tf = true;
        end
        
        function szout = propagateSize(this,szIn)
            % This method calculates the output size of stftLayer. This is
            % needed for inference engine.
            if strcmp(this.TransformMode,'realimag')
                szIn(1) = szIn(1)+szIn(1);
            end
            nwin = length(this.Window);
            hopSize = nwin-this.OverlapLength;
            nT = fix((szIn(3)-this.OverlapLength)/hopSize);
            Sdim = floor(this.FFTLength/2)+1;
            % Canonical form: SCBT
            tmpsize = [Sdim szIn(1) szIn(2) nT];
            
            % For CBT output (temporal), we need to collapse the spatial
            % (frequency) dimension into the channel dimension.
            % For 'spatial', 2D convolutional networks convolve along S
            % dimensions, so we have to perumte the array SCBT becomes SSCB
            switch this.OutputMode
                case 'temporal'
                    szout = [tmpsize(1)*tmpsize(2) tmpsize(3) tmpsize(4)];
                case 'spatial'
                    % Mark temporal dim for removal in the inference
                    % engine.
                    szout = [tmpsize(1) tmpsize(4) tmpsize(2) tmpsize(3) 0];
                otherwise
                    szout = tmpsize;
                    
            end            
        end        
    end
    
    methods(Static)
       function xBounded = iboundAwayFromZero(x)
           xBounded = x;
           thresh = eps(underlyingType(xBounded));
           xBounded(abs(xBounded) < thresh) = ...
               sign(xBounded(abs(xBounded) < thresh)).*thresh;
            
        end         
    end
end