classdef cwtLayer < nnet.layer.Layer & nnet.layer.Formattable & ...
        nnet.internal.cnn.layer.CPUFusableLayer & nnet.layer.Acceleratable        
    % cwtLayer Continuous wavelet transform layer
    %   layer = cwtLayer creates a continuous wavelet transform (CWT) layer
    %   for a signal length of 1024 samples. The layer uses 10 wavelet
    %   filters per octave and periodic boundary conditions. By default,
    %   the Morse (3,60) wavelet is used.
    %
    %   Tensor inputs to cwtLayer must be CBT arrays. The size of the 'T'
    %   dimension of the tensor input is padded to equal the value of the
    %   layer 'SignalLength' property.
    %
    %   By default, the output of cwtLayer is an SCBT array. The spatial
    %   dimension corresponds to scale, or frequency.
    %
    %   The weights in cwtLayer are initialized internally to be the
    %   wavelet filters used in the CWT. It is not recommended to
    %   initialize the weights directly.
    %
    %   layer = cwtLayer('PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
    %   parameters for the layer using name-value arguments:
    %
    %
    %   CWTLAYER properties:
    %
    %   'SignalLength'          Signal length in samples, specified as a
    %                           positive integer greater than or equal to 4.
    %                           All sequence inputs to cwtLayer are padded to
    %                           have size 'SignalLength' along the 'T'
    %                           dimension.
    %
    %   'Wavelet'               Wavelet used in the CWT, specified as a
    %                           scalar string or character array. Valid
    %                           options for 'Wavelet' are "Morse", "amor",
    %                           or "bump". If unspecified, 'Wavelet'
    %                           defaults to "Morse".
    %
    %   'Threshold'             Weight threshold, specified as a positive
    %                           real-valued scalar. The threshold value is
    %                           used to determine the significant values
    %                           for each of the CWT filters in the wavelet
    %                           filter bank prior to any weight
    %                           modification through learning. Values below
    %                           the specified threshold are set to zero and
    %                           excluded from learning. The CWT filters are
    %                           normalized so that the peak value is 2 for
    %                           each filter. If unspecified, the threshold
    %                           value defaults to 1e-8. Smaller threshold
    %                           values result in more values being retained
    %                           from the CWT filters and therefore less
    %                           weight reduction. Larger threshold values
    %                           result in more weight reduction and more
    %                           divergence between the deep learning CWT
    %                           and the transforms computed with the full
    %                           filter bank. A threshold less than realmin,
    %                           the smallest positive normalized
    %                           floating-point number in double precision,
    %                           is clipped to realmin for computing
    %                           significant filter values. Setting the
    %                           threshold to a value which results in no
    %                           values being retained for any individual
    %                           filter results in an error.
    %
    %   'IncludeLowpass'        Logical scalar to include lowpass (scaling)
    %                           filter in the CWT. 'IncludeLowpass' defaults
    %                           to false and only wavelet filters are
    %                           included.
    %
    %   'VoicesPerOctave'       Number of wavelets per octave in the CWT,
    %                           specified as a positive integer between 1
    %                           and 48. If unspecified, 'VoicesPerOctave'
    %                           defaults to 10.
    %
    %   'FrequencyLimits'       Frequency limits for the CWT, specified as
    %                           a two-element vector with positive strictly
    %                           increasing elements. The elements in
    %                           'FrequencyLimits' are interpreted as
    %                           normalized frequencies, cycles/sample. To
    %                           obtain normalized frequencies, divide your
    %                           desired frequency limits in hertz by the
    %                           sample rate in hertz. For example, if the
    %                           sample rate is 1000 Hz and your desired
    %                           frequency limits are [100,400] Hz, divide
    %                           each element by 1000 to obtain the
    %                           normalized frequencies: [100/1000,400/1000]. 
    %                           Depending on how many wavelets per octave
    %                           your layer contains, there is a minimum
    %                           required separation between the elements of
    %                           'FrequencyLimits'. See the documentation of
    %                           cwtLayer or cwtfilterbank for details. Note
    %                           that frequency limits only apply to wavelet
    %                           filters. If you additionally specify
    %                           'IncludeLowpass' as true, cwtLayer also
    %                           includes the lowpass (scaling) filter.
    %
    %   'TimeBandwidth'         Time-bandwidth product for Morse wavelets,
    %                           specified as a positive number greater than
    %                           equal to 3 and less than or equal to 120. 
    %                           If unspecified, the time-bandwidth parameter 
    %                           defaults to 60. This property is only used
    %                           when the 'Wavelet' property is "Morse". The 
    %                           time-bandwidth product is ignored for the 
    %                           "amor" and "bump" wavelets. For Morse
    %                           wavelets, the larger the time-bandwidth
    %                           parameter, the more spread out the Morse
    %                           wavelet is in time and narrower the wavelet
    %                           is in frequency.
    %
    %   'TransformMode'         Layer transform mode, specified as one of
    %                           these:
    %
    %                           'mag' - CWT magnitude (default)
    %
    %                           'squaremag' - CWT squared magnitude
    %
    %                           'realimag' - CWT real and imaginary parts
    %                           concatenated along the channel dimension
    %
    %   'WeightLearnRateFactor' Multiplier for the learning rate of the
    %                           weights, specified as a nonnegative scalar.
    %                           The weights are the reduced CWT filter
    %                           values represented as a 1-by-1-by-Nr tensor.
    %                           See CWTFILTERS2ARRAY for details.
    %                           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 ''.
    %
    %   CWTLAYER methods:
    %
    %   filterbank              Full-weight CWT filter bank
    %
    %
    %   %Example 1:
    %   %   Create a three-layer dlnetwork containing a sequence input layer,
    %   %   a CWT layer, and a 2-D max pooling layer. Run a batch of 
    %   %   10 random single-channel signals through the dlnetwork. 
    %   layers = [sequenceInputLayer(1,'Name','input','MinLength',2e3)
    %         cwtLayer('SignalLength',2e3,'WeightLearnRateFactor',1)
    %          maxPooling2dLayer([2 25],'Stride',[2 12])
    %           ];
    %   dlnet = dlnetwork(layers);
    %   dataout = dlnet.forward(dlarray(randn(1,10,2000,'single'),'CBT'));
    %   size(dataout)
    %   dims(dataout)
    %
    %   %Example 2:
    %   %   Construct a default CWT layer. Obtain and plot the filter bank.
    %   layer = cwtLayer;
    %   psif = filterbank(layer);
    %   f = 0:1/layer.SignalLength:1-1/layer.SignalLength;
    %   plot(f,psif')
    %   xlim([0 1/2])
    %   xlabel('Cycles/Sample')
    %   ylabel('Magnitude')
    %
    %
    %   See also DLCWT, CWTFILTERS2ARRAY, CWTFILTERBANK, CWT, MODWTLAYER, DLMODWT

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

    properties(Access=private,Hidden)
        pWeights
    end

    properties(Dependent)
        WeightLearnRateFactor
    end

    properties(Access=private,Dependent)
        pFrequencyLimits
    end

    properties
        % Layer output type
        TransformMode      
    end

    properties(SetAccess=private)
        SignalLength
        IncludeLowpass
        Threshold = 1e-8
        VoicesPerOctave
        Wavelet
        TimeBandwidth
        FrequencyLimits
    end

    properties(SetAccess=private,Hidden)
        fbcell = {'Boundary','periodic'}
        FilterIndices
        SamplingFrequency = 1
        NormFreqLimits
        TensorSize = []
    end

    methods
        function this = cwtLayer(NVargs)
            arguments
                NVargs.TransformMode = "mag"
                NVargs.SignalLength {mustBeNumeric} = 1024
                NVargs.WeightLearnRateFactor {mustBeNonnegative,mustBeNumeric} = 0
                NVargs.Name = ''
                NVargs.IncludeLowpass{mustBeNumericOrLogical} = false
                NVargs.Threshold = 1e-8
                NVargs.VoicesPerOctave = 10
                NVargs.Wavelet = "Morse"
                NVargs.TimeBandwidth = 60
                NVargs.FrequencyLimits = []
            end
            NVargs = nnet.internal.cnn.layer.util.gatherParametersToCPU(NVargs);
            
            this.Version = 2;
            this.Type = 'CWT';
            this.SignalLength = NVargs.SignalLength;
            this.TransformMode = NVargs.TransformMode;
            this.Name = NVargs.Name;
            this.IncludeLowpass = NVargs.IncludeLowpass;
            this.Threshold = NVargs.Threshold;
            this.VoicesPerOctave = NVargs.VoicesPerOctave;
            this.Wavelet = validatestring(NVargs.Wavelet,["Morse","amor","bump"],...
                'cwtLayer','Wavelet');
            this.FrequencyLimits = NVargs.FrequencyLimits;
            this.TimeBandwidth = NVargs.TimeBandwidth;
            this = validateProperties(this);
            [filters,filtidx] = initializeWeights(this);
            this.TensorSize = size(filters);
            this.FilterIndices = filtidx;
            % Here property set method is called
            this.Weights = filters;            
            this.WeightLearnRateFactor = NVargs.WeightLearnRateFactor;
        end

        function Z = predict(this, X)
            % Forward input data through the layer at prediction time and
            % output the result.
            arguments
                this
                X {mustBeReal}
            end
            if ~isequal(dims(X),'CBT')
                error(message('wavelet_deep:cwtLayer:NotCBT'));
            end
            cfs = dlcwt(X,real(this.Weights),this.FilterIndices);
            Z = iFormatOutput(this,cfs);            
        end

        function this = set.Weights(this,filters)
            validateattributes(filters,{'double','single'},...
                {'real','finite','nonempty'},'cwtLayer','Weights');
                dtype = underlyingType(filters);
                [rowSize,colSize,pageSize] = size(filters); 
                if rowSize ~= this.TensorSize(1) || ...
                    colSize ~= this.TensorSize(2) || ...
                    pageSize ~= this.TensorSize(3)  
                    error(message('wavelet_deep:cwtLayer:WeightSize',...
                        this.TensorSize(3),rowSize,colSize,pageSize));
                end            
            if ~startsWith(dtype,'s')
                filters = cast(filters,'single');
            end
            this.pWeights = filters;                  
        end

        function val = get.Weights(this)
            val = this.pWeights;
        end

        function flim = get.pFrequencyLimits(this)
            if strcmpi(this.Wavelet,"morse")            
                [minfreq,maxfreq] = cwtfreqbounds(this.SignalLength,...
                        this.SamplingFrequency,'Wavelet',this.Wavelet,...
                        'TimeBandwidth',this.TimeBandwidth,...
                        'VoicesPerOctave',this.VoicesPerOctave);
            else
                [minfreq,maxfreq] = cwtfreqbounds(this.SignalLength,...
                        this.SamplingFrequency,'Wavelet',this.Wavelet,...
                        'VoicesPerOctave',this.VoicesPerOctave);
            end     
            flim = [minfreq, maxfreq];
        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.TransformMode(this,transformMode)
            this.TransformMode = transformMode;
        end

        function this = set.Threshold(this,threshold)
            this.Threshold = threshold;
        end

        function fidx = get.FilterIndices(this)
            fidx = this.FilterIndices;
        end

        function psif = filterbank(this)
            % Full-weight CWT filterbank
            %   PSIF = FILTERBANK(CLAYER) returns the full-weight CWT filter
            %   bank for the cwtLayer, CLAYER.
            %
            %   %Example
            %   %   Return full-weight CWT filter bank for cwtLayer.
            %   clayer = cwtLayer('SignalLength',2048,'Wavelet','amor');
            %   psif = filterbank(clayer);
            %
            %   See also CWTFILTERBANK, DLCWT, CWTFILTERS2ARRAY
            psif = ...
                array2cwtfilters(this.Weights,this.FilterIndices);
        end
    end

    methods (Hidden)
        function this = validateProperties(this)
            validateattributes(this.SignalLength,{'numeric'},...
                {'integer','scalar','>=',4},'cwtLayer','SignalLength');
            validateattributes(this.IncludeLowpass,{'numeric','logical'},...
                {'scalar','finite','nonempty','real'},'cwtLayer','IncludeLowpass');
            this.fbcell{end+1} = 'SignalLength';
            this.fbcell{end+1} = this.SignalLength;
            this.fbcell{end+1} = 'Wavelet';
            this.fbcell{end+1} = this.Wavelet;
            validateattributes(this.Threshold,{'numeric'},...
                {'real','nonempty','scalar'},'cwtLayer','Threshold');
            validateattributes(this.VoicesPerOctave,{'numeric'},{'real','scalar',...
                '>=',1,'<=',48},'cwtLayer','VoicesPerOctave');
            this.fbcell{end+1} = 'VoicesPerOctave';
            this.fbcell{end+1} = this.VoicesPerOctave;
            this.TransformMode = validatestring(this.TransformMode,...
                ["mag","squaremag","realimag"],'cwtLayer','TransformMode');

            if ~isempty(this.TimeBandwidth) && ~startsWith(this.Wavelet,"M")
                this.TimeBandwidth = [];               
            elseif ~isempty(this.TimeBandwidth)
                validateattributes(this.TimeBandwidth,{'numeric'},...
                    {'real','scalar','>=',3,'<=',120},'cwtLayer','TimeBandwidth');
                this.fbcell{end+1} = 'TimeBandwidth';
                this.fbcell{end+1} = this.TimeBandwidth;
            end

            if  ~isempty(this.FrequencyLimits)
                validateattributes(this.FrequencyLimits,{'numeric'},...
                    {'numel',2,'increasing','>=',0,'<=',this.SamplingFrequency/2},...
                    'cwtLayer','FrequencyLimits');
                this.fbcell{end+1} = 'SamplingFrequency';
                this.fbcell{end+1} = this.SamplingFrequency;
            else
                this.FrequencyLimits = this.pFrequencyLimits;
            end
            this.fbcell{end+1} = 'FrequencyLimits';
            this.fbcell{end+1} = this.FrequencyLimits;
        end


        function Z = iFormatOutput(this,CFS)
            switch this.TransformMode
                case 'mag'
                    Z = abs(CFS);
                case 'squaremag'
                    Z = abs(CFS).^2;
                case 'realimag'
                    dimC = finddim(CFS,'C');
                    Z = cat(dimC,real(CFS),imag(CFS));
            end
        end     

        function [filters,filtidx] = initializeWeights(this)
            fb = cwtfilterbank(this.fbcell{:});
            [psifV,filtidx] = cwtfilters2array(fb,this.Threshold,...
                IncludeLowpass = this.IncludeLowpass);            
            filters = cast(real(psifV),'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 cwtLayer. This is
            % needed for inference engine.
            if strcmp(this.TransformMode,'realimag')
                szIn(1) = szIn(1)+szIn(1);
            end
            nT = double(this.FilterIndices(1,2));
            Sdim = size(this.FilterIndices,1)-1;
            szout = [Sdim szIn(1) szIn(2) nT];            
        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