function  [X,T,INFO] = cwtmag2sig(Smag,timeinfo,cwtNvargs,Nvargs)
%CWTMAG2SIG Signal reconstruction from CWT magnitude
%   X = CWTMAG2SIG(CFS) returns a time-domain real-valued signal X
%   estimated from the continuous wavelet transform (CWT) magnitude, using
%   the limited-memory BFGS (L-BFGS) optimizer. For single-channel signals,
%   CFS is a 2-D matrix, with time increasing across the columns and
%   frequency decreasing down the rows. For multichannel signals, CFS is
%   a 3-D array, where the third dimension corresponds to the channels. X
%   is a 2-D matrix, where each column corresponds to one channel. By
%   default, CWTMAG2SIG uses the Morse (3,60) wavelet with default
%   frequency range. CWTMAG2SIG assumes the CWT magnitude was obtained from
%   a real-valued signal.
%
%   CWTMAG2SIG assumes the "ExtendSignal" of the original CWT is false or
%   the Boundary of the cwtfilterbank object used in the original CWT is
%   "periodic". For optimal reconstruction results, the following CWT
%   parameters must be consistent with those used to obtain the CWT
%   magnitude. Also, retain and reuse the scaling coefficients obtained
%   from the original CWT.
%
%   CWT Parameters:
%
%   X = CWTMAG2SIG(CFS,Fs) specifies the sample rate in hertz as a positive
%   scalar.
%
%   X = CWTMAG2SIG(CFS,Ts) specifies Ts as a positive scalar duration
%   corresponding to the sampling period of X. The sample rate is
%   calculated as 1/Ts.
%
%   X = CWTMAG2SIG(...,FrequencyLimits=FLIMITS) specifies the frequency
%   limits used for reconstruction. FLIMITS is a positive two-element
%   vector. The elements of FLIMITS must be strictly increasing. The first
%   element of FLIMITS specifies the lowest peak passband frequency and
%   must be greater than or equal to the product of the wavelet peak
%   frequency in hertz and two time standard deviations divided by the
%   length of a channel. The second element of FLIMITS must be less than or
%   equal to the Nyquist frequency. The base 2 logarithm of the ratio of
%   the maximum frequency to the minimum frequency must be greater than or
%   equal to 1/NV, where NV is the number of voices per octave. If you
%   specify FLIMITS outside the permissible range, CWTMAG2SIG truncates the
%   limits to the minimum and maximum valid values.
%
%   X = CWTMAG2SIG(...,PeriodLimits=PLIMITS) specifies the period limits
%   used for reconstruction.  PLIMITS is a positive two-element duration
%   array. The elements of PLIMITS must be strictly increasing. The first
%   element of PLIMITS must be greater than or equal to 2*Ts, where Ts is
%   the sampling period. The base 2 logarithm of the ratio of the minimum
%   period to the maximum period must be less than or equal to -1/NV, where
%   NV is the number of voices per octave. The maximum period cannot exceed
%   the length of the channel divided by the product of two time standard
%   deviations of the wavelet and the wavelet peak frequency. If you
%   specify PLIMITS outside the permissible range, CWTMAG2SIG truncates the
%   limits to the minimum and maximum valid values.
%
%   X = CWTMAG2SIG(...,Wavelet=WNAME) uses the wavelet WNAME for the
%   reconstruction. Valid options for WNAME are: "morse", "amor", or
%   "bump". WNAME defaults to "morse".
%
%   X = CWTMAG2SIG(...,VoicesPerOctave=NV) specifies the number of voices
%   per octave used for reconstruction. NV is an integer between 1 and 48
%   and must agree with the value used to obtain the CWT. NV defaults to
%   10.
%
%   X = CWTMAG2SIG(...,TimeBandwidth=TB) uses the positive scalar
%   time-bandwidth parameter, TB, to reconstruct using the Morse wavelet.
%   The symmetry parameter (gamma) of the Morse wavelet is assumed to equal
%   3. TimeBandwidth defaults to 60. TimeBandwidth is valid only when
%   Wavelet is "morse". You cannot specify both the TimeBandwidth and
%   WaveletParameters.
% 
%   X = CWTMAG2SIG(...,WaveletParameters=PARAM) uses the parameters PARAM
%   to specify the Morse wavelet used for reconstruction. PARAM is a
%   two-element vector. The first element is the symmetry parameter (gamma)
%   and the second parameter is the time-bandwidth parameter.
%   WaveletParameters defaults to [3,60]. WaveletParameters is valid only
%   when Wavelet is "morse". You cannot specify both the TimeBandwidth and
%   WaveletParameters.
%   
%   X = CWTMAG2SIG(...,FilterBank=FB) specifies the CWTFILTERBANK object,
%   FB used in the CWT. FB must have Boundary="periodic". If you specify
%   FilterBank, you cannot specify any other options above for CWT
%   parameters. All options for the computation of the CWT are defined as
%   properties of the filter bank object.
%
%   Low Frequency Information:
%
%   X = CWTMAG2SIG(...,ScalingCoefficients=SCALCFS) specifies the scaling
%   coefficients, SCALCFS, used for reconstruction. You can obtain the
%   scaling coefficients as an optional output of CWT. For single-channel
%   signals, SCALCFS is a real-valued vector which is the same length as
%   the column size of the wavelet coefficient matrix. For multichannel
%   signals, concatenate the scaling coefficients of each channel along the
%   third dimension, and specify SCALCFS as that 3-D array. You cannot
%   specify both the SignalMean and ScalingCoefficients.
%
%   X = CWTMAG2SIG(...,SignalMean=MEAN) adds the scalar MEAN to the
%   reconstructed signal X. For a single-channel signal, MEAN is a
%   real-valued scalar. For a multichannel signal, MEAN is a real-valued
%   vector, where the i-th element of MEAN is the mean of the i-th channel.
%   Because the CWT does not preserve the signal mean, the reconstructed
%   signal is zero-mean by default. Note that adding a nonzero MEAN to a
%   frequency- or period-limited reconstruction adds a zero-frequency
%   component to the reconstruction. You cannot specify both the SignalMean
%   and ScalingCoefficients.
%
%   Algorithm Parameters:
%
%   X = CWTMAG2SIG(...,InitializePhaseMethod=INITPHASEMETHOD) initializes
%   the CWT phases using the method specified in INITPHASEMETHOD:
%
%       "random" - Initialize the phases as random numbers distributed
%                 uniformly in the interval [-pi,pi].
%
%       "zeros"  - Initialize the phases as zeros.
%
%   The default value of INITPHASEMETHOD is "random".
%
%   X = CWTMAG2SIG(...,InitialPhase=INITPHASE) specifies the initial CWT
%   phases. INITPHASE must have the same size as the CWT magnitude CFS, and
%   its elements must be in the range [-pi,pi]. This option is incompatible
%   with 'InitializePhaseMethod'.
%
%   X = CWTMAG2SIG(...,MaxIterations=MAXITERATIONS) specifies the maximum
%   number of optimization iterations used in the reconstruction process.
%   The process stops when the number of iterations is greater than
%   MAXITERATIONS. The default value of MAXITERATIONS is 300.
%
%   X = CWTMAG2SIG(...,InconsistencyTolerance=INCONSTOL) specifies the
%   inconsistency tolerance. The reconstruction process stops when the
%   normalized inconsistency is smaller than INCONSTOL. The tolerance is a
%   positive real value and defaults to 1e-5. The inconsistency is defined
%   as ||Sest_i-Sest_(i-1)||/CFS, where Sest_i is the estimated complex CWT
%   at the i-th iteration.
%
%   X = CWTMAG2SIG(...,Optimizer=OPTIMIZER) specifies the optimizer used in
%   the gradient descent algorithm. OPTIMIZER can be one of these:
%
%     "lbfgs"       - Limited-memory BFGS (L-BFGS) optimizer.
%
%     "sgdm"        - Stochastic gradient descent with momentum (SGDM)
%                     optimizer.
%
%     "adam"        - Adaptive moment estimation (ADAM) optimizer.
%
%     "rmsprop"     - Root mean square propagation (RMSProp) optimizer.
%
%   The default is "lbfgs". The L-BFGS algorithm generally requires the
%   fewest iterations and yields the best reconstruction results. However,
%   particularly when processing batch data inputs, the computation speed
%   can be slow. When a large number of calculations need to be performed
%   quickly, other optimizers are recommended.
%
%   X = CWTMAG2SIG(...,StepSize=STEPSIZE) specifies the step size for the
%   iterative update used by the gradient descent algorithm. STEPSIZE is a
%   real positive scalar and defaults to 0.01. You cannot specify a step
%   size if OPTIMIZER is "lbfgs".
%
%   X = CWTMAG2SIG(...,Display=DISPLAY) specifies whether to display the
%   normalized inconsistency value on the command line every 20 iterations
%   and show the final stopping information. The default is false.
%
%   [X,T] = CWTMAG2SIG(...) returns the times of the reconstructed signals.
%   When a sample rate is provided, T is a numeric vector in seconds.
%   When a sample time duration Ts is provided, T is a duration array
%   with the same unit as Ts. T is a numeric vector in sample numbers if no
%   time information is provided. The number of rows in X equals the length
%   of T.
%
%   [X,T,INFO] = CWTMAG2SIG(...) returns a structure containing this
%   convergence information:
%
%      ExitFlag:            Termination flag for each channel. A value of 0
%                           indicates the algorithm stopped when it reached
%                           the maximum number of iterations or could not
%                           converge. A value of 1 indicates the algorithm
%                           stopped when it achieved the inconsistency
%                           tolerance. For multichannel signals, ExitFlag
%                           is a vector.
%
%      NumIterations:       Total number of iterations for each channel.
%                           For multichannel signals, NumIterations is a
%                           vector.
%
%      Inconsistency:       Average relative improvement towards
%                           convergence between the final two iterations
%                           for each channel. For multichannel signals,
%                           Inconsistency is a vector.
%
%      ReconstructedCWT:    Reconstructed CWT at the final iteration.
%
%      ReconstructedPhase:  Reconstructed phase at the final iteration.
%
% % EXAMPLE 1:
%      % Reconstruct a sinusoid.
%      fs = 1000;
%      ts = 0:1/fs:2-1/fs;
%      x = cos(2*pi*2*ts).'+1;
%      [CFS,~,~,~,d] = cwt(x,ExtendSignal=false);
%      xrec = cwtmag2sig(abs(CFS),Display=true,ScalingCoefficients=d);
%      plot(ts,x,ts,xrec,"--")
%      legend("Original","Reconstructed")
% 
% % EXAMPLE 2:
%      % Reconstruct a multichannel ECG signal using a CWT filter bank.
%      load mit200
%      idx1 =  1:1000;
%      idx2 = 3001:4000;
%      x1 = ecgsig(idx1);
%      x2 = ecgsig(idx2);
%      cfb = cwtfilterbank(SignalLength=length(x1), ...
%          Boundary="periodic");
%      [CFS1,~,~,~,d1]  = cwt(x1,Filterbank=cfb);
%      [CFS2,~,~,~,d2]  = cwt(x2,Filterbank=cfb);
%      CFS = cat(3,CFS1,CFS2);
%      d = cat(3,d1,d2);
%      x = cat(2,x1,x2);
%      xrec = cwtmag2sig(abs(CFS),Display=true, ...
%          ScalingCoefficients=d,Filterbank=cfb);
%      tiledlayout(2,1)
%      nexttile
%      plot(tm(idx1),x(:,1),tm(idx1),xrec(:,1),"--")
%      legend('Original','Reconstructed')
%      nexttile
%      plot(tm(idx2),x(:,2),tm(idx2),xrec(:,2),"--")
%      legend("Original","Reconstructed")
%
%   See also CWT, CWTFILTERBANK.

%   Copyright 2023 The MathWorks, Inc.

arguments
    % Error out for dlarray
    Smag (:,:,:) {mustBeFloat,mustBeReal,mustBeA(Smag,["single","double","gpuArray"])}
 
    % CWT info
    % ----------------------------------
    % Without Filterbank
    timeinfo {mustBeScalarOrEmpty} = [];
    cwtNvargs.Wavelet {mustBeTextScalar,mustBeMember(cwtNvargs.Wavelet,["morse","amor","bump"])} ='';
    cwtNvargs.FrequencyLimits = []
    cwtNvargs.PeriodLimits duration = seconds(zeros(1,0))
    cwtNvargs.VoicesPerOctave {mustBeFloat,mustBeScalarOrEmpty,mustBeInteger,mustBePositive,mustBeLessThanOrEqual(cwtNvargs.VoicesPerOctave,48)} = [] 
    cwtNvargs.TimeBandwidth = []
    cwtNvargs.WaveletParameters = []
    % With Filterbank
    % Other arguments in this section cannot be specified.
    Nvargs.FilterBank {mustBeType(Nvargs.FilterBank,"cwtfilterbank")} =  []

    % Low frequency info
    % ----------------------------------
    Nvargs.SignalMean {mustBeFloat,mustBeVector(Nvargs.SignalMean,"allow-all-empties"),mustBeReal} = [];
    Nvargs.ScalingCoefficients {mustBeFloat,mustBeReal} = [];

    % General Info (GD)
    % ----------------------------------
    Nvargs.MaxIterations {mustBeNumeric,mustBeScalarOrEmpty,mustBePositive,mustBeInteger} = 300;
    Nvargs.Optimizer {mustBeTextScalar,mustBeMember(Nvargs.Optimizer,["adam","sgdm","rmsprop","lbfgs"])} = "lbfgs";
    Nvargs.Display = false;
    Nvargs.InconsistencyTolerance {mustBeFloat,mustBeScalarOrEmpty,mustBePositive,mustBeReal} = 1e-5

    Nvargs.StepSize {mustBeFloat,mustBeScalarOrEmpty,mustBePositive,mustBeReal} = [];
    Nvargs.InitializePhaseMethod {mustBeTextScalar,mustBeMember(Nvargs.InitializePhaseMethod,["random","zeros"])} ='';
    Nvargs.InitialPhase {mustBeFloat,mustBeValidPhase(Nvargs.InitialPhase,Smag),mustBeReal}= [];
end

% Workaround for unsupported partial-match and case-insensitive
% mustBeMemeber in argument validation
% NOTE isempty("") = true. So convert everything to char.
cwtNvargs.Wavelet = validatestring(cwtNvargs.Wavelet,{'morse','amor','bump',''},'cwtmag2sig');
Nvargs.Optimizer = validatestring(Nvargs.Optimizer,{'adam','sgdm','rmsprop','lbfgs',''},'cwtmag2sig');
Nvargs.InitializePhaseMethod = validatestring(Nvargs.InitializePhaseMethod,{'random','zeros',''},'cwtmag2sig');

% Check and handle mixed datatype
% GPU
if isgpuarray(Nvargs.ScalingCoefficients) ...
    || isgpuarray(Nvargs.SignalMean) ...
    || isgpuarray(Nvargs.InitialPhase)
    Smag = gpuArray(Smag);
end
% Single
if isUnderlyingType(Nvargs.ScalingCoefficients,"single") ...
    || isUnderlyingType(Nvargs.SignalMean,"single") ...
    || isUnderlyingType(Nvargs.InitialPhase,"single")
    Smag = single(Smag);
end

% Check Smag size (minLength)
if size(Smag,2)<4
    error(message("wavelet_deep:cwtmag2sig:InvalidCOLSizeS"));
end

% GD Method Parameters
if isempty(Nvargs.StepSize)
    Nvargs.StepSize = 1e-2; % defualt
else
    if strcmp(Nvargs.Optimizer,"lbfgs")
        error(message("wavelet_deep:cwtmag2sig:StepSizeWithLBFGS"))
    end
end

% FilterBank Specified
if ~isempty(Nvargs.FilterBank)
    % Check all other cwt NVP arguments are empty
    vars = fields(cwtNvargs);
    for ii = 1:numel(vars)
        if ~isempty(cwtNvargs.(vars{ii}))
            error(message("wavelet_deep:cwtmag2sig:FilterBankwithOtherOptions"))
        end
    end
    fb = Nvargs.FilterBank;
    % Check Fs/Ps is empty
    if ~isempty(timeinfo)
        error(message("wavelet_deep:cwtmag2sig:FilterBankwithOtherOptions"))
    end
    % Check SignalLength
    if size(Smag,2)~=fb.SignalLength
        error(message("wavelet_deep:cwtmag2sig:InvalidCOLSizeSwithFilterBank"))
    end

    % Extract timeinfo from filterbank
    if isempty(fb.SamplingPeriod)
        timeinfo =  fb.SamplingFrequency;
    else
        timeinfo =  fb.SamplingPeriod;
    end

    % Check Boundary
    % Do not support "reflection"
    if startsWith(fb.Boundary,"r")
        error(message("wavelet_deep:cwtmag2sig:InvalidBoundary"))
    end
% No FilterBank Specified
else
    % get wavelet
    if isempty(cwtNvargs.Wavelet)
        wname ="morse";
    else
        wname = cwtNvargs.Wavelet;
    end
   % Defer arguments validation into filterbank construction function
    optsCWT = {};
    optsCWT.SignalLength = size(Smag,2);
    optsCWT.Wavelet = wname;
    optsCWT.Boundary = "periodic";

    % TF
    if ~isempty(timeinfo)
        if isduration(timeinfo)
            optsCWT.SamplingPeriod  = timeinfo;
        else
            mustBePositive(timeinfo);
            optsCWT.SamplingFrequency = timeinfo;
        end
    else
        timeinfo = 1; % defaults to 1
    end
    % FrequencyLimits
    optsCWT.FrequencyLimits = cwtNvargs.FrequencyLimits;
    % PeriodLimits
    if ~isempty(cwtNvargs.PeriodLimits)
        optsCWT.PeriodLimits = cwtNvargs.PeriodLimits;
    end
    % VoicesPerOctave
    if ~isempty(cwtNvargs.VoicesPerOctave)
        optsCWT.VoicesPerOctave = cwtNvargs.VoicesPerOctave;
    end
    % TimeBandwidth
    if ~isempty(cwtNvargs.TimeBandwidth)
        optsCWT.TimeBandwidth = cwtNvargs.TimeBandwidth;
    end
    % WaveletParameters
    if ~isempty(cwtNvargs.WaveletParameters)
        optsCWT.WaveletParameters = cwtNvargs.WaveletParameters;
    end

    % ------------------------------------------------------------------
    c = namedargs2cell(optsCWT);
    try
        fb = cwtfilterbank(c{:});
        % Defer arguments validation into cwtfilterbank function
    catch e
        throw(e)
    end
end

% Check number of rwos in Smag is as expected 
nrow = numel(scales(fb));
if size(Smag,1)~= nrow
    error(message("wavelet_deep:cwtmag2sig:InvalidRowSizeS",nrow));
end
% Check size of ScalingCoefficients
if ~isempty(Nvargs.ScalingCoefficients) 
    if size(Nvargs.ScalingCoefficients,3)~=size(Smag,3)
        error(message("wavelet_deep:cwtmag2sig:InvalidChannelScalingCoefficients"));
    end
    if ~isvector(Nvargs.ScalingCoefficients(:,:,1)) || numel(Nvargs.ScalingCoefficients(:,:,1))~=size(Smag,2)
        error(message("wavelet_deep:cwtmag2sig:InvalidSizeScalingCoefficients"));
    end
end
% Check size of SignalMean
if ~isempty(Nvargs.SignalMean) && numel(Nvargs.SignalMean)~=size(Smag,3)
    error(message("wavelet_deep:cwtmag2sig:InvalidSizeSignalMean"));
end

%% Initialization Xinit
% Initialize phase
if isempty(Nvargs.InitializePhaseMethod)&& isempty(Nvargs.InitialPhase)
    Nvargs.InitializePhaseMethod ="random";
end
if strcmpi(Nvargs.InitializePhaseMethod,'random')
    initialPhase = exp(1j*(2*pi*rand(size(Smag))-pi));
    S = complex(Smag.*initialPhase);
elseif strcmpi(Nvargs.InitializePhaseMethod,'zeros')
    S = complex(Smag);
else % user input
    initialPhase = exp(1j*Nvargs.InitialPhase);
    S = complex(Smag.*initialPhase);
end

% Initialize Signal
% Prepare ICWT options
optsICWT={};
optsICWT.VoicesPerOctave = fb.VoicesPerOctave;
if ~isempty(fb.TimeBandwidth)
    optsICWT.TimeBandwidth = fb.TimeBandwidth;
end
if ~isempty(fb.WaveletParameters)
    optsICWT.WaveletParameters = fb.WaveletParameters;
end
% Use ICW to init signal
% Convert Tx1xC ScalingCoefficients into 1xTxC (always T>=4)
if ~isempty(Nvargs.ScalingCoefficients) && size(Nvargs.ScalingCoefficients,1)~=1
    Nvargs.ScalingCoefficients = permute(Nvargs.ScalingCoefficients,[2,1,3]);
end

xinit = zeros(size(S,[2,3]));
for ii=1:size(S,3)
    if ~isempty(Nvargs.SignalMean)
        optsICWT.SignalMean = Nvargs.SignalMean(ii);
    end

    if ~isempty(Nvargs.ScalingCoefficients)
        optsICWT.ScalingCoefficients = Nvargs.ScalingCoefficients(:,:,ii);
    end

    c = namedargs2cell(optsICWT);
    xinit(:,ii) = icwt(S(:,:,ii),fb.Wavelet,c{:});
end

%% Phase Retrieval GD

% Prepare phaseRetrievalGD options
opts = {};
opts.StepSize = Nvargs.StepSize;
opts.Optimizer = Nvargs.Optimizer;
opts.Display = Nvargs.Display;
opts.MaxIterations = Nvargs.MaxIterations;
opts.InconsistencyTolerance = Nvargs.InconsistencyTolerance;
opts.DataType = underlyingType(Smag);
opts.IncludeLowpass = ~isempty(Nvargs.ScalingCoefficients);

% Process ScalingCoefficients
[psifvec,filteridx] = cwtfilters2array(fb,-Inf,IncludeLowpass=opts.IncludeLowpass);
dlfun = @(x)dlcwt(real(x),psifvec,filteridx);

if opts.IncludeLowpass
    S = cat(1,S,Nvargs.ScalingCoefficients);
end

[xr,infoGD] = signalwavelet.deep.internal.tfmag2sig.phaseRetrievalGD(S,xinit,dlfun,opts,"cwt");


%% Prepare Output
% X
X = cast(xr,'like',Smag);
if ~isempty(Nvargs.SignalMean)
    X = X-mean(X,1) + Nvargs.SignalMean(:)';
end

% T
% Handle Fs/Ts information and prepare the time vector
if isnumeric(timeinfo)
    T = cast(1/timeinfo*(0:size(X,1)-1)','like',Smag);
else
    T = timeinfo*(0:size(X,1)-1)';
end

% info
% Store results
INFO = struct('ExitFlag',infoGD.ExitFlag, ...
        'NumIterations',infoGD.NumIterations, ...
        'Inconsistency',cast(infoGD.Inconsistency,'like',Smag));
if opts.IncludeLowpass
    INFO.ReconstructedCWT  = cast(infoGD.Sr(1:end-1,:,:),'like',Smag);
else
    INFO.ReconstructedCWT  = cast(infoGD.Sr,'like',Smag);
end
INFO.ReconstructedPhase = angle(INFO.ReconstructedCWT);

end


% Custom overwritten mustBeMember function
% Allow empty default string
% Allow partial match check
function mustBeMember(A,B)
if ~isempty(char(A))
    validatestring(A,B,'cwtmag2sig');
end
end
% Custom overwritten mustBeType function (for checking cwtfilterbank)
% Allow empty default 
function mustBeType(var,type)
    if ~isempty(var)
        mustBeA(var,type);
    end
end
% Custom validation function for phase
function mustBeValidPhase(Phase,Smag)
if ~isempty(Phase)
    if any(size(Phase) ~=size(Smag) )
        error(message("wavelet_deep:cwtmag2sig:InvalidInitialPhase"))
    end
    if any(abs(Phase)>pi )
        error(message("wavelet_deep:cwtmag2sig:InvalidInitialPhaseRange"))
    end
end
end