classdef modwtLayer < nnet.layer.Layer & nnet.layer.Formattable & ...
        nnet.internal.cnn.layer.CPUFusableLayer
% modwtLayer Maximal overlap discrete wavelet transform layer
%   layer = modwtLayer creates a maximal overlap discrete wavelet
%   transform (MODWT) layer. By default, the layer computes the
%   MODWT based multiresolution analysis (MODWTMRA) to level 5 using the
%   'sym4' Daubechies wavelet. 
%
%   Tensor inputs to modwtLayer must be CBT arrays. The size of the 'T'
%   dimension of the tensor input must be greater than or equal to 2^Level,
%   where Level is the transform level of the MODWT. 
%
%   The output of modwtLayer is a dlarray formatted as 'SCBT'. The spatial
%   dimension corresponds to the level of the wavelet transform. The k-th
%   row, where k<=Level, contains the k-th level detail of the signal. The
%   (Level+1)-th row contains the Level-th level smooth of the signal.
%
%   The learnable parameter 'Weights' in modwtLayer is a two-row matrix of
%   the current filter pair. The first row is the lowpass filter and the
%   second row is the highpass filter. By default, the weights are the
%   lowpass and highpass filters associated with the default wavelet and do
%   not update. It is not recommended to initialize the weights directly.
%   
%   layer = modwtLayer('PARAM1',VAL1,'PARAM2',VAL2,...) specifies optional
%   parameters for the layer using name-value arguments:
%
%   'Wavelet'               Name of an orthogonal wavelet used in the
%                           MODWT, specified as a character vector or a
%                           string scalar. Orthogonal wavelets are 
%                           designated as type 1 wavelets in the wavelet 
%                           manager. Valid built-in orthogonal wavelet
%                           families begin with 'haar', 'dbN', 'fkN',
%                           'coifN', 'blN', 'hanSR.LP', 'symN', 'vaid', or
%                           'beyl'. Use waveinfo with the wavelet family
%                           short name to see supported values for any
%                           numeric suffixes and how to interpret those
%                           values. 
%                           
%                           For example:
%                           waveinfo('db'). 
%                           
%                           For a wavelet specified by wname, the lowpass 
%                           filter Lo and highpass filter Hi are
%                           [~,~,Lo,Hi] = wfilters(wname)
%
%                           If unspecified, and the filter pair
%                           'LowpassFilter', 'HighpassFilter' are also not 
%                           specified, 'Wavelet' defaults to 'sym4',
%                           otherwise 'Wavelet' defaults to ''.
%          
%   'LowpassFilter',        Initial Lowpass (scaling) filter and highpass                         
%   'HighpassFilter'        (wavelet) filter, specified as real-valued   
%                           even-length numeric vectors.
%                         
%                           In order to satisfy the requirements of the
%                           MODWT, the filter pair must correspond to an
%                           orthogonal wavelet. You can use ISORTHWFB to
%                           determine orthogonality. 
%                           
%                           For example:
%                           [~,~,Lo,Hi] = wfilters('db2');
%                           [tf,checks] = isorthwfb(Lo,Hi);
%
%                           You cannot specify both a wavelet name and  
%                           a wavelet filter pair.  
%
%                           If unspecified, both filters default to [] and
%                           the layer uses the 'Wavelet' value to determine
%                           the filters used in MODWT.
%
%   'Level'                 Transform level to compute the MODWT, specified
%                           as a positive integer less than or equal to
%                           floor(log2(N)), where N is the size of layer
%                           input along 'T' dimension.
%                           The default is 5.
%                  
%   'SelectedLevels'        Selected levels for layer to output, specified
%                           as a vector of positive integers less than or
%                           equal to Level. 
%                           The default is 1:Level.
%
%   'IncludeLowpass'        Logical scalar to include Level-th level
%                           lowpass (scaling) coefficients in the MODWT or
%                           the Level-th level smooth in MODWTMRA.            
%                           The default is true.
%
%   'AggregateLevels'       Logical scalar to aggregate the selected levels 
%                           and lowpass level (if included) of each input  
%                           channel by summation. If this is set to true,
%                           the size of the output along the 'S' dimension 
%                           is 1. 
%                           The default is false.
%
%   'Algorithm'             The algorithm used to compute the output,
%                           specified as one of the following:
%
%                           'MODWTMRA' - Compute the maximal overlap
%                           discrete wavelet transform based 
%                           multiresolution analysis.                           
%                           (default)
%                           
%                           'MODWT' - Compute the wavelet coefficients of
%                           the maximal overlap discrete wavelet transform.
%
%   'Boundary'              Boundary condition to use in the computation
%                           of the MODWT, specified as one of these:
%
%                           'periodic' - the MODWT is computed  without
%                           making changes to the original length of input
%                           signal. 
%                           (default)
%
%                           'reflection' - the MODWT is computed using a
%                           reflected signal along the 'T' dimension. The 
%                           wavelet coefficients of MODWT will have twice
%                           the original length of the input signal.
%
%   '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:
%     % Create a maximal overlap discrete wavelet transform layer to
%     % compute the multiresolution analysis result for the input signal.
%     % Use a coiflet wavelet with order 5. Set the transform level to 8.
%     % Only keep the details at levels 3, 5, and 7, and the approximation.
%     % Then create a dlnetwork containing a sequence input layer, a MODWT
%     % layer, and a LSTM layer. For a level-8 decomposition, set the 
%     % minimum sequence length to 2^8. To work with an LSTM layer, a 
%     % a flatten layer is also needed before the LSTM layer to collapse
%     % the spatial dimension into the channel dimension. Run a batch
%     % of 10 random single-channel signals through the dlnetwork.
%
%     layer = modwtLayer('Wavelet','coif5','Level',8, ...
%                   'SelectedLevels',[3,5,7],'Name','MODWT');
%     MinLength = 2^layer.Level; 
%     layers = [sequenceInputLayer(1,'Name','input','MinLength',MinLength)
%            layer
%            flattenLayer
%            lstmLayer(10,'Name','LSTM')
%           ];
%     dlnet = dlnetwork(layers);
%     dataout = dlnet.forward(dlarray(randn(1,10,2000,'single'),'CBT'));
%     size(dataout)
%     dims(dataout)
%
%   See also DLMODWT, MODWT, MODWTMRA, CWTLAYER, DLCWT

%   Copyright 2022-2023 The MathWorks, Inc. 
   
    
    properties (Learnable,Dependent)
        % learnable parameters (with the same learning rate)
        Weights
    end

    properties(Access=private,Hidden)
        pWeights
    end
    
    properties(Dependent)
        WeightLearnRateFactor
    end
    
    properties(SetAccess = private)
        % Layer properties.
        % Input filter pair
        LowpassFilter
        HighpassFilter
        % Input wavelet name
        Wavelet
        % The level for MODWT to compute
        Level
        % Subset of levels selected for output
        SelectedLevels
        % If aggregate levels or not
        AggregateLevels
        % If include lowpass component or not
        IncludeLowpass
        % Algorithm (MRA/MODWT)
        Algorithm
        % Boudary Condition
        Boundary
    end

    
    properties(SetAccess = private,Hidden)

        % Minimum input length
        MinLength

        % The output componets of modwt/mra
        OutputComponets

        % The actual level used to compute
        ComputeLevel

        % Internal filter pair used for computation
        InternalLowpassFilter
        InternalHighpassFilter
    end
    
    methods
        function this = modwtLayer(NVargs)
            arguments
                NVargs.LowpassFilter  {modwtLayer.mustBeVectorOrEmpty,mustBeNumeric} = [] 
                NVargs.HighpassFilter  {modwtLayer.mustBeVectorOrEmpty,mustBeNumeric} = [] 
                NVargs.Wavelet {mustBeTextScalar} = '' 
                NVargs.Level {mustBePositive,mustBeInteger,mustBeScalarOrEmpty,mustBeNonempty} = 5
                NVargs.SelectedLevels  {modwtLayer.mustBeVectorOrEmpty,mustBePositive,mustBeInteger} = []
                NVargs.IncludeLowpass  {mustBeNumericOrLogical} = true
                NVargs.AggregateLevels  {mustBeNumericOrLogical} = false
                NVargs.Boundary {mustBeTextScalar} = 'periodic'
                NVargs.Algorithm {mustBeTextScalar}= 'MODWTMRA'
                NVargs.WeightLearnRateFactor {mustBeNonnegative,mustBeScalarOrEmpty,mustBeNonempty} = 0
                NVargs.Name {mustBeTextScalar} = ''
            end
            NVargs = nnet.internal.cnn.layer.util.gatherParametersToCPU(NVargs);
            
            % Layer information
            this.Version = 2;
            this.Name = NVargs.Name;
            this.Type = 'MODWT';        
            
            % Layer parameters
            this.Level = NVargs.Level;
            this.AggregateLevels = NVargs.AggregateLevels;
            this.IncludeLowpass = NVargs.IncludeLowpass;
            this.Algorithm = NVargs.Algorithm;
            this.Boundary = NVargs.Boundary;
            this.WeightLearnRateFactor = NVargs.WeightLearnRateFactor;
            
            this.Wavelet = NVargs.Wavelet;
            this.LowpassFilter = NVargs.LowpassFilter;
            this.HighpassFilter = NVargs.HighpassFilter;
            if isdlarray(this.LowpassFilter)
                this.LowpassFilter = extractdata(this.LowpassFilter);
            end
            if isdlarray(this.HighpassFilter)
                this.HighpassFilter = extractdata(this.HighpassFilter);
            end

            if ~isempty(this.LowpassFilter)
                this.LowpassFilter = this.LowpassFilter(:)';
            end
            if ~isempty(this.HighpassFilter)
                this.HighpassFilter = this.HighpassFilter(:)';
            end
            
            % Initialize the selected output levels
            if isempty(NVargs.SelectedLevels)
                this.SelectedLevels = 1:NVargs.Level;
            else
                this.SelectedLevels = NVargs.SelectedLevels; 
                % Validate SelectedLevels.<= Level
                if ~all(this.SelectedLevels<=NVargs.Level)
                    error(message('wavelet_deep:modwtLayer:SelectedLevelOutofRange',NVargs.Level))
                end
            end
            this.SelectedLevels = this.SelectedLevels(:)';

            % Set output decomposition componets
            % Reduce computation level if no need lowpass component
            if this.IncludeLowpass
                this.ComputeLevel = this.Level;
                this.OutputComponets = [this.SelectedLevels,NVargs.Level+1];
            else
                this.ComputeLevel = max(this.SelectedLevels);
                this.OutputComponets = this.SelectedLevels;
            end
            
            % Initialize the input length requirement information
            this.MinLength = 2^this.Level;
            % Add the description
            this.Description = "Custom MODWT Layer with mininum input length " + this.MinLength;   

            % Intialize the filter pair
            % When both Lo and Hi are specified
            if ~isempty(NVargs.LowpassFilter) && ~isempty(NVargs.HighpassFilter)
                if ~isempty(NVargs.Wavelet)
                    error(message('wavelet_deep:modwtLayer:InvalidBothWaveletFilters'));
                end
                this.InternalLowpassFilter = this.LowpassFilter;
                this.InternalHighpassFilter = this.HighpassFilter;
            % When only one of Lo/Hi is specified
            elseif isempty(NVargs.LowpassFilter) && ~isempty(NVargs.HighpassFilter)
                error(message('wavelet_deep:modwtLayer:BothFilters','lowpass'));
            elseif ~isempty(NVargs.LowpassFilter) && isempty(NVargs.HighpassFilter)
                error(message('wavelet_deep:modwtLayer:BothFilters','highpass'));
            % When none of Lo/Hi is specified
            else
                if isempty(NVargs.Wavelet)
                    this.Wavelet = 'sym4';                   
                end            
                try
                    wavtype = wavemngr('type',this.Wavelet);     
                    [~,~,this.InternalLowpassFilter,this.InternalHighpassFilter] = wfilters(this.Wavelet);
                catch ME
                    error(message('wavelet_deep:modwtLayer:InvalidWaveletFilter',"'"+this.Wavelet+"'"));
                end
                if wavtype ~= 1
                    error(message('wavelet_deep:modwtLayer:NotOrthogonal',"'"+this.Wavelet+"'"));               
                end
            end

            % Validate properties can not be done in argument block
            this = validateProperties(this);
            % Initialize learnable parameters, Weights
            wghts = initializeWeights(this);
            this.Weights = wghts;            
        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:modwtLayer:InvalidDataFormat'));
            end

            if size(X,3) < this.MinLength
                error(message('wavelet_deep:modwtLayer:TooShortInput',this.MinLength))
            end
            
            % Use real() to bypass complex gradient issue
            if ~isequal(this.Algorithm,'MODWT')        
                [~,Z] = dlmodwt(real(X), ...
                    real(this.Weights(1,:)),real(this.Weights(2,:)), ...
                    this.Level,Boundary=this.Boundary);
            else
                [Z,~] = dlmodwt(real(X), ...
                    real(this.Weights(1,:)),real(this.Weights(2,:)), ...
                    this.Level,Boundary=this.Boundary);
            end              

            Z = Z(this.OutputComponets,:,:,:);
             
            if this.AggregateLevels
                Z = sum(Z,1);
            end
        end

        function this = set.Weights(this,wghts)
            validateattributes(wghts,{'double','single'},...
                {'real','2d','finite','nonempty'},'modwtLayer','Weights');
            dtype = underlyingType(wghts);
            [rowSize,colSize] = size(wghts);
            expColSize = length(this.InternalLowpassFilter); 
            if rowSize ~= 2 || colSize ~= expColSize
                error(message('wavelet_deep:modwtLayer:WeightSize',...
                    expColSize,rowSize,colSize));
            end
            if ~startsWith(dtype,'s')
                wghts = cast(wghts,'single');
            end
            this.pWeights = wghts;
        end

        function val = get.Weights(this)
            val = this.pWeights;
        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 this = validateProperties(this)
            % Validate NVPs
            this.Algorithm = ...
                validatestring(this.Algorithm,...
                {'MODWTMRA','MODWT'},'modwtLayer','Algorithm');
            this.Boundary = ...
                validatestring(this.Boundary,...
                {'periodic','reflection'},'modwtLayer','Boundary');
            
            % Validate filters
            % Any filter pair with equal-length and even-length is valid 
            LenLo = numel(this.LowpassFilter);
            LenHi = numel(this.HighpassFilter);
            tfEven = ...
                ~mod(LenLo,2) && ~mod(LenHi,2);
            tfEqualLen = LenLo == LenHi;
            if  ~tfEven || ~tfEqualLen
               error(message('wavelet_deep:modwtLayer:InvalidFilters'))
            end
        end

        function wghts = initializeWeights(this)
            wghts =  cast([this.InternalLowpassFilter;this.InternalHighpassFilter],'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 modwtLayer. This is
            % needed for inference engine.

            % For CBT output (temporal), we need to collapse the spatial
            % (level) dimension into the channel dimension.
            % For 'spatial', 2D convolutional networks convolve along S
            % dimensions, so we have to perumte the array SCBT becomes SSCB
            if strcmp(this.Boundary,'reflection') && strcmp(this.Algorithm,'MODWT')
                Tdim = szIn(3) * 2;
            else
                Tdim = szIn(3);
            end

            if this.AggregateLevels
                Sdim = 1; 
            else
                Sdim = length(this.OutputComponets);
            end
            % Canonical form: SCBT
            szout = [Sdim szIn(1) szIn(2) 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

       function mustBeVectorOrEmpty(arg)
            if ~isempty(arg) && ~isvector(arg)
                error(message('wavelet_deep:modwtLayer:NotVectorNorEmpty'))
            end
        end
    end

end


