classdef istftLayer < nnet.layer.Layer & nnet.layer.Formattable...
        & nnet.layer.Acceleratable & nnet.internal.cnn.layer.CPUFusableLayer
    % istftLayer Inverse short-time Fourier transform layer
    %   layer = istftLayer creates an inverse short-time Fourier transform
    %   (ISTFT) layer with a 128-sample Hann window. There are 96
    %   overlapped samples between adjoining segments. Equivalently, the stride
    %   between adjoining segments is 32 samples.
    %
    %   Tensor inputs to istftLayer must be real-valued CBT or SCBT arrays.
    %   The size of the first dimension of both SCBT and CBT inputs
    %   must be divisible by floor(FFTLength/2)+1. For CBT inputs, the
    %   channel (C) dimension must be even.
    %
    %   The output of istftLayer is a real-valued CBT array.
    %
    %   layer = istftLayer(PARAM1=VAL1,PARAM2=VAL2,...) specifies optional
    %   parameters for the layer using name-value arguments:
    %
    %   Window                  Window used to compute the ISTFT, specified
    %                           as a vector with two or more elements. For
    %                           perfect reconstruction, use the same window
    %                           used in stftLayer. The default is a Hann
    %                           window of length 128. The weights in
    %                           istftLayer 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
    %                           positive integer strictly smaller than the
    %                           length of the window. This argument defaults to
    %                           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 compute
    %                           the discrete Fourier transform, specified as an
    %                           integer greater than or equal to the window
    %                           length. This argument defaults to the window
    %                           length.
    %
    %   WeightLearnRateFactor   Multiplier for the learning rate of the
    %                           weights, specified as a nonnegative scalar.
    %                           This argument defaults to 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 ''.
    %
    %   Method                  Overlap-add method to use for
    %                           reconstruction. Valid options are: "ola"
    %                           (Overlap-add) and 'wola' (Weighted overlap-add).
    %                           If unspecified, Method defaults to "wola".
    % 
    %   ExpectedOutputSize      Expected number of channels and samples
    %                           output by istftLayer, specified as "none"
    %                           or a two-element real-valued vector. If
    %                           unspecified, ExpectedOutputSize defaults to
    %                           "none" and istftLayer does not check the
    %                           output size of the ISTFT. If
    %                           ExpectedOutputSize is specified as a
    %                           two-element vector, the first element is
    %                           the expected number of channels and the
    %                           second element is the expected number of
    %                           time samples. istftLayer errors if the
    %                           inverse short-time Fourier transform for
    %                           the given inputs does not match
    %                           ExpectedOutputSize in the number of
    %                           channels and samples. See the documentation
    %                           for details on the size of the ISTFT.
    %
    %   % Example 1: 
    %   %   Create an inverse short-time Fourier transform layer using a
    %   %   64-sample Hamming window. Specify 63 overlapped samples
    %   %   between adjoining segments.
    %
    %   layer = istftLayer(Window=hamming(64),OverlapLength=63);
    % 
    %    % Example 2:
    %    %   Create an array of 5 layers, containing a sequence input layer,
    %    %   an STFT layer, an LSTM layer, an ISTFT layer, and a regression
    %    %   layer. There is one feature in the sequence input. Set the minimum
    %    %   signal length in the sequence input layer to 2048 samples. Use the
    %    %   default window of length 128 for both STFT and ISTFT layers.
    %    %   Analyze the layers as a dlnetwork with a batch of 10 signals.
    %    
    %    layers = [sequenceInputLayer(1,MinLength=2048)
    %              stftLayer(TransformMode='realimag')
    %              lstmLayer(130)     
    %              fullyConnectedLayer(130)
    %              istftLayer
    %             ];
    % 
    %    analyzeNetwork(layers,dlarray(randn(1,10,2048,'single'),'CBT'),targetusage="dlnetwork");
    %
    %
    %   See also DLSTFT, STFT, STFTLAYER, DLISTFT, ISTFT. 

    % Copyright 2023 The MathWorks, Inc.
    
    properties (Learnable,Dependent)
        Weights
    end

    properties(Access=private,Hidden)
        pWeights
        pWindowLength     
        CachedInputSizes (1,1) signal.deep.internal.CacheISTFTSize
    end
    
    properties(Dependent)
        WeightLearnRateFactor
    end

    properties(SetAccess = private)
        % Layer properties.
        % Spectral window
        Window
        % Number of overlapped samples
        OverlapLength
        % Number of DFT points
        FFTLength
        % Method
        Method
        % OutputSize
        ExpectedOutputSize       
    end

    properties(Hidden,Dependent,Access=private)
        HopSize
        FreqSamples       
        CheckOutputSize
    end   

    methods
        function obj = istftLayer(NVargs)            
            arguments
                NVargs.Window {mustBeFloat,mustBeNonempty,mustBeFinite,...
                    mustBeReal,mustBeVector} = hann(128,'periodic');
                NVargs.FFTLength = []
                NVargs.OverlapLength = []
                NVargs.Name = ''
                NVargs.WeightLearnRateFactor (1,1) {mustBeNumeric} = 0                
                NVargs.Method {mustBeMember(NVargs.Method,{'wola','ola'})} = 'wola'
                NVargs.ExpectedOutputSize = 'none'
            end
            NVargs = nnet.internal.cnn.layer.util.gatherParametersToCPU(NVargs);
            obj.Type = 'ISTFT';
            obj.Name = NVargs.Name;
            obj.Window = NVargs.Window(:);
            if isdlarray(obj.Window)
                obj.Window = extractdata(obj.Window);
            end
            obj.OverlapLength = NVargs.OverlapLength;
            obj.pWindowLength = length(obj.Window);
            obj.Method = validatestring(NVargs.Method,{'wola','ola'});
            NVargs.ExpectedOutputSize;
            if isStringScalar(NVargs.ExpectedOutputSize) || ischar(NVargs.ExpectedOutputSize)
                obj.ExpectedOutputSize = validatestring(NVargs.ExpectedOutputSize,{'none'},'istftLayer','ExpectedOutputSize');
            end    
            if ~isStringScalar(obj.ExpectedOutputSize) && ~ischar(obj.ExpectedOutputSize)                     
                validateattributes(NVargs.ExpectedOutputSize,{'numeric'},{'vector',...
                    'real','finite','numel',2},'istftLayer','ExpectedOutputSize');
                obj.ExpectedOutputSize = NVargs.ExpectedOutputSize;                      
            end                     
            if isempty(NVargs.FFTLength)
                obj.FFTLength = length(obj.Window);
            else
                obj.FFTLength = NVargs.FFTLength;
            end
            if isempty(NVargs.OverlapLength)
                obj.OverlapLength = floor(0.75*length(obj.Window));
            end
            obj = validateProperties(obj);
            obj.Weights = single(obj.Window);   
            obj.WeightLearnRateFactor = NVargs.WeightLearnRateFactor;
            obj.CachedInputSizes = signal.deep.internal.CacheISTFTSize();
        end

        function xrec = predict(obj,s)    
            arguments
                obj
                s {mustBeReal}
            end
            % dims(s) must be 'CBT' or 'SCBT'
            if ~ismember(dims(s),{'CBT','SCBT'})
                error(message('signal_deep:stftLayer:NotCBTSCBT'));
            end  
            % Validate the size of the STFT is compatible with dlistft
            if startsWith(dims(s),'S')
                if size(s,1) ~= obj.FreqSamples
                    error((message('signal_deep:dlstft:InvalidFirstDimSCBT', ...
                    num2str(size(s,1)), num2str(obj.FreqSamples))));
                end
            %Only option left is CBT    
            elseif mod(size(s,1),obj.FreqSamples)
                error((message('signal_deep:dlstft:InvalidFirstDimCBT', ...
                    num2str(size(s,1)), num2str(obj.FreqSamples))));
            end        
            if isreal(s) && startsWith(dims(s),'S')
                if signalwavelet.internal.isodd(size(s,2))
                    error((message('signal_deep:dlstft:InvalidChannelSize', ...
                        num2str(size(s,2)))));
                end
            elseif isreal(s) && startsWith(dims(s),'C')
                if mod(size(s,1)/obj.FreqSamples/2,1)
                    error((message('signal_deep:dlstft:InvalidRealCBTFirstDim', ...
                        num2str(size(s,1)), num2str(obj.FreqSamples))));
                end
            end           
            if obj.CheckOutputSize
                inChannels = obj.ExpectedOutputSize(1);
                expectedTime = obj.ExpectedOutputSize(2);
                [match,numChannels,siglen] = signal.deep.internal.checkSTFTSize(s,obj.FreqSamples,...
                    obj.pWindowLength,obj.HopSize,inChannels,expectedTime);
                if ~match
                    error(message('signal_deep:dlstft:InvalidInputSize', ...
                    num2str(numChannels), num2str(expectedTime), num2str(inChannels),...
                    num2str(siglen)));
                end
            end            
            if strcmpi(dims(s),'CBT')
                xrec = signal.deep.internal.istftCBTInput(real(s),...
                    real(obj.Weights),obj.FFTLength,obj.HopSize,obj.Method);
            else
                xrec = signal.deep.internal.istftSCBTInput(real(s),...
                    real(obj.Weights),obj.FFTLength,obj.HopSize,obj.Method);
            end          
        end        
        
        function obj = set.Weights(obj,weights)
            validateattributes(weights,{'double','single'},{'vector','real',...
                'finite'},'istftLayer','Weights');
            [rowsize,colsize] = size(weights);
            weights = weights(:);            
            if length(weights) ~= obj.pWindowLength         
                error(message('signal_deep:stftLayer:WeightSize',...
                    obj.pWindowLength,rowsize,colsize));
            end      
            obj.pWeights = single(weights);            
        end

        function obj = set.FreqSamples(obj, nfreq)
            validateattributes(nfreq,{'numeric'},...
                {'scalar','integer','nonnegative','real','nonnan','finite',...
                'nonempty','>',1,'<=',obj.FFTLength},'istftLayer','FreqSamples'); 
            obj.FreqSamples = nfreq;
        end

        function obj = set.HopSize(obj,hopsize)
            validateattributes(hopsize,{'numeric'},...
                {'scalar','integer','nonnegative','real','nonnan','finite',...
                'nonempty','>',0,'<=',obj.pWindowLength},'istftLayer','HopSize');  
            obj.HopSize = hopsize;
        end

        function val = get.Weights(obj)
            val = obj.pWeights;
        end
        
        function val = get.HopSize(obj)
            val = length(obj.Window)-obj.OverlapLength;
        end

        function val = get.FreqSamples(obj)
            val = fix(obj.FFTLength/2) +1;
        end    

        function val = get.CheckOutputSize(obj)
            if ~isempty(obj.ExpectedOutputSize) && isvector(obj.ExpectedOutputSize)...
                    && isnumeric(obj.ExpectedOutputSize)
                val = true;
            else
                val = false;
            end
        end
        
        function this = set.WeightLearnRateFactor(this,lrate)
            this = setLearnRateFactor(this,"Weights",lrate);
        end
        
        function lrate = get.WeightLearnRateFactor(this)
            lrate = getLearnRateFactor(this,"Weights");
        end        
    end

    methods (Hidden)        
        function obj = validateProperties(obj)      
            % Window
            nwin = length(obj.Window);
            try
                validateattributes(nwin,{'numeric'},...
                    {'integer','>',1},'istftLayer');
            catch e
                error(message('signal_deep:stftLayer:windowlength'));
            end
            % OverlapLength
            validateattributes(obj.OverlapLength,{'numeric'},{'scalar',...
                'integer','nonnegative','real','nonempty','>=',0,'<',nwin},...
                'istftLayer','OverlapLength');
            
            % FFTLength
            validateattributes(obj.FFTLength,{'numeric'},...
                {'scalar','integer','nonnegative','real','nonnan',...
                'finite','nonempty','>=',nwin},'istftLayer','FFTLength');  
        end

        function Z = cpuInferenceEnginePredict(obj, X)
            % Prediction method for the CPU inference engine.
            hasSpatialDimension = numel(obj.CachedInputSizes.Value) > 3;
            if hasSpatialDimension
                X = dlarray(X,'SCBT');
            else
                X = dlarray(X,'CBT');
            end
            Z = obj.predict(X);
            Z = extractdata(Z);
        end
        
        function layerArgs = getFusedArguments(obj)
            % getFusedArguments  Returned the arguments needed to call the
            % layer in a fused network.
            % 1,1 refers to 1 input, 1 output argument
            layerArgs = { 'external', @obj.cpuInferenceEnginePredict, ...
                1, 1, @obj.propagateSize};
        end
        
        function tf = isFusable(~)
            % isFusable  Indicates if the layer is fusable in a given network.
            tf = true;
        end

        function szout = propagateSize(obj,szIn) 
            % This method calculates the output size of istftLayer. This is
            % needed for inference engine.       
            obj.CachedInputSizes.fillCache(szIn);
            if numel(szIn) == 4
                nseg = szIn(4);
                nchan = szIn(2);
                channelSize = nchan/2;
                nbatch = szIn(3);                
            else
                nseg = szIn(3);
                nchan = szIn(1);
                channelSize = nchan/(2*obj.FreqSamples);
                nbatch = szIn(2);
            end
            Tdim = obj.pWindowLength+(nseg-1)*obj.HopSize;                    
            % Canonical form: CBT
            szout = [channelSize nbatch Tdim];
        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
    end
end

function mustBeMember(A,B)
    validatestring(A,B,'isftLayer','Method');
end
