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. 
    %   
    %   The output of stftLayer is an SCBT array. The spatial
    %   dimension corresponds to frequency.
    %
    %
    %   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. The weights
    %                           in stftLayer are equal to the window vector
    %                           in single precision. It is not recommended 
    %                           to initialize the weights directly.
    %
    %   '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. 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
    %
    %   '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. 
    %
    %   layer = stftLayer('Window',kaiser(250,2.5),'OverlapLength',100, ...
    %       'TransformMode','logsquaremag');
    %
    %   %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. Run a batch of 10 random 
    %   %   single-channel signals through the dlnetwork.
    %
    %   layers = [sequenceInputLayer(1,'Name','input','MinLength',1e3)
    %             stftLayer('Window',hamming(100),...
    %                    'WeightLearnRateFactor',1)];
    %   dlnet = dlnetwork(layers);
    %   dataout = forward(dlnet,dlarray(randn(1,10,1000,'single'),'CBT'));
    %   size(dataout)
    %   dims(dataout)
    %
    %   See also DLSTFT, STFT
    %   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>
    
    
    %   Copyright 2021-2023 The MathWorks, Inc.
    
    properties (Learnable,Dependent)
        Weights
    end

    properties(Access=private,Hidden)
        pWeights
        pWindowLength
    end
    
    properties(Dependent)
        WeightLearnRateFactor
    end
    
    properties 
        % Layer output type
        TransformMode
    end

    properties (Hidden)
        % 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  
        
    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
            NVargs = nnet.internal.cnn.layer.util.gatherParametersToCPU(NVargs);
            this.Version = 2;
            this.Name = NVargs.Name;
            this.Type = 'STFT';
            this.Window = NVargs.Window(:);
            if isdlarray(this.Window)
                this.Window = extractdata(this.Window);
            end
            this.OverlapLength = NVargs.OverlapLength;
            this.pWindowLength = length(this.Window);
            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.Weights = single(this.Window);
            this.WeightLearnRateFactor = NVargs.WeightLearnRateFactor;            
        end
        
        function Z = predict(this, X)
            arguments
                this
                X {mustBeReal}
            end
            if ~isequal(dims(X),'CBT')
                error(message('signal_deep:stftLayer:NotCBT'));
            end
            
            T = size(X,3);
            nwin = length(this.Weights);
            if T < nwin
                error(message('signal_deep:stftLayer:LessThanWindow'));
            end

            S = signalwavelet.deep.internal.stft.computeDLSTFT(real(X), ...
                real(this.Weights), ...
                this.OverlapLength, ...
                this.FFTLength);
            
            Z = iFormatOutput(this,S);
            if ~strcmp(this.OutputMode,'spatiotemporal')
                Z = iFormatTensor(this,Z);
            end
        end
        
        function this = set.WeightLearnRateFactor(this,lrate)
            this = setLearnRateFactor(this,"Weights",lrate);
        end
        
        function this = set.Weights(this,weights)
            validateattributes(weights,{'double','single'},{'vector','real',...
                'finite'},'stftLayer','Weights');
            [rowsize,colsize] = size(weights);
            weights = weights(:);            
            if length(weights) ~= this.pWindowLength         
                error(message('signal_deep:stftLayer:WeightSize',...
                    this.pWindowLength,rowsize,colsize));
            end      
            this.pWeights = single(weights);            
        end

        function val = get.Weights(this)
            val = this.pWeights;
        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,S)
            switch this.TransformMode
                case 'mag'
                    Z = abs(S);
                case 'logmag'
                    Z = log(stftLayer.iboundAwayFromZero(abs(S)));
                case 'squaremag'
                    Z = abs(S).*abs(S);
                case 'logsquaremag'
                    Z = log(stftLayer.iboundAwayFromZero(abs(S).*abs(S)));
                case 'realimag'
                    dimC = finddim(S,'C');
                    Z = cat(dimC,real(S),imag(S));
            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 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 
        
        function obj = saveobj(this)
            obj = this;
            obj.Weights = gather(obj.Weights);
            obj.WeightLearnRateFactor = gather(obj.WeightLearnRateFactor);
        end
    end

    methods(Static,Hidden)
        function obj = loadobj(obj)
            obj.Weights = gather(obj.Weights);
            obj.WeightLearnRateFactor = gather(obj.WeightLearnRateFactor);
        end

      function xBounded = iboundAwayFromZero(x)
           xBounded = nnet.internal.cnn.util.boundAwayFromZero(x);                    
       end
    end
end
