function Beta = ADMMimplTall(X_,Y_,W_,Beta,rho,lambda,doridge,...
    N_chunks,Wk_,hmarkforreuse,hupdateforreuse,hfixedchunkfun,...
    betaTol,gradTol,admmIterations,tallPassLimit,progressF,verbose,...
    lossfun_ADMMLBFGS,betaTol_ADMMLBFGS,gradTol_ADMMLBFGS,...
    iterationlimit_ADMMLBFGS, doridge_ADMMLBFGS,...
    hessianHistorySize_ADMMLBFGS,dowolfe_ADMMLBFGS,doBias_ADMMLBFGS,...
    FM,expType,sigma,varargin)
% ADMMimpl ADMM for tall.
%
%   ADMMimpl is a simplified version to
%   classreg.learning.linearutils.ADMMimpl with support to be nested into a
%   gather-free scheduler.

%   Copyright 2017-2018 The MathWorks, Inc

iContainer = internal.stats.bigdata.Scheduler.catchJobContainerFromInputArgs(varargin);
if isempty(iContainer)
    Beta = internal.stats.bigdata.Scheduler.createAndRun(...
       @internal.stats.bigdata.ADMMimplTall,...
       X_,Y_,W_,Beta,rho,lambda,doridge,...
       N_chunks,Wk_,hmarkforreuse,hupdateforreuse,hfixedchunkfun,...
       betaTol,gradTol,admmIterations,tallPassLimit,progressF,verbose,...
       lossfun_ADMMLBFGS,betaTol_ADMMLBFGS,gradTol_ADMMLBFGS,...
       iterationlimit_ADMMLBFGS, doridge_ADMMLBFGS,...
       hessianHistorySize_ADMMLBFGS,dowolfe_ADMMLBFGS,doBias_ADMMLBFGS,...
       FM,expType,sigma);
   return
end

%-%-% SCHEDULER: PRELUDE  %-%-% 

% non-initialized shared vars between blocks
Beta_old=[];gra=[];RK_squared_mag=[];

dbeta_mag = Inf;
g_mag = Inf;
iter = 0;

% Evaluate objective function on initial values, note output is lazy, i.e.
% actual computation is deferred, so pass over data will be combined with the
% Beta update in the ADMM algorithm.
objgraF_ = progressF.LazyObjGraFunctor;

if admmIterations<=1
    progressF.Solver = 'INIT';
else
    progressF.Solver = 'ADMM';
end

% detect if it is classification or regression
doClass = 0; %regression
if ~iscell(lossfun_ADMMLBFGS)  % when epsilon is given it is augmented to the lossfun
    if any(strcmp(lossfun_ADMMLBFGS,{'hinge','logit'}))
        doClass = 2; %binary classification (one class is not supported)
    end
end

% Initialize Primal and Dual Variables
BetaK_ = hfixedchunkfun(1,@(x) Beta',Y_);
DeltaK_ = hfixedchunkfun(1,@(x) zeros(1,numel(Beta)),Y_);

BetaK_DeltaK_ = [BetaK_ DeltaK_];

if admmIterations>1
    Wk_ =  hmarkforreuse(Wk_);
    [BetaK_DeltaK_] = hmarkforreuse(BetaK_DeltaK_);
end
p = numel(Beta);

%-%-%  SCHEDULER: SETUP %-%-% 
iJob = internal.stats.bigdata.Job(@ch1,@pushVar,iContainer);

%-%-%  SCHEDULER: GATHER-FREE BLOCKS %-%-%

function iOutORiNext = ch1() % convergence testing 
    if (dbeta_mag<=betaTol || g_mag<=gradTol || iter>=admmIterations || progressF.DataPass>=tallPassLimit)
        iOutORiNext = {Beta};
    else
        iOutORiNext = @fh1;
    end
end

function iNext = fh1() % inside loop before gather

    iter = iter + 1;
    
    [~,gra_] = objgraF_(Beta);   % get gra for older Beta in the same data pass
    
    % Pull DeltaK_ and BetaK_ from the cached variable BetaK_DeltaK_
    BetaK_ = BetaK_DeltaK_(:,1:p);
    DeltaK_ =  BetaK_DeltaK_(:,(p+1):(2*p));  
    
    %%%%%%%%%%%%%%%%%%%%%%%%% DUAL VARIABLE UPDATE (DeltaK) %%%%%%%%%%%%%%%
    
    DeltaK_ = DeltaK_ + BetaK_- Beta';
    
    %%%%%%%%%%%%%%%%%%%%%%% LOCAL BETA UPDATE (BetaK) %%%%%%%%%%%%%%%%%%%%%
    
    BetaK_ = hfixedchunkfun(1, @(x,y,w,delta)  ...
        chunkBetaUpdateFun(x,y,w,delta,lossfun_ADMMLBFGS,Beta',...
        rho*(iter>1),...
        betaTol_ADMMLBFGS,gradTol_ADMMLBFGS,iterationlimit_ADMMLBFGS,...
        max(0,verbose-1),...
        doridge_ADMMLBFGS,...
        hessianHistorySize_ADMMLBFGS,dowolfe_ADMMLBFGS,doBias_ADMMLBFGS,...
        FM,expType,sigma,doClass), ...
        X_,Y_,W_, DeltaK_);  
    
    % Indicate that distributed variables will be reused so the are cached
    if admmIterations>1
        BetaK_DeltaK_ = hupdateforreuse(BetaK_DeltaK_,[BetaK_ DeltaK_]);
    end
    
    %%%%%%%%%%%%%%%%%%%%%%%%% GLOBAL BETA UPDATE (Beta) %%%%%%%%%%%%%%%%%%%
    
    BetaU_ = hfixedchunkfun(1,@calculateChunkWeigthedContributions,BetaK_+DeltaK_,Wk_);
    Beta_ = sum(BetaU_,1);   % weighted mean of BetaK_+DeltaK_
  
    % Apply regularization
    if doridge
        if iter>1 % do not regularize on the warm step
            % Regularize coefficients but not the bias
            Beta_ = Beta_./[1 repmat((1+lambda./rho./N_chunks),1,numel(Beta)-1)];
        end
    else % L1
        error(message('stats:tall:fitclinear:LassoNotSuported'))
    end
    
    RK_squared_magU_ = hfixedchunkfun(1,@calculateChunkWeigthedContributions,sum((BetaK_- Beta').^2,2),Wk_);
    RK_squared_mag_ = sum(RK_squared_magU_,1); % weigthed mean of the primal residual squared magnitude
    
    %%%%%%%%%%%%%%%%%%%%%%%%% TRIGGER COMPUTATIONS  %%%%%%%%%%%%%%%%%%%
    Beta_old = Beta; % will use a copy for later to check for convergence
    
    Beta = Beta_;
    gra = gra_;
    RK_squared_mag = RK_squared_mag_;
    %[Beta,gra,RK_squared_mag] = gather(Beta_,gra_,RK_squared_mag_);
    deferGather(iJob,{'Beta','gra','RK_squared_mag'},{Beta_,gra_,RK_squared_mag_});
    iNext = @fh2;
end

function iNext = fh2() % inside loop after gather
    Beta = Beta';

    %%%%%%%%%%%%%% CONVERGENCE AND PROGRESS TRACKING %%%%%%%%%%%%%%%%%%
    g_mag = sqrt(sum(gra.^2));
    r_mag = sqrt(RK_squared_mag);                 % Primal residual magnitude
    s_mag = sqrt(sum((rho.*(Beta-Beta_old)).^2)); % Dual residual magnitude
    beta_mag = sqrt(Beta'*Beta);
    dbeta_mag = sqrt(sum((Beta-Beta_old).^2))/beta_mag; % Relative change in Beta
    progressF.Beta = Beta_old;
    progressF.Solver = 'ADMM';
    progressF.IterationNumber = iter;
    progressF.PrimalResidual = r_mag;
    progressF.DualResidual = s_mag;
    iNext = @ch1;
end

function pushVar(iJob_vn,iJob_v)
    eval([iJob_vn '=iJob_v;']);
end

end

function betaWeigthed = calculateChunkWeigthedContributions(beta,wk)
    if wk==0
        % in the extreme case a coefficient is inf, but the weight of the
        % chunk is 0 we treat it separately to avoid inf.*0->NaN
        betaWeigthed = zeros(size(beta),'like',beta);
    else
        betaWeigthed = (beta).*wk;
    end
end
    
function BetaK = chunkBetaUpdateFun(...
    x,y,w,DeltaK,lossfun,Beta,...
    rho,...
    betaTol,gradTol,iterationlimit,...
    verbose,doridge,...
    historysize,dowolfe,fitbias,...
    FM,expType,sigma,doClass)

assert(iscell(x) && (numel(x)==1)); % sanity check, we assume previous cellification

x = x{1};
y = y{1};
w = w{1};

if isempty(x)
    % Setting all coefficients to 0 as they will be averaged with weight 0 anyways
    BetaK = zeros(size(Beta),'like',Beta);
    return;
end

% Set lineSearchType
% Note that currently we only support 'weakwolfe', setting line search to
% 'backtrack' is intended only exploratory work.
if dowolfe
    lineSearchType = 'weakwolfe';
else
    lineSearchType = 'backtrack';
end

% Get epsilon if loss is 'epsiloninsensitive'
if iscell(lossfun)
    epsilon = lossfun{2};
    lossfun = lossfun{1};
else
    epsilon = [];
end

% Expansion
if strcmpi(expType,'none')
    xm = x;
else
    xm = map(FM,x,sigma);
end

obj.Impl = classreg.learning.impl.LinearImpl.make(doClass,...
    (Beta(2:end)-DeltaK(2:end))',Beta(1)-DeltaK(1),...
    xm',y,w./sum(w),...
    lossfun,...
    doridge,...
    0,...   % Lambda
    [],...  % PassLimit
    [],...  % BatchLimit
    [],...  % NumCheckConvergence
    [],...  % BatchIndex
    [],...  % BatchSize
    {'lbfgs'},...
    betaTol,...
    gradTol,...
    1e-6,... % DeltaGradientTolerance
    [],...   % LearnRate
    [],...   % OptimizeLearnRate
    [],[],[],... % Validation
    iterationlimit,...
    [],...   % TruncationPeriod,...
    fitbias,... % FitBias
    false,... % PostFitBias
    epsilon,... % Epsilon
    historysize,... % hessian History Size
    lineSearchType,... % Line search
    rho, ... % Consensus
    [],... % Stream
    verbose);

BetaK = [obj.Impl.Bias,obj.Impl.Beta(:)'];

end

