function OptimInfo = minimizeWithSparsity(Estimator,SparsOpt)
%minimizeWithSparsity Minimize nonlinear least squares with added norm penalty using ADMM.
% Penalty takes the form of 0, 1, 2, or Inf-norm of a weighted parameter
% vector w'*p. p may represent groups. 
%
% Estimator: An idestimatorpack.AbstractEstimator object (can perform minimization of
%            NLSQ objective with no additional penalty) 
% SparsOpt: idoptions.ProximalMinimization object.

% Options:
% - some choice of algorithm (Barratt/Sharma, ...)
% - parameter groups and weights; not all parameters are penalized
% - ADMM related knobs - max iter, tolerances, acceleration
% - penalty size, adaptiveness
% (see idoptions.ProximalMinimization, idoptions.NLARXSparsification)
% Result: structure with estimation results, termination conditions.

% Copyright 2021-2022 The MathWorks, Inc.

% Approach: 
% Ref: Barratt, Sharma, "Interpreting a Decomposition Method for Sparse Nonconvex Optimization",
% ICML 2018.
%
% I.C.: p(0), z(0), u(0)
% 1. p(k+1) = argmin_p(V(p) + rho/2 ||p-z(k)+u(k)||_2^2)
% 2. z(k+1) = argmin_z(\lambda||z||_n + \rho/2 ||p(k+1)-z-u(k)||_2^2), where np is 0, 1, inf
%    For n=1, the solution in coordinate-wise soft thresholding: 
%    z(k+1)_i = S(p(k+1)_i + u(k)_i)), <--- prox operator 
%               where:
%               S(a)= (a-m)+ - (-a-m)+, m = \lambda/rho
%    (similar for other p; see idpack.proxop)
% 3. u(k+1) = u(k) + p(k+1) - z(k+1)

% Algorithmic aspects
%  - Stopping criteria in coordinate descent
%  - Lambda search using cross validation 
%  - Parallelized bagging for regressor preference

pinfo = Estimator.Info;
p0 = pinfo.Value;
npFree = numel(p0);
LambdaUser = SparsOpt.Lambda;
regnames = Estimator.Options.Utility.RegInfo.RegNames;
g = Estimator.Options.Utility.RegParGroupIndex;
AllGroups = unique(g,'stable');
AllGroups(isnan(AllGroups)) = [];
nr0 = numel(regnames);
regnames = regnames(AllGroups); % remove regressors already not in use
nr = numel(regnames);

% Not all free parameters are regularized. Excluded are the parameters that do not multiply the
% regressors ("others"), or the ones that contain a fixed but non-zero one in their group.
RPF = Estimator.Options.Utility.RegParFree; % true if "others" or regularized

% determine which parameters should be subject to sparsification regularization
iRegularized = isfinite(g)' & RPF; % regularized
g(~iRegularized) = [];
% remove truly/originally fixed parameters (not fixed owing to group induction) as dictated by
% pinfo.Free, since those are not considered parameters at all
iRegularized = iRegularized(pinfo.Free); % now length(p0) = length(iRegularized) = npFree

Measure = SparsOpt.SparsityMeasure;
ReweightedHeuristics = Measure=="log-sum";
if ReweightedHeuristics
   Measure = "l1";
end
maxiter = SparsOpt.MaxIterations;
if ~any(iRegularized)
   maxiter = 1;
end
Estimator.Options.Regularization.Lambda = 0;
R0 = zeros(1,npFree);
R0(iRegularized) = 1;
% (don't set R which may be all 0)
Estimator.Options.Regularization.Nominal = 'model';

% initial guesses
zk = p0(iRegularized); 
uk = zeros(size(zk));

ActiveParameters = true(numel(zk),1);
ctr = zeros(npFree,1);
Estimator.Center = ctr;
fcnct = 0;
p_new = p0;
AllActive = true(npFree,1);
Min = Estimator.Info.Minimum;
Max = Estimator.Info.Maximum;

Disp = ~strcmpi(Estimator.Options.Display,'off');
StatusTitle = getString(message('Ident:estimation:strSparseSearh'));
if Disp
   PW = Estimator.Options.ProgressWindow;
   WaitFig = getWidget(PW);
else
   WaitFig = uifigure('Name',StatusTitle);
   matlab.graphics.internal.themes.figureUseDesktopTheme(WaitFig);
   WaitFig.Position(4) = WaitFig.Position(4)/2;
   cleanupObj = onCleanup(@()delete(WaitFig));
   PW = [];
end

wt = ones(1,length(zk));
k = 1;
GNorm = SparsOpt.GroupNorm;
while k<=maxiter % begin ADMM
   Estimator.Options.ProgressWindow.STOP = false;
   try
      OptimInfo_k = minimize(Estimator); % first minimization is unregularized
   catch E
      Estimator.Options.ProgressWindow.STOP = false;
      if k==1
         throw(E)
      else
         break; % accept last good as final
      end
   end
   fcnct = fcnct + OptimInfo_k.Output.funcCount;
   p_temp = OptimInfo_k.X;
   if k==1
      if ~any(iRegularized)
         Lambda = 1e-9;
         LambdaOverRho = 1;
      else
         % Note: we should not be computing sum(log(abs(p))) literally; that quantity can be negative
         parwt = localGroupNorms(p_temp(iRegularized),g,GNorm);
         % Pick a factor that puts the problem at the boundary of sparsity, assuming user-specified
         % penalty is 1.
         LambdaScale = OptimInfo_k.Cost*1e-2; % Lambda_Used = LambdaScale*User_specified_value
         LambdaOverRho = LambdaUser*max(1e-6,min(parwt));
         Lambda = LambdaUser*LambdaScale;
      end
      rho = Lambda/LambdaOverRho;
      Estimator.Options.Regularization.Lambda = rho/2; 
      
      mx = Estimator.Options.SearchOptions.MaxIterations;
      Estimator.Options.SearchOptions.MaxIterations = min(5,mx);      

      % wait bar
      msg = getString(message('Ident:estimation:strSparseProgress',0,0,nr));
      stopstr = getString(message('Ident:estimation:strStopSparse'));
      h = uiprogressdlg(WaitFig,'Message',msg,'Title',StatusTitle,'Value',0,...
         'Cancelable',true,'CancelText',stopstr,'ShowPercentage',true);
   end
   p_new(AllActive) = p_temp;
   
   Shrink_Vector = p_new(iRegularized)+uk;
   Shrink_Vector(~ActiveParameters) = 0;
   [z_new, GroupNorms] = idpack.proxop.prox_Group(Measure,GNorm,Shrink_Vector,g,LambdaOverRho,wt);
   ActiveGroups = GroupNorms>0;   
   u_new = uk + p_new(iRegularized)-z_new;

   if ReweightedHeuristics
      % apply reweighted heuristics 
      % Ref: Candes, Wakin, Boyd, "Enhancing Sparsity by Reweighted l1 Minimization", Journal of
      % Fourier Analysis and Applications, 2007.
      wt = 1./GroupNorms;
      wt(~ActiveGroups) = 1;
      ActiveParameters = z_new~=0;
   else
      ActiveParameters = ActiveParameters & z_new~=0;
   end
   
   p_new_k = p_new(iRegularized);
   p_new_k(~ActiveParameters) = 0;
   p_new(iRegularized) = p_new_k;
   
   uk = u_new;
   zk = z_new;
   
   ActiveInPenalty = true(size(p_new_k));
   ActiveInPenalty(~ActiveParameters) = false;
   AllActive = true(npFree,1);
   AllActive(iRegularized) = ActiveInPenalty;
   if ~any(ActiveGroups) && numel(ActiveGroups)==nr
      % stop
      localDispError('Ident:estimation:proxMinPenaltyTooLarge',PW,Disp)
   elseif ~all(AllActive)
      f0 = pinfo.Free;
      f = true(npFree,1);
      f(~AllActive) = false;
      f0(f0) = f;
      Estimator.Info.Free = f0;
      Estimator.Info.Value = p_new(AllActive);
      Estimator.Info.Minimum = Min(AllActive);
      Estimator.Info.Maximum = Max(AllActive);

      % regularization (works on non-participating parameters with weight zero and on participating
      % but non-zero parameters with weight 1.
      R = R0;
      R(~AllActive) = [];
      if ~any(R)
         Estimator.Options.Regularization.Lambda = 0;
      else
         Estimator.Options.Regularization.Lambda = rho/2;
         Estimator.Options.Regularization.R = diag(R); 
      end           
      ctr = zeros(npFree,1);
      ctr(iRegularized) = zk-uk;
      ctr(~AllActive) = [];
      Estimator.Center = ctr;
   end

   if h.CancelRequested
      break;
   else
      msg = getString(message('Ident:estimation:strSparseProgress',...
         k,sum(~ActiveGroups),nr));
      h.Message = msg;
      h.Value = k/maxiter;
      if Disp
         idDisplayEstimationInfo('Progress',{msg,' '},PW)   
      end
   end
   k = k+1;
   %sum(~ActiveGroups)
end % end ADMM

I = ismember(AllGroups,unique(g,'stable'));
AllGroups = true(size(AllGroups));
AllGroups(I) = ActiveGroups;
Regs = regnames(AllGroups);
OptimInfo.ActiveRegs = Regs;
OptimInfo.X = p_new;
if ~any(ActiveGroups) && numel(ActiveGroups)==nr
   localDispError('Ident:estimation:proxMinPenaltyTooLarge',PW,Disp)
   return
end
if Disp
   msg = getString(message('Ident:estimation:strSparseProgressFinal',...
      sum(~ActiveGroups),nr));
   idDisplayEstimationInfo('Result',msg,PW)
   stopstr = getString(message('Ident:estimation:strReducedRegSet'));
   idDisplayEstimationInfo('Result',[stopstr,':'],PW)
   S = idguis.internal.utils.utWrapForDisp(Regs);
   idDisplayEstimationInfo('Result',S,PW);
   idDisplayEstimationInfo('Result',{' '},PW)
end

Estimator.Options.Regularization.Lambda = 0;

end

%---------------------------------------------------------------------------------------
function localDispError(id,PW,Disp)
msg = message(id);
if Disp
   msgstr = getString(msg);
   idDisplayEstimationInfo('Error',msgstr,PW)   
end
error(msg)
end

%----------------------------------------------------------------------------------------
function V = localGetParameterNorm(p,g,GNorm,Measure)
% Compute penalty
gu = unique(g,'stable');
NoGroups = length(gu)==length(g);
if NoGroups
   switch Measure
      case "l1"
         V = norm(p,1);
      case "l0"
         V = sum(p~=0);
   end
else
   L = length(gu);
   v = zeros(1,L);
   for ct = 1:L
      I = g==gu(ct);
      v(ct) = norm(p(I),GNorm);
   end
   V = localGetParameterNorm(v,gu,[],Measure);
end
end

%----------------------------------------------------------------------------------------
function V = localGroupNorms(p,g,GroupNorm)
% Compute penalty.
gu = unique(g,'stable');
NoGroups = length(gu)==length(g);
if NoGroups
   V = abs(p);
else
   L = length(gu);
   V = zeros(1,L);
   for ct = 1:L
      I = g==gu(ct);
      V(ct) = norm(p(I),GroupNorm);
   end
end
end