 function [xs, info] = mmd_pl_nsqs_os_fly(x, Gb, data, R, varargin)
%|function [xs, info] = mmd_pl_nsqs_os_fly(x, Gb, data, R, [options])
%|
%| penalized-likelihood minimization for multi-material decomposition with 
%| constraints on pixels.
%| 1. each pixel contains at most three materials
%| 2. box constraints on each pixel of each material [pixmin_l, pixmax_l]
%| 3. sum to one constraint sum_i x_i = 1
%|
%| cost function:
%| negative log-likelihood with convex non-quadratic penalty
%| cost(x) = sum_i h(data_i; [Gx]_i) + R(x)
%|
%| optimization algorithm:
%| ordered subset optimization transfer using pixel-wise separable quadratic surrogates
%| (materials are still coupled) and generalized sequential minimization (GSMO).
%| 
%| see mmd_pl_decom.m for usage 
%| in
%|	x	[np,nl]		initial estimate
%|	Gb	[nd,np]		system matrix (see Gblock)
%|	data	{cell}		whatever data is needed for the likelihood
%|				data{1} must have size "[nb,na]", so for 3D
%|				one must use reshaper(yi, '2d').
%|	R			penalty object (see Robject.m)
%|
%| options
%|	dercurv	{function_handle} function returning derivatives and curvatures
%|				of negative log-likeihood via the call:
%|		[deriv curv] = dercurv(data, Gb{m}*x, curvtype, iblock, nblock)
%|				or choose 'trl' or 'wls' (default)
%|	niter 	#		# total iterations (default: 1+0, pure OS!)
%|	riter	#		# of penalty subiterations (default: 3).
%|	curvtype ''		curvature type:
%|					'pc' : precomputed (default)
%|					'oc' : optimal (ensures monotonic)
%|	pixmax	 [nl, 1]		upper bound on pixel values (default: 1)
%|	pixmin	 [nl, 1]		lower bound on pixel values (default: 0)
%|	isave	[]		list of iterations to archive
%|				(default: [] 'last)
%|	gi	[nd]		precomputed Gb * 1 factors.
%|	chat	(value)		verbosity (default: 0)
%|	userfun			user defined function handle (see default below)
%|  mtype    materials types for multi-material decomposition
%|                  'air' has to be one 
%|  triplets   library of possible material triplets
%|  how  ''  method for constrained optimization given material types
%|           'cp' chambolle-pock (cp) algorithm
%|           'cgsmo' generalized sequential miniminzation algorithm (GSMO) 
%|                  http://cmp.felk.cvut.cz/~xfrancv/libqp/html/
%|                  the solver is libqp_gsmo.c
%|                   one pixel at a time
%|           'mgsmo' vectorized GSMO by Yong Long 
%|   ig  image geometry" structure for turning short column(s) into array(s)
%|  
%| out
%|	xs	[np,nsave]	estimates each (saved) iteration
%|	info	[niter+1,?]	userfun output.  default is cpu times
%|
%|
%| on the fly curvature
%| 2013-12-21


if nargin < 4, help(mfilename), error(mfilename), end

cpu tic

Gb = block_ob(Gb, 'ensure'); % make it a block object (if not already)
nblock = block_ob(Gb, 'n');
starts = subset_start(nblock);

[np LL] = size(x);

% defaults
arg.dercurv = @de_pl_dercurv;
arg.curvtype = 'fly';
arg.riter = 3;
arg.niter = 1; % default is pure OS, no IOT
arg.isave = [];
arg.pixmin = zeros(LL, 1);
arg.pixmax = ones(LL, 1); % volume fraction maximum is 1
arg.gi = [];
arg.chat = false;
arg.userfun = @userfun_default;
arg.mtype = {'water', 'bone', 'air'};
arg.triplets = nchoosek(arg.mtype, 3); 
arg.how = 'cgsmo'; 
arg.ig = [];
arg.coupled_material = 1; % materials are coupled 
arg.xtrue = 1;
% options
arg = vararg_pair(arg, varargin);
arg.isave = iter_saver(arg.isave, arg.niter);
xtrue = arg.xtrue;

if ~isa(arg.dercurv, 'function_handle') % fix: or inline?
	error 'dercurv not function handle?'
end

[hi  I0 ftab] = deal(data{:});
[nb na MM] = size(hi);
ns = nb * na;


ticker(mfilename, 1, arg.niter)

ntri = size(arg.triplets, 1);

% flag for existence of material types 
arg.mflg = zeros(ntri, LL);
for t = 1 : ntri%loop over triplet materials
    ttype = arg.triplets(t, :);     
    for i = 1 : 3 %triplet 
        for j = 1 : LL %mtype
            if isequal(arg.mtype{j}, ttype{i})
                arg.mflg(t, j) = 1;
            end
        end
    end
end



% only on the fly curvature is available
% 2013-12-21
if streq(arg.curvtype, 'fly')
   ldenom = zeros(np, LL , LL);
else
    error 'use on the fly curvature'
end


xs = zeros(np, LL, length(arg.isave));
if any(arg.isave == 0)
    xs(:, :, arg.isave == 0) = x;
end

%
% SPS-OS iterations 
%

for iter = 1:arg.niter
	for iset = 1:nblock
		ticker([mfilename ': os'], [iter iset], [arg.niter nblock])
	
		iblock = starts(iset);
		ia = iblock:nblock:na;
        
        for ii = 1 : LL-1
            li = Gb{iblock} * x(:,ii);	% l=G*x "line integrals"
            sil(:,:,ii) = reshape(li, nb, length(ia));
        end
        
       % gradient of cost function, [np, L-1]
        [dh curv] = feval(arg.dercurv, data, sil, Gb, ...
            arg.curvtype, arg.coupled_material, iblock, nblock);
        %pad zeros for the air compoent
        dh = [dh zeros(np, 1)];
        % pad zeros for the air compoent
        for ll = 1 : LL-1
            ldenom(:, :, ll) = [squeeze(curv(:, :, ll)) zeros(np,1)];
        end
        
        x = inner_update(x, dh*nblock, ldenom*nblock, R, arg);
    end
    
    if any(arg.isave == iter)
		xs(:, :, arg.isave == iter) = x;
    end
    
	info(1+iter,:) = feval(arg.userfun);
    if arg.chat
        for ii = 1 : LL
            printf('%g-th material, Range %g %g', ii, min(x(:,ii)), max(x(:,ii)));
        end
    end
end

if ~isempty(arg.ig)
    xs = arg.ig.embed(xs);
end
%
% inner_update()
% given PWSPS for likelihood, update x using regularization subiterations
% penalize on air component too
function xew = inner_update(x, dh, ldenom, R, arg)

[np LL] = size(x);

for ii = 1 : arg.riter 
	if isempty(R)
        H = ldenom; % [np L L]
        dhdr = dh;
    else
        H = ldenom; % [np L L]
       for jj = 1 : LL
           H(:, jj, jj) = H(:, jj, jj) + R{jj}.denom(R{jj}, x(:,jj)); %[np L]
           dr(:, jj) = R{jj}.cgrad(R{jj}, x(:,jj)); %[np L]         
       end
       dhdr = dh + dr;         
    end
    
    Htmp = reshape(H, [np LL*LL]); 
    xtmp1 = repmat(x, [1 LL]); %[np LL*LL](x1 x2 x3 x1 x2 x3 ...)
    p = reshape(sum(reshape(xtmp1 .* Htmp, [np LL LL]), 2), [np LL]);
    p = dhdr - p;
   
    xtmp2 = reshape(repmat(x, [LL 1]), [np LL*LL]); % (x1 x1... x2 x2....x3 x3...)
    fold = sum(xtmp1 .* xtmp2 .* Htmp, 2); %[np 1]
    fold = 0.5 * fold + sum(p .* x, 2);

    xew = cons_3ofm_update(x, p, H, fold, arg.pixmin, arg.pixmax, arg.triplets, arg.mflg, arg.how);
end

%
% cons_3ofm_update
% enforce box [pixmin_l, pixmax_l] and volume conservation constraints
%
%
function xnew = cons_3ofm_update(xold, p, H, fold, pixmin, pixmax, triplets, mflg, how)

[np L] = size(xold);

ntri = size(triplets, 1);% # triplets
xtri = zeros(np, L, ntri);
ftri = zeros(np, ntri); % minimum f function vlaues

for t = 1 : ntri%loop over triplet materials
    mind = find(mflg(t,:) == 1);
    
    % compute p corresponding current material index
    [x ftri(:,t)] = cons_m3_update(xold(:, mind), p(:, mind), H(:, mind, mind), pixmin(mind), pixmax(mind), how);
  
    xtri(:, mind, t) = x;
    
end


[fnew indt] = min(ftri, [], 2); 

% make sure cost function goes downhill
% indx = fnew < fold; 
if 0
indx = (fnew - fold) < -10^-6;
indx = find(indx == 1);
npx = length(indx); %# pixels need constrained minimizers  

xnew = xold;

if npx == 0
    return
end


for j = 1 : npx   
    ind_update = indx(j);
    xnew(ind_update, :) = xtri(ind_update, :, indt(ind_update));
end
end

xnew = xold;
for j = 1 : np  
    xnew(j, :) = xtri(j, :, indt(j));
end



%
% cons_m3_update_cp
%
function  [mu fvalue]  = cons_m3_update(xold, p, H, pixmin, pixmax, how)

if streq(how, 'cp')
   error 'not completed'
elseif streq(how, 'cgsmo') %
   [mu fvalue] = cgsmo_update(xold, p, H, pixmin, pixmax);
elseif streq(how, 'mgsmo')
    [mu fvalue] = mgsmo_update(xold, p, H, pixmin, pixmax);
else
    error 'unknow method'
end


%
% mgsmo_update
%
% generalized sequential miniminzation algorithm (GSMO) 
% vectorized GSMO by Yong Long
% 
% min 0.5*x'*H*x + p'*x  
%    x                                      
%
%   subject to   a'*x = b 
%                LB(i) <= x(i) <= UB(i)   for all i=1:n
% 
% convert the optimziation problem to the above form
%
function  [mu fvalue] = mgsmo_update(xold, p, H, pixmin, pixmax)

[np L] = size(xold);

if L ~= 3 
    error '3 unknown materials only'
end

if size(H, 2) ~= 3
    error '3 unknown materials only'
end


opt.TolKKT = 10^-3; % Determines relaxed KKT conditions (default TolKKT=0.001)
opt.MaxIter = 500; % Maximal number of iterations
opt.verb = 0; %if > 0 then prints info every verb-th iterations (default 0)

a = ones(np, L);%[L x 1]
b = 1;
LB = repmat(pixmin', [np 1]); % [np x L]
UB = repmat(pixmax', [np 1]);
H = reshape(H, [np L*L]); %make it two-dimensional

% if any xold is not an feasible solution
% find one for those pixels
ind = (xold < LB) | (xold > UB);
ind = max(ind, [], 2);
ind = ind | (xold * a(1,:)' ~= b);
ind = find(ind > 0);

%create feasible solutions for pixels needing updates
% todo, use nearest neighbor for initialization instead of
% this one that favors the first material
nup = length(ind);
x0 = zeros(nup, L); 
xa = zeros(nup, 1); 
i = 0;
while any(sum(x0.* a(ind, :), 2)) ~= b
    i = i + 1;
    if i > length(a)
      error('No feasible solution exists.');
    end
     x0(:, i) = min(UB(ind,i),max(LB(ind,i),(b-xa)./ a(ind,i)));
     xa = xa + x0(:,i) .* a(ind, i);
end

xold(ind,:) = x0;

[mu, stat] = quad_prog_gsmo(H, p, a, b, LB, UB, xold, opt);
fvalue = stat.QP;


%
% cgsmo_update
%
% generalized sequential miniminzation algorithm (GSMO) 
% http://cmp.felk.cvut.cz/~xfrancv/libqp/html/
%
% min 0.5*x'*H*x + p'*x  
%    x                                      
%
%   subject to   a'*x = b 
%                LB(i) <= x(i) <= UB(i)   for all i=1:n
% 
% convert the optimziation problem to the above form
%

function  [mu  fvalue] = cgsmo_update(xold, p, H, fold, pixmin, pixmax)

[np L] = size(xold);

if L ~= 3 
    error '3 unknown materials only'
end

if size(H, 2) ~= 3
    error '3 unknown materials only'
end

opt.TolKKT = 10^-3; % Determines relaxed KKT conditions (default TolKKT=0.001)
opt.MaxIter = 500; % Maximal number of iterations
opt.verb = 0; %if > 0 then prints info every verb-th iterations (default 0)

a = ones(np, L);%[L x 1]
b = 1;
LB = repmat(pixmin', [np 1]); % [np x L]
UB = repmat(pixmax', [np 1]);

% check whether gammas are feasible solution
% only undapte the pixels whose gammas are not feasible solution 
% 
ind = (xold < LB) | (xold > UB);
ind = max(ind, [], 2);
ind = ind | (xold * a(1,:)' ~= b);
ind = find(ind > 0);

% create feasible solutions for pixels needing updates
nup = length(ind);
x0 = zeros(nup, L); 
xa = zeros(nup, 1); 
QP = zeros(nup, 1); 

Hind = H(ind, :, :);
dhdrind = dhdr(ind, :);
aind = a(ind, :);
LBind = LB(ind, :);
UBind = UB(ind, :);
xoldind = xold(ind, :);

i = 0;
while any(sum(x0.* aind, 2)) ~= b
    i = i + 1;
    if i > length(a),
      error('No feasible solution exists.');
    end
    x0(:, i) = min(UB(ind,i),max(LBind(:,i),(b-xa)./ aind(:,i)));
    xa = xa + x0(:,i) .* aind(:, i);
end

mu = xold;
fvalue = fold;

%parfor j = 1 : nup
parfor (j = 1 : nup, 12)   
  
    
    [x(j, :), QP(j), dummy1, dummy2] = ...
        libqp_gsmo_mex(squeeze(Hind(j,:,:) ), p(j,:)', aind(j,:)', b, LBind(j,:)', UBind(j,:)', x0(j,:)',...
        opt.MaxIter, opt.TolKKT, opt.verb);
end

mu(ind,:) = x;
fvalue(ind) = QP; 








