function [x cost info] = trustregions(problem, x, options)
% Riemannian trust-regions solver for optimization on manifolds.
%
% function [x cost info] = trustregions(problem)
% function [x cost info] = trustregions(problem, x0)
% function [x cost info] = trustregions(problem, x0, options)
% function [x cost info] = trustregions(problem, [], options)
%
% This is the Riemannian Trust-Region solver (with tCG inner solve), named
% RTR. This solver will attempt to minimize the cost function described in
% the problem structure. It requires the availability of the cost function
% and of its gradient. It will issue calls for the Hessian. If no Hessian
% nor approximate Hessian is provided, a standard approximation of the
% Hessian based on the gradient will be computed. If a preconditioner for
% the Hessian is provided, it will be used.
%
% For a description of the algorithm and theorems offering convergence
% guarantees, see the references below. Documentation for this solver is
% available online at:
%
% http://www.manopt.org/solver_documentation_trustregions.html
%
%
% The initial iterate is x0 if it is provided. Otherwise, a random point on
% the manifold is picked. To specify options whilst not specifying an
% initial iterate, give x0 as [] (the empty matrix).
%
% The two outputs 'x' and 'cost' are the last reached point on the manifold
% and its cost. Notice that x is not necessarily the best reached point,
% because this solver is not forced to be a descent method. In particular,
% very close to convergence, it is sometimes preferable to accept very
% slight increases in the cost value (on the order of the machine epsilon)
% in the process of reaching fine convergence. In practice, this is not a
% limiting factor, as normally one does not need fine enough convergence
% that this becomes an issue.
%
% The output 'info' is a struct-array which contains information about the
% iterations:
%   iter (integer)
%       The (outer) iteration number, or number of steps considered
%       (whether accepted or rejected). The initial guess is 0.
%	cost (double)
%       The corresponding cost value.
%	gradnorm (double)
%       The (Riemannian) norm of the gradient.
%	numinner (integer)
%       The number of inner iterations executed to compute this iterate.
%       Inner iterations are truncated-CG steps. Each one requires a
%       Hessian (or approximate Hessian) evaluation.
%	time (double)
%       The total elapsed time in seconds to reach the corresponding cost.
%	rho (double)
%       The performance ratio for the iterate.
%	rhonum, rhoden (double)
%       Regularized numerator and denominator of the performance ratio:
%       rho = rhonum/rhoden. See options.rho_regularization.
%	accepted (boolean)
%       Whether the proposed iterate was accepted or not.
%	stepsize (double)
%       The (Riemannian) norm of the vector returned by the inner solver
%       tCG and which is retracted to obtain the proposed next iterate. If
%       accepted = true for the corresponding iterate, this is the size of
%       the step from the previous to the new iterate. If accepted is
%       false, the step was not executed and this is the size of the
%       rejected step.
%	Delta (double)
%       The trust-region radius at the outer iteration.
%	cauchy (boolean)
%       Whether the Cauchy point was used or not (if useRand is true).
%   And possibly additional information logged by options.statsfun.
% For example, type [info.gradnorm] to obtain a vector of the successive
% gradient norms reached at each (outer) iteration.
%
% The options structure is used to overwrite the default values. All
% options have a default value and are hence optional. To force an option
% value, pass an options structure with a field options.optionname, where
% optionname is one of the following and the default value is indicated
% between parentheses:
%
%   tolgradnorm (1e-6)
%       The algorithm terminates if the norm of the gradient drops below
%       this. For well-scaled problems, a rule of thumb is that you can
%       expect to reduce the gradient norm by 8 orders of magnitude
%       (sqrt(eps)) compared to the gradient norm at a "typical" point (a
%       rough initial iterate for example). Further decrease is sometimes
%       possible, but inexact floating point arithmetic will eventually
%       limit the final accuracy. If tolgradnorm is set too low, the
%       algorithm may end up iterating forever (or at least until another
%       stopping criterion triggers).
%   maxiter (1000)
%       The algorithm terminates if maxiter (outer) iterations were executed.
%   maxtime (Inf)
%       The algorithm terminates if maxtime seconds elapsed.
%	miniter (3)
%       Minimum number of outer iterations (used only if useRand is true).
%	mininner (1)
%       Minimum number of inner iterations (for tCG).
%	maxinner (problem.M.dim() : the manifold's dimension)
%       Maximum number of inner iterations (for tCG).
%	Delta_bar (problem.M.typicaldist() or sqrt(problem.M.dim()))
%       Maximum trust-region radius. If you specify this parameter but not
%       Delta0, then Delta0 will be set to 1/8 times this parameter.
%   Delta0 (Delta_bar/8)
%       Initial trust-region radius. If you observe a long plateau at the
%       beginning of the convergence plot (gradient norm VS iteration), it
%       may pay off to try to tune this parameter to shorten the plateau.
%       You should not set this parameter without setting Delta_bar.
%	useRand (false)
%       Set to true if the trust-region solve is to be initiated with a
%       random tangent vector. If set to true, no preconditioner will be
%       used. This option is set to true in some scenarios to escape saddle
%       points, but is otherwise seldom activated.
%	kappa (0.1)
%       Inner kappa convergence tolerance.
%	theta (1.0)
%       Inner theta convergence tolerance.
%	rho_prime (0.1)
%       Accept/reject ratio : if rho is at least rho_prime, the outer
%       iteration is accepted. Otherwise, it is rejected. In case it is
%       rejected, the trust-region radius will have been decreased.
%       To ensure this, rho_prime must be strictly smaller than 1/4.
%   rho_regularization (1e3)
%       Close to convergence, evaluating the performance ratio rho is
%       numerically challenging. Meanwhile, close to convergence, the
%       quadratic model should be a good fit and the steps should be
%       accepted. Regularization lets rho go to 1 as the model decrease and
%       the actual decrease go to zero. Set this option to zero to disable
%       regularization (not recommended). See in-code for the specifics.
%   statsfun (none)
%       Function handle to a function that will be called after each
%       iteration to provide the opportunity to log additional statistics.
%       They will be returned in the info struct. See the generic Manopt
%       documentation about solvers for further information. statsfun is
%       called with the point x that was reached last, after the
%       accept/reject decision. See comment below.
%   stopfun (none)
%       Function handle to a function that will be called at each iteration
%       to provide the opportunity to specify additional stopping criteria.
%       See the generic Manopt documentation about solvers for further
%       information.
%   verbosity (2)
%       Integer number used to tune the amount of output the algorithm
%       generates during execution (mostly as text in the command window).
%       The higher, the more output. 0 means silent. 3 and above includes a
%       display of the options structure at the beginning of the execution.
%   debug (false)
%       Set to true to allow the algorithm to perform additional
%       computations for debugging purposes. If a debugging test fails, you
%       will be informed of it, usually via the command window. Be aware
%       that these additional computations appear in the algorithm timings
%       too.
%   storedepth (20)
%       Maximum number of different points x of the manifold for which a
%       store structure will be kept in memory in the storedb. If the
%       caching features of Manopt are not used, this is irrelevant. If
%       memory usage is an issue, you may try to lower this number.
%       Profiling may then help to investigate if a performance hit was
%       incured as a result.
%
% Notice that statsfun is called with the point x that was reached last,
% after the accept/reject decision. Hence: if the step was accepted, we get
% that new x, with a store which only saw the call for the cost and for the
% gradient. If the step was rejected, we get the same x as previously, with
% the store structure containing everything that was computed at that point
% (possibly including previous rejects at that same point). Hence, statsfun
% should not be used in conjunction with the store to count operations for
% example. Instead, you could use a global variable and increment that
% variable directly from the cost related functions. It is however possible
% to use statsfun with the store to compute, for example, alternate merit
% functions on the point x.
%
% See also: steepestdescent conjugategradient manopt/examples

% This file is part of Manopt: www.manopt.org.
% This code is an adaptation to Manopt of the original GenRTR code:
% RTR - Riemannian Trust-Region
% (c) 2004-2007, P.-A. Absil, C. G. Baker, K. A. Gallivan
% Florida State University
% School of Computational Science
% (http://www.math.fsu.edu/~cbaker/GenRTR/?page=download)
% See accompanying license file.
% The adaptation was executed by Nicolas Boumal.
%
% Change log:
%
%   NB April 3, 2013:
%       tCG now returns the Hessian along the returned direction eta, so
%       that we do not compute that Hessian redundantly: some savings at
%       each iteration. Similarly, if the useRand flag is on, we spare an
%       extra Hessian computation at each outer iteration too, owing to
%       some modifications in the Cauchy point section of the code specific
%       to useRand = true.
%
%   NB Aug. 22, 2013:
%       This function is now Octave compatible. The transition called for
%       two changes which would otherwise not be advisable. (1) tic/toc is
%       now used as is, as opposed to the safer way:
%       t = tic(); elapsed = toc(t);
%       And (2), the (formerly inner) function savestats was moved outside
%       the main function to not be nested anymore. This is arguably less
%       elegant, but Octave does not (and likely will not) support nested
%       functions.
%
%   NB Dec. 2, 2013:
%       The in-code documentation was largely revised and expanded.
%
%   NB Dec. 2, 2013:
%       The former heuristic which triggered when rhonum was very small and
%       forced rho = 1 has been replaced by a smoother heuristic which
%       consists in regularizing rhonum and rhoden before computing their
%       ratio. It is tunable via options.rho_regularization. Furthermore,
%       the solver now detects if tCG did not obtain a model decrease
%       (which is theoretically impossible but may happen because of
%       numerical errors and/or because of a nonlinear/nonsymmetric Hessian
%       operator, which is the case for finite difference approximations).
%       When such an anomaly is detected, the step is rejected and the
%       trust region radius is decreased.
%
%   NB Dec. 3, 2013:
%       The stepsize is now registered at each iteration, at a small
%       additional cost. The defaults for Delta_bar and Delta0 are better
%       defined. Setting Delta_bar in the options will automatically set
%       Delta0 accordingly. In Manopt 1.0.4, the defaults for these options
%       were not treated appropriately because of an incorrect use of the
%       isfield() built-in function.

% Verify that the problem description is sufficient for the solver.
if ~canGetCost(problem)
    warning('manopt:getCost', ...
            'No cost provided. The algorithm will likely abort.');
end
if ~canGetGradient(problem)
    warning('manopt:getGradient', ...
            'No gradient provided. The algorithm will likely abort.');
end
if ~canGetHessian(problem)
    warning('manopt:getHessian:approx', ...
            'No Hessian provided. Using an approximation instead.');
end

% Define some strings for display
tcg_stop_reason = {'negative curvature',...
                   'exceeded trust region',...
                   'reached target residual-kappa',...
                   'reached target residual-theta',...
                   'dimension exceeded',...
                   'model increased'};

% Set local defaults here
localdefaults.verbosity = 2;
localdefaults.maxtime = inf;
localdefaults.miniter = 3;
localdefaults.maxiter = 1000;
localdefaults.mininner = 1;
localdefaults.maxinner = problem.M.dim();
localdefaults.tolgradnorm = 1e-6;
localdefaults.kappa = 0.1;
localdefaults.theta = 1.0;
localdefaults.rho_prime = 0.1;
localdefaults.useRand = false;
localdefaults.rho_regularization = 1e3;

% Merge global and local defaults, then merge w/ user options, if any.
localdefaults = mergeOptions(getGlobalDefaults(), localdefaults);
if ~exist('options', 'var') || isempty(options)
    options = struct();
end
options = mergeOptions(localdefaults, options);

% Set default Delta_bar and Delta0 separately to deal with additional
% logic: if Delta_bar is provided but not Delta0, let Delta0 automatically
% be some fraction of the provided Delta_bar.
if ~isfield(options, 'Delta_bar')
    if isfield(problem.M, 'typicaldist')
        options.Delta_bar = problem.M.typicaldist();
    else
        options.Delta_bar = sqrt(problem.M.dim());
    end
end
if ~isfield(options,'Delta0')
    options.Delta0 = options.Delta_bar / 8;
end

% Check some option values
assert(options.rho_prime < 1/4, ...
        'options.rho_prime must be strictly smaller than 1/4.');
assert(options.Delta_bar > 0, ...
        'options.Delta_bar must be positive.');
assert(options.Delta0 > 0 && options.Delta0 < options.Delta_bar, ...
        'options.Delta0 must be positive and smaller than Delta_bar.');

% It is sometimes useful to check what the actual option values are.
if options.verbosity >= 3
    disp(options);
end

% Create a store database
storedb = struct();

tic();

% If no initial point x is given by the user, generate one at random.
if ~exist('x', 'var') || isempty(x)
    x = problem.M.rand();
end

%% Initializations

% k counts the outer (TR) iterations. The semantic is that k counts the
% number of iterations fully executed so far.
k = 0;

% initialize solution and companion measures: f(x), fgrad(x)
[fx fgradx storedb] = getCostGrad(problem, x, storedb);
norm_grad = problem.M.norm(x, fgradx);

% initialize trust-region radius
Delta = options.Delta0;

% Save stats in a struct array info, and preallocate
% (see http://people.csail.mit.edu/jskelly/blog/?x=entry:entry091030-033941)
if ~exist('used_cauchy', 'var')
    used_cauchy = [];
end
stats = savestats(problem, x, storedb, options, k, fx, norm_grad, Delta);
info(1) = stats;
info(min(10000, options.maxiter+1)).iter = [];

% ** Display:
if options.verbosity == 2
   fprintf(['%3s %3s      %5s                %5s     ',...
            'f: %e   |grad|: %e\n'],...
           '   ','   ','     ','     ', fx, norm_grad);
elseif options.verbosity > 2
   fprintf('************************************************************************\n');
   fprintf('%3s %3s    k: %5s     num_inner: %5s     %s\n',...
           '','','______','______','');
   fprintf('       f(x) : %e       |grad| : %e\n', fx, norm_grad);
   fprintf('      Delta : %f\n', Delta);
end

% **********************
% ** Start of TR loop **
% **********************
while true

	% Start clock for this outer iteration
    tic();

    % Run standard stopping criterion checks
    [stop reason] = stoppingcriterion(problem, x, options, info, k+1);

    % If the stopping criterion that triggered is the tolerance on the
    % gradient norm but we are using randomization, make sure we make at
    % least miniter iterations to give randomization a chance at escaping
    % saddle points.
    if stop == 2 && options.useRand && k < options.miniter
        stop = 0;
    end

    if stop
        if options.verbosity >= 1
            fprintf([reason '\n']);
        end
        break;
    end

    if options.verbosity > 2 || options.debug > 0
        fprintf('************************************************************************\n');
    end

    % *************************
    % ** Begin TR Subproblem **
    % *************************

    % Determine eta0
    if ~options.useRand
        % Pick the zero vector
        eta = problem.M.zerovec(x);
    else
        % Random vector in T_x M (this has to be very small)
        eta = problem.M.lincomb(x, 1e-6, problem.M.randvec(x));
        % Must be inside trust-region
        while problem.M.norm(x, eta) > Delta
            eta = problem.M.lincomb(x, sqrt(sqrt(eps)), eta);
        end
    end

    % solve TR subproblem
    [eta Heta numit stop_inner storedb] = ...
                     tCG(problem, x, fgradx, eta, Delta, options, storedb);
    srstr = tcg_stop_reason{stop_inner};

    % This is only computed for logging purposes, because it may be useful
    % for some user-defined stopping criteria. If this is not cheap for
    % specific application (compared to evaluating the cost), we should
    % reconsider this.
    norm_eta = problem.M.norm(x, eta);

    if options.debug > 0
        testangle = problem.M.inner(x, eta, fgradx) / (norm_eta*norm_grad);
    end

    % If using randomized approach, compare result with the Cauchy point.
    % Convergence proofs assume that we achieve at least the reduction of
    % the Cauchy point. After this if-block, either all eta-related
    % quantities have been changed consistently, or none of them have
    % changed.
    if options.useRand
        used_cauchy = false;
        % Check the curvature,
        [Hg storedb] = getHessian(problem, x, fgradx, storedb);
        g_Hg = problem.M.inner(x, fgradx, Hg);
        if g_Hg <= 0
            tau_c = 1;
        else
            tau_c = min( norm_grad^3/(Delta*g_Hg) , 1);
        end
        % and generate the Cauchy point.
        eta_c  = problem.M.lincomb(x, -tau_c * Delta / norm_grad, fgradx);
        Heta_c = problem.M.lincomb(x, -tau_c * Delta / norm_grad, Hg);

        % Now that we have computed the Cauchy point in addition to the
        % returned eta, we might as well keep the best of them.
        mdle  = fx + problem.M.inner(x, fgradx, eta) ...
                   + .5*problem.M.inner(x, Heta,   eta);
        mdlec = fx + problem.M.inner(x, fgradx, eta_c) ...
                   + .5*problem.M.inner(x, Heta_c, eta_c);
        if mdle > mdlec
            eta = eta_c;
            Heta = Heta_c; % added April 11, 2012
            used_cauchy = true;
        end
    end

	% Compute the retraction of the proposal
	x_prop  = problem.M.retr(x, eta);

	% Compute the function value of the proposal
	[fx_prop storedb] = getCost(problem, x_prop, storedb);

	% Will we accept the proposed solution or not?
    % Check the performance of the quadratic model against the actual cost.
    rhonum = fx - fx_prop;
    rhoden = -problem.M.inner(x, fgradx, eta) ...
             -.5*problem.M.inner(x, eta, Heta);

    % Heuristic -- added Dec. 2, 2013 (NB) to replace the former heuristic.
    % This heuristic is documented in the book by Conn Gould and Toint on
    % trust-region methods, section 17.4.2.
    % rhonum measures the difference between two numbers. Close to
    % convergence, these two numbers are very close to each other, so
    % that computing their difference is numerically challenging: there may
    % be a significant loss in accuracy. Since the acceptance or rejection
    % of the step is conditioned on the ratio between rhonum and rhoden,
    % large errors in rhonum result in a large error in rho, hence in
    % erratic acceptance / rejection. Meanwhile, close to convergence,
    % steps are usually trustworthy and we should transition to a Newton-
    % like method, with rho=1 consistently. The heuristic thus shifts both
    % rhonum and rhoden by a small amount such that far from convergence,
    % the shift is irrelevant and close to convergence, the ratio rho goes
    % to 1, effectively promoting acceptance of the step.
    % The rationale is that close to convergence, both rhonum and rhoden
    % are quadratic in the distance between x and x_prop. Thus, when this
    % distance is on the order of sqrt(eps), the value of rhonum and rhoden
    % is on the order of eps, which is indistinguishable from the numerical
    % error, resulting in badly estimated rho's.
    % For abs(fx) < 1, this heuristic is invariant under offsets of f but
    % not under scaling of f. For abs(fx) > 1, the opposite holds. This
    % should not alarm us, as this heuristic only triggers at the very last
    % iterations if very fine convergence is demanded.
    rho_reg = max(1, abs(fx)) * eps * options.rho_regularization;
    rhonum = rhonum + rho_reg;
    rhoden = rhoden + rho_reg;

    if options.debug > 0
        fprintf('DBG:     rhonum : %e\n', rhonum);
        fprintf('DBG:     rhoden : %e\n', rhoden);
    end

    % This is always true if a linear, symmetric operator is used for the
    % Hessian (approximation) and if we had infinite numerical precision.
    % In practice, nonlinear approximations of the Hessian such as the
    % built-in finite difference approximation and finite numerical
    % accuracy can cause the model to increase. In such scenarios, we
    % decide to force a rejection of the step and a reduction of the
    % trust-region radius. We test the sign of the regularized rhoden since
    % the regularization is supposed to capture the accuracy to which
    % rhoden is computed: if rhoden were negative before regularization but
    % not after, that should not be (and is not) detected as a failure.
    model_decreased = (rhoden >= 0);

    if ~model_decreased
        srstr = [srstr ', model did not decrease']; %#ok<AGROW>
    end

    rho = rhonum / rhoden;

    if options.debug > 0
        m = @(x, eta) ...
          getCost(problem, x, storedb) + ...
          getDirectionalDerivative(problem, x, eta, storedb) + ...
          .5*problem.M.inner(x, getHessian(problem, x, eta, storedb), eta);
        zerovec = problem.M.zerovec(x);
        actrho = (fx - fx_prop) / (m(x, zerovec) - m(x, eta));
        fprintf('DBG:   new f(x) : %e\n', fx_prop);
        fprintf('DBG: actual rho : %e\n', actrho);
        fprintf('DBG:   used rho : %e\n', rho);
    end

    % Choose the new TR radius based on the model performance
    trstr = '   ';
    % If the actual decrease is smaller than 1/4 of the predicted decrease,
    % then reduce the TR radius.
    if rho < 1/4 || ~model_decreased
        trstr = 'TR-';
        Delta = Delta/4;
    % If the actual decrease is at least 3/4 of the precicted decrease and
    % the tCG (inner solve) hit the TR boundary, increase the TR radius.
    elseif rho > 3/4 && (stop_inner == 1 || stop_inner == 2)
        trstr = 'TR+';
        Delta = min(2*Delta, options.Delta_bar);
    end
    % Otherwise, keep the TR radius constant.

    % Choose to accept or reject the proposed step based on the model
    % performance.
    if model_decreased && rho > options.rho_prime
        accept = true;
        accstr = 'acc';
        x = x_prop;
        fx = fx_prop;
        [fgradx storedb] = getGradient(problem, x, storedb);
        norm_grad = problem.M.norm(x, fgradx);
    else
        accept = false;
        accstr = 'REJ';
    end

    % Make sure we don't use too much memory for the store database
    storedb = purgeStoredb(storedb, options.storedepth);

    % k is the number of iterations we have accomplished.
    k = k + 1;

    % Log statistics for freshly executed iteration.
    % Everything after this in the loop is not accounted for in the timing.
    stats = savestats(problem, x, storedb, options, k, fx, norm_grad, ...
                      Delta, info, rho, rhonum, rhoden, accept, numit, ...
                      norm_eta, used_cauchy);
    info(k+1) = stats; %#ok<AGROW>

    % ** Display:
    if options.verbosity == 2,
        fprintf(['%3s %3s   k: %5d     num_inner: %5d     ', ...
        'f: %e   |grad|: %e   %s\n'], ...
        accstr,trstr,k,numit,fx,norm_grad,srstr);
    elseif options.verbosity > 2,
        if options.useRand && used_cauchy,
            fprintf('USED CAUCHY POINT\n');
        end
		fprintf('%3s %3s    k: %5d     num_inner: %5d     %s\n', ...
				accstr, trstr, k, numit, srstr);
		fprintf('       f(x) : %e     |grad| : %e\n',fx,norm_grad);
		if options.debug > 0
			fprintf('      Delta : %f          |eta| : %e\n',Delta,norm_eta);
		end
		fprintf('        rho : %e\n',rho);
    end
    if options.debug > 0,
        fprintf('DBG: cos ang(eta,gradf): %d\n',testangle);
        if rho == 0
            fprintf('DBG: rho = 0, this will likely hinder further convergence.\n');
        end
    end

end  % of TR loop (counter: k)

% Restrict info struct-array to useful part
info = info(1:k+1);

if (options.verbosity > 2) || (options.debug > 0),
   fprintf('************************************************************************\n');
end
if (options.verbosity > 0) || (options.debug > 0)
    fprintf('Total time is %f [s] (excludes statsfun)\n', info(end).time);
end

% Return the best cost reached
cost = fx;

end

% Routine in charge of collecting the current iteration stats
function stats = savestats(problem, x, storedb, options, k, fx, ...
                           norm_grad, Delta, info, rho, rhonum, ...
                           rhoden, accept, numit, norm_eta, used_cauchy)
    stats.iter = k;
    stats.cost = fx;
    stats.gradnorm = norm_grad;
    stats.Delta = Delta;
    if k == 0
        stats.time = toc();
        stats.rho = inf;
        stats.rhonum = NaN;
        stats.rhoden = NaN;
        stats.accepted = true;
        stats.numinner = NaN;
        stats.stepsize = NaN;
        if options.useRand
            stats.cauchy = false;
        end
    else
        stats.time = info(k).time + toc();
        stats.rho = rho;
        stats.rhonum = rhonum;
        stats.rhoden = rhoden;
        stats.accepted = accept;
        stats.numinner = numit;
        stats.stepsize = norm_eta;
        if options.useRand,
          stats.cauchy = used_cauchy;
        end
    end

    % See comment about statsfun above: the x and store passed to statsfun
    % are that of the most recently accepted point after the iteration
    % fully executed.
    stats = applyStatsfun(problem, x, storedb, options, stats);

end


function p = prior_sqinvsinvchi2(varargin)
%PRIOR_SQINVSINVCHI2  Scaled-Inv-Chi^2 prior structure
%
%  Description
%    P = PRIOR_SQINVSINVCHI2('PARAM1', VALUE1, 'PARAM2', VALUE2, ...)
%    creates Scaled-Inv-Chi^2 prior structure for square inverse of
%    the parameter in which the named parameters have the specified
%    values. Any unspecified parameters are set to default values.
%
%    P = PRIOR_SQINVSINVCHI2(P, 'PARAM1', VALUE1, 'PARAM2', VALUE2, ...)
%    modify a prior structure with the named parameters altered
%    with the specified values.
%
%    Parameterisation is done by Bayesian Data Analysis,
%    second edition, Gelman et.al 2004.
%
%    Parameters for Scaled-Inv-Chi^2 [default]
%      s2       - scale squared (variance) [1]
%      nu       - degrees of freedom [4]
%      s2_prior - prior for s2 [prior_fixed]
%      nu_prior - prior for nu [prior_fixed]
%
%  See also
%    PRIOR_*

% Copyright (c) 2000-2001,2010,2012 Aki Vehtari
% Copyright (c) 2010 Jaakko Riihim�ki

% This software is distributed under the GNU General Public
% License (version 3 or later); please refer to the file
% License.txt, included with the software, for details.

  ip=inputParser;
  ip.FunctionName = 'PRIOR_SQINVSINVCHI2';
  ip.addOptional('p', [], @isstruct);
  ip.addParamValue('s2',1, @(x) isscalar(x) && x>0);
  ip.addParamValue('s2_prior',[], @(x) isstruct(x) || isempty(x));
  ip.addParamValue('nu',4, @(x) isscalar(x) && x>0);
  ip.addParamValue('nu_prior',[], @(x) isstruct(x) || isempty(x));
  ip.parse(varargin{:});
  p=ip.Results.p;

  if isempty(p)
    init=true;
    p.type = 'SqInv-S-Inv-Chi2';
  else
    if ~isfield(p,'type') && ~isequal(p.type,'SqInv-S-Inv-Chi2')
      error('First argument does not seem to be a valid prior structure')
    end
    init=false;
  end

  % Initialize parameters
  if init || ~ismember('s2',ip.UsingDefaults)
    p.s2 = ip.Results.s2;
  end
  if init || ~ismember('nu',ip.UsingDefaults)
    p.nu = ip.Results.nu;
  end
  % Initialize prior structure
  if init
    p.p=[];
  end
  if init || ~ismember('s2_prior',ip.UsingDefaults)
    p.p.s2=ip.Results.s2_prior;
  end
  if init || ~ismember('nu_prior',ip.UsingDefaults)
    p.p.nu=ip.Results.nu_prior;
  end

  if init
    % set functions
    p.fh.pak = @prior_sqinvsinvchi2_pak;
    p.fh.unpak = @prior_sqinvsinvchi2_unpak;
    p.fh.lp = @prior_sqinvsinvchi2_lp;
    p.fh.lpg = @prior_sqinvsinvchi2_lpg;
    p.fh.recappend = @prior_sqinvsinvchi2_recappend;
  end

end

function [w, s] = prior_sqinvsinvchi2_pak(p)

  w=[];
  s={};
  if ~isempty(p.p.s2)
    w = log(p.s2);
    s=[s; 'log(SqInv-Sinvchi2.s2)'];
  end
  if ~isempty(p.p.nu)
    w = [w log(p.nu)];
    s=[s; 'log(SqInv-Sinvchi2.nu)'];
  end
end

function [p, w] = prior_sqinvsinvchi2_unpak(p, w)

  if ~isempty(p.p.s2)
    i1=1;
    p.s2 = exp(w(i1));
    w = w(i1+1:end);
  end
  if ~isempty(p.p.nu)
    i1=1;
    p.nu = exp(w(i1));
    w = w(i1+1:end);
  end
end

function lp = prior_sqinvsinvchi2_lp(x, p)

  lJ = -log(x)*3 + log(2);  % log(-2/x^3) log(|J|) of transformation
  xt = x.^-2;               % transformation
  lp = -sum((p.nu./2+1) .* log(xt) + (p.s2.*p.nu./2./xt) + (p.nu/2) .* log(2./(p.s2.*p.nu)) + gammaln(p.nu/2)) + sum(lJ);

  if ~isempty(p.p.s2)
    lp = lp + p.p.s2.fh.lp(p.s2, p.p.s2) + log(p.s2);
  end
  if ~isempty(p.p.nu)
    lp = lp + p.p.nu.fh.lp(p.nu, p.p.nu) + log(p.nu);
  end
end

function lpg = prior_sqinvsinvchi2_lpg(x, p)
  lJg = -3./x;              % gradient of log(|J|) of transformation
  xt  = x.^-2;              % transformation
  xtg = -2/x.^3;             % derivative of transformation
  lpg = xtg.*(-(p.nu/2+1)./xt +p.nu.*p.s2./(2*xt.^2)) + lJg;

  if ~isempty(p.p.s2)
    lpgs2 = (-sum(p.nu/2.*(1./xt-1./p.s2)) + p.p.s2.fh.lpg(p.s2, p.p.s2)).*p.s2 + 1;
    lpg = [lpg lpgs2];
  end
  if ~isempty(p.p.nu)
    lpgnu = (-sum(0.5*(log(xt) + p.s2./xt + log(2./p.s2./p.nu) - 1 + digamma1(p.nu/2))) + p.p.nu.fh.lpg(p.nu, p.p.nu)).*p.nu + 1;
    lpg = [lpg lpgnu];
  end
end

function rec = prior_sqinvsinvchi2_recappend(rec, ri, p)
% The parameters are not sampled in any case.
  rec = rec;
  if ~isempty(p.p.s2)
    rec.s2(ri,:) = p.s2;
  end
  if ~isempty(p.p.nu)
    rec.nu(ri,:) = p.nu;
  end
end


function [y]=funcrs(tt,fun,eps,y,nswp,varargin)
%Cross approximation of a function of a TT-tensor, Method 1
%   [Y]=FUNCRS(TT,FUN,EPS,Y,NSWP)
%   Computes approximation to the function FUN(TT) with accuracy EPS
%   Auxiliary parameters:  Y (initial approximation), NSWP
%   (number of sweeps in the method
%   Much faster then usual cross by vectorized computation of subtensors
%
%
% TT-Toolbox 2.2, 2009-2012
%
%This is TT Toolbox, written by Ivan Oseledets et al.
%Institute of Numerical Mathematics, Moscow, Russia
%webpage: http://spring.inm.ras.ru/osel
%
%For all questions, bugs and suggestions please mail
%ivan.oseledets@gmail.com
%---------------------------

%PARAMETERS SECTION
rmin=1;
dropsweeps = 4;
kick_rank=10;
verb=true;
if (~isempty(y))
   yold=y;
end
%For this procedure we need only local (!) indices, since it
%is more or less equivalent to orthogonalization;
d=tt.d;
ps=tt.ps;
core=tt.core;
n=tt.n;
r=tt.r;

ry=y.r;
psy=y.ps;
cry=y.core;
phi=cell(d+1,1);
phi{d+1}=1;
phi{1}=1;
%Warmup procedure: orthogonalize from right-to-left & maxvol
pos1=psy(d);
%Cores i
for i=d-1:-1:1
   cr=cry(pos1:pos1+ry(i+1)*n(i+1)*ry(i+2)-1);
   cr2=cry(pos1-ry(i)*n(i)*ry(i+1):pos1-1);
   cr2=reshape(cr2,[ry(i)*n(i),ry(i+1)]);
   cr=reshape(cr,[ry(i+1),n(i+1)*ry(i+2)]);
   cr=cr.';
   [cr,rm]=qr(cr,0);
   ry(i+1)=size(cr,2);
   ind=maxvol2(cr);
   r1=cr(ind,:);
   cr=cr/r1;
   cr=cr.';
   cry(pos1:pos1+ry(i+1)*n(i+1)*ry(i+2)-1)=cr(:);
   pos1=pos1-ry(i)*n(i)*ry(i+1);
   cr2=cr2*(r1*rm).';
   cry(pos1:pos1+ry(i)*n(i)*ry(i+1)-1)=cr2(:);
   %Take phi matrix; convolve from right with current cors of V
   cr0=core(ps(i+1):ps(i+2)-1);
   cr0=reshape(cr0,[r(i+1)*n(i+1),r(i+2)]);
   cr0=cr0*phi{i+2}; %cr0 is now r(i)*n(i)*ry(i+1);
   cr0=reshape(cr0,[r(i+1),n(i+1)*ry(i+2)]);
   phi{i+1}=cr0(:,ind);
 %  psy=cumsum([1;n.*ry(1:d).*ry(2:d+1)]);
%y.core=cry;
%y.r=ry;
%y.ps=psy;
%keyboard

   %Orthogonalization & maxvol is "local operation" (touches two
   %cores)
   %Computation of elements requries storage of phi matrices
end
%Truncate cry
%pos1=pos1-n(1)*r(2);
cry=cry(pos1:numel(cry));
y.core=cry;
y.r=ry;
y.ps=psy;
%keyboard
swp=1;

yold=[];
not_converged = true;
pos1=1;
last_sweep = false;
while ( swp < nswp && not_converged )
    max_er=0;
cry_old=cry;
psy=cumsum([1;n.*ry(1:d).*ry(2:d+1)]);
pos1=1;
 for i=1:d-1
     %fprintf('i=%d \n',i);
     %We care for two cores, with number i & number i+1, and use
     %psi(i) and psi(i+2) as a basis; also we will need to recompute
     %psi(i+1)
     ps1=phi{i}; ps2=phi{i+2};
     %Take current core; convolve it with
     %core=core(ps
     %Compute (!) superblock and function (!) of it
     cr1=core(ps(i):ps(i+1)-1);
     cr2=core(ps(i+1):ps(i+2)-1);
     cr1=reshape(cr1,[r(i),n(i)*r(i+1)]);
     cr1=ps1*cr1;
     cr2=reshape(cr2,[r(i+1)*n(i+1),r(i+2)]);
     cr2=cr2*ps2;
     cr1=reshape(cr1,[ry(i)*n(i),r(i+1)]);
     cr2=reshape(cr2,[r(i+1),n(i+1)*ry(i+2)]);
     cr=cr1*cr2;
     cr=reshape(cr,[ry(i)*n(i),n(i+1)*ry(i+2)]);
     cr=fun(cr); %Elements are evaluated here!
     cr=reshape(cr,[ry(i)*n(i),n(i+1)*ry(i+2)]);
     %Check for local approximation of cr for the error
     cry1=cry(pos1:pos1+ry(i)*n(i)*ry(i+1)-1);

     cry2=cry_old(psy(i+1):psy(i+2)-1);

     cry1=reshape(cry1,[ry(i)*n(i),ry(i+1)]);
     cry2=reshape(cry2,[ry(i+1),n(i+1)*ry(i+2)]);
     appr=cry1*cry2;
     er=norm(appr-cr,'fro')/norm(cr,'fro');
     max_er=max(er,max_er);
     %Compute SVD of cr

     if (mod(swp,dropsweeps)~=0)&&(swp>1)&&(~last_sweep)
         [u,s,v]=svd(cr-appr,'econ');
     else
         [u,s,v]=svd(cr,'econ');
     end;
     s=diag(s);
     r2=my_chop2(s,eps*norm(cr,'fro')/sqrt(d-1));
     s=s(1:r2); u=u(:,1:r2); v=v(:,1:r2);

     v=conj(v)*diag(s);

     if (mod(swp,dropsweeps)~=0)&&(swp>1)&&(~last_sweep)
         u = [cry1, u];
         v = [cry2.', v];
         [u,rv]=qr(u,0);
         v = v*(rv.');
     else
         if (~last_sweep)
             %Kick rank of u
             ur=randn(size(u,1),kick_rank);
             %Orthogonalize ur to u by Golub-Kahan reorth
             u=reort(u,ur);
             radd=size(u,2)-r2;
             if ( radd > 0 )
                 vr=zeros(size(v,1),radd);
                 v=[v,vr];
             end
         end;
     end;

     %vr=randn(size(v,1),kick_rank);
     %v=reort(v,vr);
     %radd=size(v,2)-rnew;
     %if ( radd > 0 )
     %    ur=zeros(size(u,1),radd);
     %    u=[u,ur];
     %end
     %rnew=rnew+radd;

     %u=[u,uadd]; v=[v,vadd];
     %[u,ru]=qr(u,0);
     %v=v*(ru.');
     r2=size(u,2);

     ind=maxvol2(u);
     r1=u(ind,:);
     u=u/r1; v=v*r1'; v=v.';
     %Compute new phi
     ry(i+1)=r2;
     cry(pos1:pos1+ry(i)*n(i)*ry(i+1)-1)=u(:);
     pos1=pos1+ry(i)*n(i)*ry(i+1);
     cry(pos1:pos1+ry(i+1)*n(i+1)*ry(i+2)-1)=v(:);

     %Now we have to: cr1 with phi from the left
     phi{i+1}=cr1(ind,:); %phi{i+1}=phi{i+1};

 end
  %Truncate local memory
  cry=cry(1:pos1+ry(d)*n(d)*ry(d+1)-1);
  psy=cumsum([1;n.*ry(1:d).*ry(2:d+1)]);
y.core=cry;
y.r=ry;
y.ps=psy;
%keyboard;
cry_old=cry;
%m=numel(cry);
%cry=[zeros(size(cry)),cry];
%pos1=pos1+m;
%cry2=cry_old(psy(i+1):psy(i+2)-1);
cry=cry(psy(d):psy(d+1)-1); %Start--only two cores
 for i=d-1:-1:1
     %fprintf('i=%d \n',i);
     %We care for two cores, with number i & number i+1, and use
     %psi(i) and psi(i+2) as a basis; also we will need to recompute
     %psi(i+1)
     ps1=phi{i}; ps2=phi{i+2};
     %Take current core; convolve it with
     %core=core(ps
     %Compute (!) superblock and function (!) of it
     cr1=core(ps(i):ps(i+1)-1);
     cr2=core(ps(i+1):ps(i+2)-1);
     cr1=reshape(cr1,[r(i),n(i)*r(i+1)]);
     cr1=ps1*cr1;
     cr2=reshape(cr2,[r(i+1)*n(i+1),r(i+2)]);
     cr2=cr2*ps2;
     cr1=reshape(cr1,[ry(i)*n(i),r(i+1)]);
     cr2=reshape(cr2,[r(i+1),n(i+1)*ry(i+2)]);
     cr=cr1*cr2;
     cr=reshape(cr,[ry(i)*n(i),n(i+1)*ry(i+2)]);
     cr=fun(cr); %Elements are evaluated here!
     cr=reshape(cr,[ry(i)*n(i),n(i+1)*ry(i+2)]);
     %Check for local approximation of cr for the error
     %cry1=cry(pos1-ry(i)*n(i)*ry(i+1):pos1-1);
     cry1=cry_old(psy(i):psy(i+1)-1);
     %cry2=cry(ry(:ry(i+1)*n(i+1)*ry(i+2));
     %cry1=cry(1:ry(i)*n(i)*ry(i+1));
     %cry2=cry(ry(i)*n(i)*ry(i+1)+1:ry(i)*n(i)*ry(i+1)+ry(i+1)*n(i+1)*ry(i+2));
     cry2=cry(1:ry(i+1)*n(i+1)*ry(i+2));
     cry(1:ry(i+1)*n(i+1)*ry(i+2))=[];
     %cry2=cry_old(psy(i+1):psy(i+2)-1);
     %cry(1:(ry(i+1)*n(i+1)*ry(i+2)))=[]; %Delete both of first cores
     %cry(1:ry(i)*n(i)*ry(i+1)+ry(i+1)*n(i+1)*ry(i+2))=[];
     cry1=reshape(cry1,[ry(i)*n(i),ry(i+1)]);
     cry2=reshape(cry2,[ry(i+1),n(i+1)*ry(i+2)]);
     appr=cry1*cry2;
     er=norm(appr-cr,'fro')/norm(cr,'fro');
     %er
     max_er=max(er,max_er);

     if (mod(swp,dropsweeps)~=0)&&(swp>1)&&(~last_sweep)
         [u,s,v]=svd(cr-appr,'econ');
     else
         %Compute SVD of cr
         [u,s,v]=svd(cr,'econ');
     end;
     s=diag(s);
     r2=my_chop2(s,eps*norm(cr,'fro')/sqrt(d-1));
     s=s(1:r2); u=u(:,1:r2); v=conj(v(:,1:r2));
     %Make it standard
     u=u*diag(s);

     if (mod(swp,dropsweeps)~=0)&&(swp>1)&&(~last_sweep)
         u = [cry1, u];
         v = [cry2.', v];
         [v,rv]=qr(v,0);
         u = u*(rv.');
     else
         if (~last_sweep)
             %Kick rank
             vr=randn(size(v,1),kick_rank);
             v=reort(v,vr);
             radd=size(v,2)-r2;
             if ( radd > 0 )
                 ur=zeros(size(u,1),radd);
                 u=[u,ur];
             end
         end;
     end;

     r2=size(v,2);
     ind=maxvol2(v);
     r1=v(ind,:);
     v=v/r1; v=v.';
     u=u*r1';
     %Compute new phi;
     ry(i+1)=r2;
     u=u(:); u=u'; v=v(:); v=v';
     cry=[u,v,cry];
     %keyboard;
     %We need new memory for
     %cry(pos1:pos1+ry(i+1)*n(i+1)*ry(i+2)-1)=v(:); %Here memory has to be (?) enlarged
     %if ( pos1 <= ry(i)*n(i)*ry(i+1) ) %Have to enlarge memory
     %  cry=cry(pos1:numel(cry));
     %  cry=[u(:)',cry];
     %  pos1=ry(i)*n(i)*ry(i+1)+1;
     %else
     %   pos1=pos1-ry(i)*n(i)*ry(i+1);
      %  cry(pos1:pos1+ry(i)*n(i)*ry(i+1)-1)=u(:);
     %end
     %Now we have to: cr1 with phi from the left
     phi{i+1}=cr2(:,ind); phi{i+1}=phi{i+1};

 end
 %Truncate local memory
 %cry=cry(pos1:numel(cry));

% keyboard;
  psy=cumsum([1;n.*ry(1:d).*ry(2:d+1)]);

y.core=cry;
y.r=ry;
y.ps=psy;
%keyboard;
if ( isempty(yold) )
 yold=y;
 er_nrm=1;
else

   er_nrm=norm(yold-y)/norm(y);
   yold=y;
end
if ( verb )
 fprintf('sweep=%d, er=%3.2e er_nrm=%3.2e \n',swp,max_er,er_nrm);
end
if (last_sweep)
    break;
end;
if (er_nrm<eps)
    last_sweep=true;
end;
swp=swp+1;
end
  psy=cumsum([1;n.*ry(1:d).*ry(2:d+1)]);

y.core=cry;
y.r=ry;
y.ps=psy;

end


function Cube = InitCube
% initialize the cube
Cube = [];
SmCube = ...
  [[0 0 0]; [1 0 0]; [0 1 0]; [1 1 0]; ...
   [0 0 1]; [1 0 1]; [0 1 1]; [1 1 1]];
Cube.Nodes = [];
for x=1:3,
  for y=1:3,
    for z=1:3,
      Cube.Nodes(:,1,(z-1)*9+(y-1)*3+x) = SmCube(:,1)*2+2*x-5;
      Cube.Nodes(:,2,(z-1)*9+(y-1)*3+x) = SmCube(:,2)*2+2*y-5;
      Cube.Nodes(:,3,(z-1)*9+(y-1)*3+x) = SmCube(:,3)*2+2*z-5;
    end
  end
end
Cube.Order = [[1 5 7 3];[2 4 8 6];[1 2 6 5];...
              [3 7 8 4];[1 3 4 2];[5 6 8 7]];
NumColor = [[1:3:25];[3:3:27];[1 2 3 10 11 12 19 20 21];...
  [7 8 9 16 17 18 25 26 27];[1:9];[19:27]];
ValColor = [[1 0 0];[1 0 1];[0 0 1];[0 1 0];[1 1 0];[1 1 1]];
Cube.Color = zeros(6,3,27); % black
for x=1:6,
  for y=1:9,
    Cube.Color(x,:,NumColor(x,y)) = ValColor(x,:);
  end
end
Cube.RotCubeCube = ... % for rotate cube
  [[ 7  8  9 16 17 18 25 26 27  4  5  6 13 14 15 22 23 24  1  2  3 10 11 12 19 20 21];...
   [19 20 21 10 11 12  1  2  3 22 23 24 13 14 15  4  5  6 25 26 27 16 17 18  7  8  9];...
   [19 10  1 22 13  4 25 16  7 20 11  2 23 14  5 26 17  8 21 12  3 24 15  6 27 18  9];...
   [ 3 12 21  6 15 24  9 18 27  2 11 20  5 14 23  8 17 26  1 10 19  4 13 22  7 16 25];...
   [ 3  6  9  2  5  8  1  4  7 12 15 18 11 14 17 10 13 16 21 24 27 20 23 26 19 22 25];...
   [ 7  4  1  8  5  2  9  6  3 16 13 10 17 14 11 18 15 12 25 22 19 26 23 20 27 24 21]];
Cube.RotCubeFlat = ... % for rotate cube
  [[1 2 5 6 4 3];...
   [1 2 6 5 3 4];...
   [6 5 3 4 1 2];...
   [5 6 3 4 2 1];...
   [3 4 2 1 5 6];...
   [4 3 1 2 5 6]];
Cube.RotLayerCube = ... % for rotate layer
  [[1:3:25];[3:3:27];[1 2 3 10 11 12 19 20 21];[7 8 9 16 17 18 25 26 27];[1:9];[19:27]];
Cube.NumColor = NumColor;
% Mouse polygons for any rotations
Cube.PgRotCubeX = ... % for rotate the cube around OX, OY, OZ
  [[168;114;171;225;168],[292;292;346;346;292],...
   [232;292;346;287;232],[292;347;347;292;292],...
   [232;171;119;180;232],[119;180;180;119;119],...
   [340;278;225;287;340],[119;180;180;119;119],...
   [240;240;180;180;240],[400;400;346;346;400],...
   [60;60;119;119;60],[292;292;240;240;292]];
Cube.PgRotCubeY(:,:,1) = ...
  [[373;345;323;350;373],[273;214;239;299;273],...
   [296;273;299;324;296],[89;117;176;148;89],...
   [296;323;294;272;296],[107;85;146;170;107],...
   [352;375;350;324;352],[234;209;272;294;234],...
   [123;187;209;146;123],[203;265;239;176;203],...
   [256;191;170;234;256],[148;214;187;123;148]];
Cube.PgRotLayX = ... % for rotate the layer
  [[240;292;232;240],[240;292;292;240],...
   [60;171;114;60],[240;292;292;240],...
   [220;168;225;220],[400;346;346;400],...
   [400;287;346;400],[400;347;347;400],...
   [400;340;287;400],[240;180;180;240],...
   [240;232;180;240],[240;180;180;240],...
   [60;119;171;60],[60;119;119;60],...
   [220;225;278;220],[60;119;119;60],...
   [60;60;119;60],[240;240;292;240],...
   [240;180;240;240],[400;346;400;400],...
   [240;240;180;240],[400;400;347;400],...
   [60;119;60;60],[240;292;240;240]];
Cube.PgRotLayY = ...
  [[246;273;296;246],[60;89;148;60],...
   [319;323;345;319],[246;214;273;246],...
   [398;373;350;398],[329;299;239;329],...
   [329;324;299;329],[142;176;117;142],...
   [329;352;324;329],[246;272;209;246],...
   [246;296;272;246],[60;146;85;60],...
   [319;294;323;319],[133;107;170;133],...
   [398;350;375;398],[319;234;294;319],...
   [319;256;234;319],[246;187;214;246],...
   [246;209;187;246],[329;239;265;329],...
   [60;123;146;60],[142;203;176;142],...
   [133;170;191;133],[60;148;123;60]];
Cube.hand = ...
  [[NaN,NaN,NaN,NaN,NaN,NaN,NaN,  1,  1,NaN,NaN,NaN,NaN,NaN,NaN,NaN];...
   [NaN,NaN,NaN,NaN,  1,  1,  1,  2,  2,  1,  1,  1,NaN,NaN,NaN,NaN];...
   [NaN,NaN,NaN,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,NaN,NaN,NaN];...
   [NaN,NaN,NaN,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,NaN,NaN,NaN];...
   [NaN,NaN,NaN,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  1,  1,NaN];...
   [NaN,  1,  1,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1,  2,  2,  1];...
   [  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1];...
   [NaN,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1];...
   [NaN,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1,NaN];...
   [NaN,NaN,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1,NaN];...
   [NaN,NaN,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  1,NaN]];
return


function a = r8mat_diag_add_scalar ( n, a, s )

%*****************************************************************************80
%
%% R8MAT_DIAG_ADD_SCALAR adds a scalar to the diagonal of a matrix.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    18 April 2005
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer N, the number of rows and columns of the matrix.
%
%    Input, real A(N,N), the matrix to be modified.
%
%    Input, real S, the value to be added to the diagonal
%    of the matrix.
%
%    Output, real A(N,N), the modified matrix.
%
  for i = 1 : n
    a(i,i) = a(i,i) + s;
  end

  return
end


function [rxnInLoops, N, loopInfo] = findMinNull(model, formulation, varargin)
% Find a minimal null-space for all internal cycles by solving a MILP,
% as proposed in Chan et al., 2017.
%
% USAGE:
%    [rxnInLoops, N, loopInfo] = findMinNull(model, formulation, parameters)
%
% INPUT:
%    model:               COBRA model structure
%
% OPTIONAL INPUTS:
%    formulation:         1: solve the MILP problem for a minimal null-space basis
%                            in an interlaced fashion by presolving some relaxed LPs to simultaneously
%                            determine the directions of reactions that participate in internal cycles
%                            (default, the quickest way as of 2017 Nov)
%                         2: directly solve the MILP. Then determine the directions
%                            of reactions that participate in internal cycles
%    parameters:          solver-specific parameter structure or name-value pair argument for solverCobraMILP
%
% OUTPUTS:
%    rxnInLoops:          #rxns-by-2 matrix. rxnInLoops(j, 1) = true => reverse direction of rxn j in loops
%                         rxnInLoops(j, 2) = true => forward direction of rxn j in loops
%    N:                   Minimal feasible null-space matrix for internal cycles
%    loopInfo:            structure containing the following parameters/information:
%                         *.M:          the bound for minimum/maximum flux
%                         *.minFlux:    minimum flux required for a reaction to be active
%                         *.ignoreRxns: rxns with small coefficients that are prechecked
%                                       before solving the MILP to avoid numerical issues
%                         *.nsTime:      wall time for finding the null-space
%                         *.nsCPU:      CPU time for finding the null-space
%                         *.loopPreprocessCPU:  CPU time for finding the directions of reactions participating in loops
%                         *.loopPreprocessTime: wall time for finding the directions of reactions participating in loops

if nargin < 2 || isempty(formulation)
	formulation = 1;
else
	if ~isscalar(formulation)
        % first varargin recognized as formulation
		varargin = [{formulation}; varargin(:)];
		formulation = 1;
	end
end

% filter out reactions with very small stoichiometric coefficents. They are
% usually not in loops. This can accelerate MILP solution time.
% (part of nullspace preprocessing)
cpu0 = cputime;
t = tic;

bigM = 100;
minFlux = 0.1;
rxnIn = sum(model.S ~= 0, 1) > 1;
rxnInIDs = find(rxnIn);
nR = sum(rxnIn);
metIn = any(model.S(:, rxnIn), 2);
nM = sum(metIn);
feasTol = getCobraSolverParams('LP', 'feasTol');

% rxnInLoops(j, 1) = true => reverse direction in loops
% rxnInLoops(j, 2) = true => forward direction in loops
rxnInLoops = false(numel(model.rxns), 2);

nsCPU = cputime - cpu0;
nsTime = toc(t);

if formulation == 1
    % Formulation 1: the procedure stated in SI Methods section 7 in Chan et al. (2017)
    % Determine reactions in loops and their feasible directions in loops
    % in an interlaced fashion

    % find maximum number of active reactions
    % (part of nullspace preprocessing)
    cpu0 = cputime;
    t = tic;

    MILP.A = [model.S(metIn, rxnIn),  sparse(nM, nR * 2); ... Sv = 0
             -speye(nR),             -bigM * speye(nR)    sparse(nR, nR); ...  -v - M z_pos <= -minFlux
              speye(nR),              sparse(nR, nR),    -bigM * speye(nR)];  % v - M z_neg <= -minFlux
    MILP.b = [zeros(nM, 1); -minFlux * ones(nR * 2, 1)];
    MILP.c = [zeros(nR, 1); ones(nR * 2, 1)];

    % Step 2 in SI Methods section 7.
    % Find all irreversible reactions that are in loops
    % fix z_pos = 1 if lb < 0, i.e., v may be negative
    % fix z_neg = 1 if ub > 0, i.e., v may be positive
    % (=> z_pos = z_neg = 1 for truly reversible reactions)
    MILP.lb = [-bigM * (model.lb(rxnIn) < 0); model.lb(rxnIn) < 0; model.ub(rxnIn) > 0];
    MILP.ub = [bigM * (model.ub(rxnIn) > 0); ones(nR * 2, 1)];
    MILP.osense = 1;
    MILP.csense = char(['E' * ones(1, nM), 'L' * ones(1, nR * 2)]);
    sol.irr = solveCobraLP(MILP, varargin{:});

    % Step 3 in SI Methods section 7.
    % Record identified reactions
    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | sol.irr.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | sol.irr.full(1:nR) > 1e-3;

    % Step 4 in SI Methods section 7.
    % Check reversible reactions whose forward direction not yet found to be in loops
    rxnCk = ~rxnInLoops(rxnIn, 2) & (model.lb(rxnIn) < 0 & model.ub(rxnIn) >0);
    % fix irreversible reactions known to be not in loops
    MILP.lb((model.lb(rxnIn) >= 0 | model.ub(rxnIn) <= 0) & ~any(rxnInLoops(rxnIn, :), 2)) = 0;
    MILP.ub((model.lb(rxnIn) >= 0 | model.ub(rxnIn) <= 0) & ~any(rxnInLoops(rxnIn, :), 2)) = 0;
    % fix z_pos = 1 for reactions not in rxnCk
    MILP.lb((nR + 1):end) = [~rxnCk; ones(nR, 1)];
    sol.rev1 = solveCobraLP(MILP, varargin{:});
    % Record identified reactions
    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | sol.rev1.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | sol.rev1.full(1:nR) > 1e-3;

    % Step 5 in SI Methods section 7.
    % Check reversible reactions whose reverse direction not yet found to be in loops
    rxnCk = ~rxnInLoops(rxnIn, 1) & (model.lb(rxnIn) < 0 & model.ub(rxnIn) >0);
    % fix z_neg = 1 for reactions not in rxnCk
    MILP.lb((nR + 1):end) = [ones(nR, 1); ~rxnCk];
    sol.rev2 = solveCobraLP(MILP, varargin{:});
    % Record identified reactions
    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | sol.rev2.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | sol.rev2.full(1:nR) > 1e-3;

    % This step is not mentioned in SI.
    % Filter reactions with small coefficients before solving MILP.
    % They are usually not in loops. This can accelerate MILP solution time.
    % reactions to be checked
    rxnCk = find(any(abs(model.S(:, rxnIn)) > 0 & abs(model.S(:, rxnIn)) < 1e-3, 1)' ...
        & model.lb(rxnIn) < 0 & model.ub(rxnIn) > 0 & ~any(rxnInLoops(rxnIn, :), 2));
    % fix all z_pos and z_neg
    MILP.lb((nR + 1):end) = 1;
    MILP.c(:) = 0;
    for j = 1:numel(rxnCk)
        % check feasibility in the forward direction
        if ~rxnInLoops(rxnInIDs(rxnCk(j)), 2) && MILP.ub(rxnCk(j)) > 0
            bd0 = MILP.lb(rxnCk(j));
            MILP.lb(rxnCk(j)) = minFlux / 10;
            solJ = solveCobraLP(MILP, varargin{:});
            if checkSolFeas(MILP, solJ) <= feasTol
                rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
                rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
            end
            MILP.lb(rxnCk(j)) = bd0;
        end
        % check feasibility in the reverse direction
        if ~rxnInLoops(rxnInIDs(rxnCk(j)), 1) && MILP.lb(rxnCk(j)) < 0
            bd0 = MILP.ub(rxnCk(j));
            MILP.ub(rxnCk(j)) = -minFlux / 10;
            solJ = solveCobraLP(MILP, varargin{:});
            if checkSolFeas(MILP, solJ) <= feasTol
                rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
                rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
            end
            MILP.ub(rxnCk(j)) = bd0;
        end
    end
    ignoreRxns = rxnCk(~any(rxnInLoops(rxnInIDs(rxnCk), :), 2));
    MILP.lb(ignoreRxns) = 0;
    MILP.ub(ignoreRxns) = 0;
    ignoreRxns = rxnInIDs(ignoreRxns);

    % Step 6 in SI Methods section 7.
    % Solve MILP for the forward directions
    rxnCk = ~rxnInLoops(rxnIn, 2) & (model.lb(rxnIn) < 0 & model.ub(rxnIn) >0);
    % fix z_pos for those not rxnCk. Fix all z_neg
    MILP.lb((nR + 1):end) = [~rxnCk; ones(nR, 1)];
    MILP.vartype = char(['C' * ones(1, nR), 'B' * ones(1, nR * 2)]);
    MILP.x0 = [];
    % minimize z_pos + z_neg
    MILP.c((nR + 1):end) = 1;
    sol.rev3 = solveCobraMILP(MILP, varargin{:});
    % Record identified reactions
    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | sol.rev3.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | sol.rev3.full(1:nR) > 1e-3;

    % end of the first part of null-space preprocessing
    nsCPU = nsCPU + (cputime - cpu0);
    nsTime = nsTime + toc(t);

    % Step 7 in SI Methods section 7.
    % Check whether the reactions with negative fluxes in sol.rev3.x can
    % actual have positive fluxes (part of looping reaction preprocessing)
    cpu0 = cputime;
    t = tic;

    rxnCk = find(sol.rev3.full(1:nR) < -1e-3 & rxnCk);
    if ~isempty(rxnCk)
        % fix all z when solving LP
        MILP.lb((nR + 1):end) = 1;
        MILP.c(:) = 0;
        for j = 1:numel(rxnCk)
            if ~rxnInLoops(rxnInIDs(rxnCk(j)), 2)
                bd0 = MILP.lb(rxnCk(j));
                MILP.lb(rxnCk(j)) = minFlux / 10;
                solJ = solveCobraLP(MILP, varargin{:});
                if checkSolFeas(MILP, solJ) <= feasTol
                    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
                    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
                end
                MILP.lb(rxnCk(j)) = bd0;
            end
        end
    end

    % end of the first part of loop preprocessing time
    loopPreprocessCPU = cputime - cpu0;
    loopPreprocessTime = toc(t);

    % Step 8 in SI Methods section 7.
    % Solve MILP for the reverse direction. (2nd part of nullspace preprocessing)
    cpu0 = cputime;
    t = tic;

    rxnCk = ~rxnInLoops(rxnIn, 1) & (model.lb(rxnIn) < 0 & model.ub(rxnIn) >0);
    % fix z_neg for those not rxnCk. Fix all z_pos
    MILP.lb((nR + 1):end) = [ones(nR, 1); ~rxnCk];
    MILP.c((nR + 1):end) = 1;
    sol.rev4 = solveCobraMILP(MILP, varargin{:});

    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | sol.rev4.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | sol.rev4.full(1:nR) > 1e-3;

    nsCPU = nsCPU + (cputime - cpu0);
    nsTime = nsTime + toc(t);

    % Step 9 in SI Methods section 7.
    % Check whether the reactions with positive fluxes in sol.rev4.cont can
    % actual have negative fluxes (part of looping reaction preprocessing)
    cpu0 = cputime;
    t = tic;

    rxnCk = find(sol.rev4.full(1:nR) > 1e-3 & rxnCk);
    if ~isempty(rxnCk)
        % fix all z when solving LP
        MILP.lb((nR + 1):end) = 1;
        MILP.c(:) = 0;
        for j = 1:numel(rxnCk)
            if ~rxnInLoops(rxnInIDs(rxnCk(j)), 1)
                bd0 = MILP.ub(rxnCk(j));
                MILP.ub(rxnCk(j)) = -minFlux / 10;
                solJ = solveCobraLP(MILP, varargin{:});
                if checkSolFeas(MILP, solJ) <= feasTol
                    rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
                    rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
                end
                MILP.ub(rxnCk(j)) = bd0;
            end
        end
    end
    loopPreprocessCPU = loopPreprocessCPU + (cputime - cpu0);
    loopPreprocessTime = loopPreprocessTime + toc(t);
else
    % Formulation 2: directly solve the MILP and find the feasible reaction
    % direction participating in loops for each reaction

    % find maximum number of active reactions (nullspace preprocessing)
    cpu0 = cputime;
    t = tic;

    % filter reactions with small coefficients before solving MILP
    % reactions to be checked
    LP.A = model.S(metIn, rxnIn);
    LP.b = zeros(nM, 1);
    LP.c = zeros(nR, 1);
    LP.lb = -bigM * (model.lb(rxnIn) < 0);
    LP.ub = bigM * (model.ub(rxnIn) > 0);
    LP.csense = char('E' * ones(1, nM));
    [lb0, ub0] = deal(LP.lb, LP.ub);
    % reactions to be checked
    rxnCk = find(any(abs(model.S(:, rxnIn)) > 0 & abs(model.S(:, rxnIn)) < 1e-3, 1));
    for j = 1:numel(rxnCk)
        % check feasibility in the forward direction
        if LP.ub(rxnCk(j)) > 0
            LP.lb(rxnCk(j)) = minFlux / 10;
            solJ = solveCobraLP(LP, varargin{:});
            if checkSolFeas(LP, solJ) <= feasTol
                rxnInLoops(rxnInIDs(rxnCk(j)), 2) = true;
            end
            LP.lb(rxnCk(j)) = lb0(rxnCk(j));
        end
        % check feasibility in the reverse direction
        if ~rxnInLoops(rxnInIDs(rxnCk(j)), 2) && LP.lb(rxnCk(j)) < 0
            LP.ub(rxnCk(j)) = -minFlux / 10;
            solJ = solveCobraLP(LP, varargin{:});
            if checkSolFeas(LP, solJ) <= feasTol
                rxnInLoops(rxnInIDs(rxnCk(j)), 1) = true;
            end
            LP.ub(rxnCk(j)) = ub0(rxnCk(j));
        end
    end
    clear LP
    ignoreRxns = rxnInIDs(rxnCk(~any(rxnInLoops(rxnInIDs(rxnCk), :), 2)));
    rxnIn(ignoreRxns) = false;
    nR = sum(rxnIn);
    metIn = any(model.S(:, rxnIn), 2);
    nM = sum(metIn);
    rxnInIDs = find(rxnIn);

    % solve MILP
    MILP.A = [model.S(metIn, rxnIn),  sparse(nM, nR * 2); ... Sv = 0
             -speye(nR),             -bigM * speye(nR)    sparse(nR, nR); ...  -v - M z_pos <= -1
              speye(nR),              sparse(nR, nR),    -bigM * speye(nR); ...   % v - M z_neg <= -1
              sparse(nR, nR),         speye(nR),          speye(nR)];   % z_pos + z_neg >= 1
    MILP.b = [zeros(nM, 1); -minFlux * ones(nR * 2, 1); ones(nR, 1)];
    MILP.c = [zeros(nR, 1); ones(nR * 2, 1)];
    % fix z_pos = 1 for reactions with ub <= 0. Fix z_neg = 1 for reactions with lb >= 0
    MILP.lb = [-bigM * (model.lb(rxnIn) < 0); model.ub(rxnIn) <= 0; model.lb(rxnIn) >= 0];
    MILP.ub = [bigM * (model.ub(rxnIn) > 0); ones(nR * 2, 1)];
    MILP.osense = 1;
    MILP.csense = char(['E' * ones(1, nM), 'L' * ones(1, nR * 2), 'G' * ones(1, nR)]);
    MILP.vartype = char(['C' * ones(1, nR), 'B' * ones(1, nR * 2)]);
    % z_pos for reactions with lb >= 0 and z_neg for reactions with ub <= 0 can be treated as continuous
    MILP.vartype(nR + find(model.lb(rxnIn) >= 0)) = 'C';
	MILP.vartype(nR * 2 + find(model.ub(rxnIn) <= 0)) = 'C';
    MILP.x0 = [zeros(nR, 1); ones(nR * 2, 1)];
    sol = solveCobraMILP(MILP, varargin{:});
    if isempty(sol.full)
		varargin2 = [{'feasTol'; 1e-7}; varargin(:)];
		sol = solveCobraMILP(MILP, varargin2{:});
    end

    nsCPU = nsCPU + (cputime - cpu0);
    nsTime = nsTime + toc(t);

    cpu0 = cputime;
    t = tic;

    % Check whether the reversible reactions with nonzero fluxes in sol.cont can
    % have fluxes in the opposite sign (part of looping reaction preprocessing)

    rxnInLoops(rxnIn, 1) = sol.full(1:nR) < -1e-3;
    rxnInLoops(rxnIn, 2) = sol.full(1:nR) > 1e-3;

    rxnCk = find(sol.full(1:nR) < -1e-3 | sol.full(1:nR) > 1e-3 & model.lb(rxnIn) < 0 & model.ub(rxnIn) > 0);

    LP.A = model.S(metIn, rxnIn);
    LP.b = zeros(nM, 1);
    LP.c = zeros(nR, 1);
    LP.lb = MILP.lb(1:nR) .* any(rxnInLoops(rxnIn, :), 2);
    LP.ub = MILP.ub(1:nR) .* any(rxnInLoops(rxnIn, :), 2);
    LP.csense = MILP.csense(1:nM);
    clear MILP

    for j = 1:numel(rxnCk)
        if ~rxnInLoops(rxnInIDs(rxnCk(j)), 2)
			bd0 = LP.lb(rxnCk(j));
            LP.lb(rxnCk(j)) = minFlux / 10;
            solJ = solveCobraLP(LP, varargin{:});
            if checkSolFeas(LP, solJ) <= feasTol
                rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
				rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
            end
            LP.lb(rxnCk(j)) = bd0;
        end
        if ~rxnInLoops(rxnInIDs(rxnCk(j)), 1)
			bd0 = LP.ub(rxnCk(j));
            LP.ub(rxnCk(j)) = -minFlux / 10;
            solJ = solveCobraLP(LP, varargin{:});
            if checkSolFeas(LP, solJ) <= feasTol
                rxnInLoops(rxnIn, 1) = rxnInLoops(rxnIn, 1) | solJ.full(1:nR) < -1e-3;
				rxnInLoops(rxnIn, 2) = rxnInLoops(rxnIn, 2) | solJ.full(1:nR) > 1e-3;
            end
            LP.ub(rxnCk(j)) = bd0;
        end
    end

    loopPreprocessCPU = cputime - cpu0;
    loopPreprocessTime = toc(t);
end

% get the nullspace matrix (part of nullspace preprocessing)
cpu0 = cputime;
t = tic;

rxnCyc = any(rxnInLoops, 2);
N = sparseNull(model.S(any(model.S(:, rxnCyc), 2), rxnCyc));
[row, col, entry] = find(N);
rxnCyc = find(rxnCyc);
row = rxnCyc(row);
N = sparse(row, col, entry, size(model.S, 2), size(N, 2));

nsCPU = nsCPU + (cputime - cpu0);
nsTime = nsTime + toc(t);

loopInfo.M = bigM;
loopInfo.minFlux = minFlux;
loopInfo.nsCPU = nsCPU;
loopInfo.nsTime = nsTime;
loopInfo.loopPreprocessCPU = loopPreprocessCPU;
loopInfo.loopPreprocessTime = loopPreprocessTime;
loopInfo.ignoreRxns = model.rxns(ignoreRxns);
end


function [Results, OPTIONS, HeadModel] = bst_wmne(HeadModel,OPTIONS)
% BST_WMNE: Compute the whitened and weighted minimum-norm operator (wMNE imaging kernel)
%
% USAGE:  [Results,OPTIONS] = bst_wmne(HeadModel, OPTIONS) : Compute mininum operator
%                   OPTIONS = bst_wmne()                   : Return default options
%
% DESCRIPTION:
%     This program computes the whitened and weighted minimum-norm operator,
%     (the wMNE imaging kernel), which is used to compute whitened
%     and weighted minimum-norm estimates (MNE).
%         (e.g., J=wMNEoperator*B; where B is the unwhitened data).
%     It can also compute the whitened and noise-normalized dynamic
%     statistical parametric mapping (dSPM) inverse operator, and/or the
%     whitened standardized low resolution brain electromagnetic tomography
%     (sLORETA) inverse operator, which are used to compute whitened source
%     activity dSPM and sLORETA maps.
%         (e.g., S_dSPM=dSPMoperator*B; where B is the unwhitened data).
%         (e.g., S_sLORETA=sLORETAoperator*B; where B is the unwhitened data).
%
%     The function was written with the goal of providing some of the same
%     functionality of the MNE software written by Matti Hamalainen, but no
%     guarantees are made that it performs all computations in the same
%     exact way. It also provides some functionalities not available in the
%     MNE software.
%
% INPUTS:
%    - HeadModel: Array of Brainstorm head model structures
%         |- Gain       : Forward field matrix for all the channels (unconstrained source orientations)
%         |- GridOrient : Dipole orientation matrix
%         |- area       : Vector with the areas (or possibly volumes) associated with the vertices of the source space.
%    - OPTIONS: structure
%         |- NoiseCov      : NoiseCov is the noise covariance matrix.
%         |- ChannelTypes  : Type of each channel (for each row of the Leadfield and the NoiseCov matrix)
%         |- InverseMethod : {'wmne', 'dspm', 'sloreta'}
%         |- SourceOrient  : String or a cell array of strings specifying the type of orientation constraints for each HeadModel (default: 'fixed')
%         |- SNR        : Signal-to noise ratio defined as in MNE (default: 3).
%         |- diagnoise  : Flag to discard off-diagonal elements of NoiseCov (assuming heteroscedastic uncorrelated noise) (default: 0)
%         |- loose      : Value that weights the source variances of the dipole components defining the tangent space of the cortical surfaces (default: []).
%         |- depth      : Flag to do depth weighting (default: 1).
%         |- weightexp  : Order of the depth weighting. {0=no, 1=full normalization, default=0.8}
%         |- weightlimit: Maximal amount depth weighting (default: 10).
%         |- magreg     : Amount of regularization of the magnetometer noise covariance matrix
%         |- gradreg    : Amount of regularization of the gradiometer noise covariance matrix.
%         |- eegreg     : Amount of regularization of the EEG noise covariance matrix.
%         |- ecogreg    : Amount of regularization of the ECOG noise covariance matrix.
%         |- seegreg    : Amount of regularization of the SEEG noise covariance matrix.
%         |- fMRI       : Vector of fMRI values are the source points.
%         |- fMRIthresh : fMRI threshold. The source variances of source points with OPTIONS.fMRI smaller
%         |               than fMRIthresh will be multiplied by OPTIONS.fMRIoff.
%         |- fMRIoff    : Weight assigned to non-active source points according to fMRI and fMRIthresh.
%
% OUTPUTS:
%    - Results : structure with the wMNE inverse operator and possibly the dSPM
%                and/or sLORETA inverse operators, and other information:

% NOTES:
%     - More leadfield matrices can be used: the solution will combine all
%       leadfield matrices appropriately.
%
%     - This leadfield structure allows to combine surface and volume
%       source spaces with and without dipole orientation constraints,and
%       with and without area or volumetric current density computations.
%       If using a single sphere headmodel, the silent radial
%       component could be eliminated using the SVD (e.g., use
%       bst_remove_silent.m).
%
%     - NoisCov: This should be computed from the pre-stimulus period for
%       averaged ERF data (e.g., using MNE), or from an empty room recording
%       for unaveraged spontaneous or induced data.
%
%     - Orientation constrains for dipoles (.SourceOrient field)
%         - "fixed"     : Dipoles constrained to point normal to cortical surfaces
%         - "free"      : No constraints, dipoles point in x,y,z directions
%         - "loose"     : Source variances of dipole components pointing tangentially to the cortical surfaces are multipled by OPTIONS.loose
%         - "truncated" : An SVD of the gain matrix for each source point is
%                         used to remove the dipole component with least variance, which for
%                         the Single Sphere Head Model, corresponds to the radialsilent component).
%       => For dealing with multiple source spaces with different types of orientation constraints use for example,
%          OPTION.SourceOrient{1}='fixed';
%          OPTION.SourceOrient{2}='free';
%          This has to correspond with the HeadModel Structure
%
%     - The output HeadModel param is used here in return to save LOTS of memory in the bst_wmne function,
%       event if it seems to be absolutely useless. Having a parameter in both input and output have the
%       effect in Matlab of passing them "by referece". So please, do NOT remove it from the function description
%
%     - sLORETA: Output values are multiplied by 1e12 for display in Brainstorm (time series and cortical maps).

% @=============================================================================
% This function is part of the Brainstorm software:
% https://neuroimage.usc.edu/brainstorm
%
% Copyright (c) University of Southern California & McGill University
% This software is distributed under the terms of the GNU General Public License
% as published by the Free Software Foundation. Further details on the GPLv3
% license can be found at http://www.gnu.org/copyleft/gpl.html.
%
% FOR RESEARCH PURPOSES ONLY. THE SOFTWARE IS PROVIDED "AS IS," AND THE
% UNIVERSITY OF SOUTHERN CALIFORNIA AND ITS COLLABORATORS DO NOT MAKE ANY
% WARRANTY, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO WARRANTIES OF
% MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, NOR DO THEY ASSUME ANY
% LIABILITY OR RESPONSIBILITY FOR THE USE OF THIS SOFTWARE.
%
% For more information type "brainstorm license" at command prompt.
% =============================================================================@
%
% Copyright (C) 2010 - Rey Rene Ramirez
%
% Authors:  Rey Rene Ramirez, Ph.D.   e-mail: rrramirez at mcw.edu
%           Francois Tadel, 2010-2013

%% ===== DEFINE DEFAULT OPTIONS =====
Def_OPTIONS.NoiseCov      = [];
Def_OPTIONS.InverseMethod = 'wmne';
Def_OPTIONS.ChannelTypes  = {};
Def_OPTIONS.SNR           = 3;
Def_OPTIONS.diagnoise     = 0;
%Def_OPTIONS.SourceOrient= {'free'};
Def_OPTIONS.SourceOrient= {'fixed'};
Def_OPTIONS.loose       = 0.2;
Def_OPTIONS.depth       = 1;
Def_OPTIONS.weightexp   = 0.5;
Def_OPTIONS.weightlimit = 10;
Def_OPTIONS.regnoise    = 1;
Def_OPTIONS.magreg      = .1;
Def_OPTIONS.gradreg     = .1;
Def_OPTIONS.eegreg      = .1;
Def_OPTIONS.ecogreg     = .1;
Def_OPTIONS.seegreg     = .1;
Def_OPTIONS.fMRI        = [];
Def_OPTIONS.fMRIthresh  = [];
Def_OPTIONS.fMRIoff     = 0.1;
Def_OPTIONS.pca         = 1;
% Return the default options
if (nargin == 0)
    Results = Def_OPTIONS;
    return
end
% Make the default for all the leadfields
numL = size(HeadModel,2);
Def_OPTIONS.SourceOrient = repmat(Def_OPTIONS.SourceOrient, [1 numL]);
% Copy default options to OPTIONS structure (do not replace defined values)
OPTIONS = struct_copy_fields(OPTIONS, Def_OPTIONS, 0);

%% ===== CHECK FOR INVALID VALUES =====
% Detect if the input noise covariance matrix is or should be diagonal
C_noise = OPTIONS.NoiseCov;
variances = diag(C_noise);
if isequal(C_noise, diag(variances))
    OPTIONS.diagnoise = 1;
    disp('wMNE> Detected diagonal noise covariance: setting diagnoise to 1');
end
% If OPTIONS.diagnoise is 1, then OPTIONS.pca=0
if OPTIONS.diagnoise
    OPTIONS.pca=0;
    disp('wMNE> If using diagonal noise covariance, PCA option should be off. Setting PCA option off.')
end
if isempty(OPTIONS.NoiseCov)
    error('You need to input the noise covariance in the NoiseCov field of OPTIONS.')
end
if (numL ~= length(OPTIONS.SourceOrient))
    error('The number of elements in the HeadModel structure should equal the length of the cell array OPTIONS.SourceOrient.')
end
if ~isempty(OPTIONS.loose) && (OPTIONS.loose>=1 || OPTIONS.loose<=0)
    error('loose value should be smaller than 1 and bigger than 0, or empty for no loose orientations.')
end
if OPTIONS.weightexp>1 || OPTIONS.weightexp<0
    error('weightexp should be a scalar between 0 and 1')
end
if OPTIONS.magreg>1 || OPTIONS.magreg<0
    error('magreg should be a scalar between 0 and 1')
end
if OPTIONS.eegreg>1 || OPTIONS.eegreg<0
    error('eegreg should be a scalar between 0 and 1')
end
if OPTIONS.ecogreg>1 || OPTIONS.ecogreg<0
    error('ecogreg should be a scalar between 0 and 1')
end
if OPTIONS.seegreg>1 || OPTIONS.seegreg<0
    error('seegreg should be a scalar between 0 and 1')
end

%% ===== NOISE COVARIANCE RANK =====
% Get indices of MEG and EEG channels
iMeg  = find(strncmpi(OPTIONS.ChannelTypes,'MEG',3));
iEeg  = find(strncmpi(OPTIONS.ChannelTypes,'EEG',3));
iEcog = find(strncmpi(OPTIONS.ChannelTypes,'ECOG',3));
iSeeg = find(strncmpi(OPTIONS.ChannelTypes,'SEEG',3));
% Diagonal noisecov
if OPTIONS.diagnoise
    C_noise = diag(variances);
    rnkC_noise_meg = length(iMeg);
    rnkC_noise_eeg = length(iEeg);
    rnkC_noise_ecog = length(iEcog);
    rnkC_noise_seeg = length(iSeeg);
    disp('wMNE> Setting off diagonal elements of the noise covariance to zero.');
    disp(['wMNE> Rank of noise covariance is ' num2str(size(C_noise,1))]);
% Full noisecov
else
    % Estimate noise covariance matrix rank separately for sensor types
    if ~isempty(iMeg)
        rnkC_noise_meg = rank(single(C_noise(iMeg,iMeg))); % Rey added this. Separate rank of MEG. 3/23/11
        disp(['wMNE> Rank of MEG part of noise covariance is ' num2str(rnkC_noise_meg)]);
    end
    if ~isempty(iEeg)
        rnkC_noise_eeg = rank(single(C_noise(iEeg,iEeg))); % Rey added this. Separate rank of EEG. 3/23/11
        disp(['wMNE> Rank of EEG part of noise covariance is ' num2str(rnkC_noise_eeg)]);
    end
    if ~isempty(iEcog)
        rnkC_noise_ecog = rank(single(C_noise(iEcog,iEcog)));  % FT added 21-Feb-13
        disp(['wMNE> Rank of ECOG part of noise covariance is ' num2str(rnkC_noise_ecog)]);
    end
    if ~isempty(iSeeg)
        rnkC_noise_seeg = rank(single(C_noise(iSeeg,iSeeg)));  % FT added 21-Feb-13
        disp(['wMNE> Rank of SEEG part of noise covariance is ' num2str(rnkC_noise_seeg)]);
    end
    % Sets off-diagonal terms to zero. Rey added this. 3/23/11
    C_noise_new = 0 * C_noise;
    if ~isempty(iMeg)
        C_noise_new(iMeg,iMeg) = C_noise(iMeg,iMeg);
    end
    if ~isempty(iEeg)
        C_noise_new(iEeg,iEeg) = C_noise(iEeg,iEeg);
    end
    if ~isempty(iEcog)
        C_noise_new(iEcog,iEcog) = C_noise(iEcog,iEcog);
    end
    if ~isempty(iSeeg)
        C_noise_new(iSeeg,iSeeg) = C_noise(iSeeg,iSeeg);
    end
    C_noise = C_noise_new;
end

%% ===== REGULARIZE NOISE COVARIANCE MATRIX =====
% Only if option is selected
if OPTIONS.regnoise
    listTypes = unique(OPTIONS.ChannelTypes);
    % Loop on all the required data types (MEG MAG, MEG GRAD, EEG)
    for iType = 1:length(listTypes)
        % Get channel indices
        iChan = find(strcmpi(OPTIONS.ChannelTypes, listTypes{iType}));
        % Regularize noise covariance matrix
        switch listTypes{iType}
            case 'MEG GRAD', reg = OPTIONS.gradreg;
            case 'MEG MAG',  reg = OPTIONS.magreg;
            case 'MEG',      reg = OPTIONS.gradreg;
            case 'EEG',      reg = OPTIONS.eegreg;
            case 'ECOG',     reg = OPTIONS.ecogreg;
            case 'SEEG',     reg = OPTIONS.seegreg;
        end
        C_noise(iChan,iChan) = C_noise(iChan,iChan) + (reg * mean(variances(iChan)) * eye(length(iChan)));
    end
end

% FT added 06-Nov-2014
% Make sure the noise covariance matrix is strictly symmetrical
% (the previous operations may cause rounding errors that make the matrix not exactly symmetrical)
if ~isequal(C_noise, C_noise')
    C_noise = tril(C_noise) + (tril(C_noise)' - diag(diag(C_noise)));
end

%% ===== WHITENING OPERATOR =====
% Rey added all of this, 3/23/11
% Modified FT 21-Feb-2013
% Whitening of each modality separately (MEG,EEG,ECOG,SEEG), which assumes
% zero covariance between them (i.e., a block diagonal noise covariance). This
% was recommended by Matti as EEG does not measure all the signals from the same
% environmental noise sources as MEG.
nChan = size(C_noise,2);
W = zeros(0, nChan);
if ~isempty(iMeg)
    W_meg = CalculateWhitener('MEG', C_noise, iMeg, rnkC_noise_meg, OPTIONS.pca);
    W_tmp = zeros(size(W_meg,1), nChan);
    W_tmp(:,iMeg) = W_meg;
    W = [W; W_tmp];
end
if ~isempty(iEeg)
    W_eeg = CalculateWhitener('EEG', C_noise, iEeg, rnkC_noise_eeg, OPTIONS.pca);
    W_tmp = zeros(size(W_eeg,1), nChan);
    W_tmp(:,iEeg) = W_eeg;
    W = [W; W_tmp];
end
if ~isempty(iEcog)
    W_ecog = CalculateWhitener('ECOG', C_noise, iEcog, rnkC_noise_ecog, OPTIONS.pca);
    W_tmp = zeros(size(W_ecog,1), nChan);
    W_tmp(:,iEcog) = W_ecog;
    W = [W; W_tmp];
end
if ~isempty(iSeeg)
    W_seeg = CalculateWhitener('SEEG', C_noise, iSeeg, rnkC_noise_seeg, OPTIONS.pca);
    W_tmp = zeros(size(W_seeg,1), nChan);
    W_tmp(:,iSeeg) = W_seeg;
    W = [W; W_tmp];
end
% Check for whitener integrity
if any(isnan(W(:))) || any(isinf(W(:)))
    error('Invalid noise covariance matrix.')
end
% Display rank of the whitener
rnkC_noise = size(W,1);
display(['wMNE> Total rank is ' num2str(rnkC_noise) '.'])

%% ===== PROCESSING LEAD FIELD MATRICES, WEIGHTS, AREAS & ORIENTATIONS =====
% Initializing.
spl = zeros(numL,1);
numdipcomp = spl;
for k = 1:numL
    sL = size(HeadModel(k).Gain, 2);
    switch OPTIONS.SourceOrient{k}
        case 'fixed',      numdipcomp(k)=1;
        case 'free',       numdipcomp(k)=3;
        case 'loose',      numdipcomp(k)=3;
        case 'truncated',  numdipcomp(k)=2;
    end
    spl(k) = (sL / 3) * numdipcomp(k); % This is a vector with the total number of dipole components per source space.
end
sspl = sum(spl); % This is the total number of dipole components across all source spaces.
L = zeros(rnkC_noise,sspl);
w = ones(sspl,1);
if isfield(HeadModel, 'area')
    areas = w;
else
    areas = [];
end
itangential = [];
start = 0;
Q_Cortex = cell(1,numL);
for k = 1:numL
    start = start + 1;
    endd = start + spl(k) - 1;
    Lk = HeadModel(k).Gain;
    HeadModel(k).Gain = [];
    %% ===== COMPUTE POWER =====
    szL = size(Lk);
    if OPTIONS.depth
        % display(['wMNE> Computing power of gain matrices at each source point for source space ' num2str(k) '.'])
        % Computing power
        wk = squeeze(sum(sum((reshape(Lk,[szL(1) 3 szL(2)/3])) .^2,1),2));
        wk = repmat(wk',[numdipcomp(k) 1]);
        wk = reshape(wk,[spl(k) 1]);
        w(start:endd) = wk;
        clear wk
    end
    switch OPTIONS.SourceOrient{k}
        case 'fixed'
            display('wMNE> Appying fixed dipole orientations.')
            Lk = bst_gain_orient(Lk,HeadModel(k).GridOrient);
        case 'free'
            display('wMNE> Using free dipole orientations. No constraints.')
        case 'loose'
            display('wMNE> Transforming lead field matrix to cortical coordinate system.')
            [Lk, Q_Cortex{k}] = bst_xyz2lf(Lk, HeadModel(k).GridOrient');
            % Getting indices for tangential dipoles.
            itangentialtmp = start:endd;
            itangentialtmp(1:3:end) = [];
            itangential = [itangential itangentialtmp];  %#ok<AGROW>
        case 'truncated'
            display('wMNE> Truncating the dipole component pointing in the direction with least variance (i.e., silent component for single sphere head model.')
            [Lk, Q_Cortex{k}] = bst_remove_silent(Lk);
    end
    %% ===== WHITEN LEAD FIELD MATRIX =====
    % Whiten lead field.
    % display(['wMNE> Whitening lead field matrix for source space ' num2str(k) '.'])
    Lk = W * Lk;
    if isfield(HeadModel(k),'area') && ~isempty(HeadModel(k).area)
        areav = HeadModel(k).area;
        areav = repmat(areav', [numdipcomp(k) 1]);
        areav = reshape(areav, [spl(k) 1]);
        areas(start:endd) = areav;
    end
    L(:,start:endd) = Lk;
    start = endd;
end
% Computing reciprocal of power.
w = 1 ./ w;
% Clear memory
clear Lk endd start itangentialtmp sL szL

%% ===== APPLY AREAS =====
if ~isempty(areas)
    display('wMNE> Applying areas to compute current source density.')
    areas = areas.^2;
    w = w .* areas;
end
clear areas

%% ===== APPLY DEPTH WEIGHTHING =====
if OPTIONS.depth
    % ===== APPLY WEIGHT LIMIT =====
    % Applying weight limit.
    % display('wMNE> Applying weight limit.')
    weightlimit2 = OPTIONS.weightlimit .^ 2;
    %limit=min(w(w>min(w)*weightlimit2));  % This is the Matti way.
    limit = min(w) * weightlimit2;  % This is the Rey way (robust to possible weight discontinuity).
    w(w>limit) = limit;

    % ===== APPLY WEIGHT EXPONENT =====
    % Applying weight exponent.
    % display('wMNE> Applying weight exponent.')
    w = w .^ OPTIONS.weightexp;
    clear limit weightlimit2
end

%% ===== APPLY LOOSE ORIENTATIONS =====
if ~isempty(itangential)
    display(['wMNE> Applying loose dipole orientations. Loose value of ' num2str(OPTIONS.loose) '.'])
    w(itangential) = w(itangential) * (OPTIONS.loose);
end

%% ===== APPLY fMRI PRIORS =====
% Apply fMRI Priors
if ~isempty(OPTIONS.fMRI)
    display('wMNE> Applying fMRI priors.')
    ifmri = (OPTIONS.fMRI < OPTIONS.fMRIthresh);
    w(ifmri) = w(ifmri) * OPTIONS.fMRIoff;
end

%% ===== ADJUSTING SOURCE COVARIANCE MATRIX =====
% Adjusting Source Covariance matrix to make trace of L*C_J*L' equal to number of sensors.
% display('wMNE> Adjusting source covariance matrix.')
C_J = speye(sspl, sspl);
C_J = spdiags(w, 0, C_J);
trclcl = trace(L * C_J * L');
C_J = C_J * (rnkC_noise / trclcl);
Rc = chol(C_J, 'lower');
LW = L * Rc;
clear C_J trclcl sspl itangential rnkC_noise

%% ===== SINGULAR VALUE DECOMPOSITION =====
% Set regularization parameter based on SNR
lambda2 = OPTIONS.SNR ^ (-2);
% Compute SVD.
% display('wMNE> Computing SVD of whitened and weighted lead field matrix.')
[U,S,V] = svd(LW,'econ');
s = diag(S);
ss = s ./ (s.^2 + lambda2);
display(sprintf('wMNE> Effective Max and Mean SNR: %.2f and %.2f', s(1)^2/lambda2, mean(s.^2)/lambda2));
clear LW lambda2 s

%% ===== WHITENED MNE IMAGING KERNEL =====
% Compute whitened MNE operator.
Kernel = Rc * V * diag(ss) * U';
clear Rc V ss U

%% ===== NORMALIZATIONS =====
start = 0;
for k = 1:numL
    start = start + 1;
    endd = start + spl(k) - 1;

    % ===== dSPM OPERATOR =====
    if strcmpi(OPTIONS.InverseMethod, 'dspm')
        dspmdiag = sum(Kernel(start:endd,:) .^2, 2);
        if (numdipcomp(k) == 1)
            dspmdiag = sqrt(dspmdiag);
        elseif (numdipcomp(k)==3 || numdipcomp(k)==2)
            dspmdiag = reshape(dspmdiag, [numdipcomp(k), spl(k)/numdipcomp(k)]);
            dspmdiag = sqrt(sum(dspmdiag)); % Taking trace and sqrt.
            dspmdiag = repmat(dspmdiag, [numdipcomp(k), 1]);
            dspmdiag = reshape(dspmdiag, [spl(k), 1]);
        end
        Kernel(start:endd,:) = bst_bsxfun(@rdivide, Kernel(start:endd,:), dspmdiag);

    % ===== sLORETA OPERATOR =====
    elseif strcmpi(OPTIONS.InverseMethod, 'sloreta')
        if (numdipcomp(k) == 1)
            sloretadiag = sqrt(sum(Kernel(start:endd,:) .* L(:,start:endd)', 2));
            Kernel(start:endd,:) = bst_bsxfun(@rdivide, Kernel(start:endd,:), sloretadiag);
        elseif (numdipcomp(k)==3 || numdipcomp(k)==2)
            for spoint = start:numdipcomp(k):endd
                R = Kernel(spoint:spoint+numdipcomp(k)-1,:) * L(:,spoint:spoint+numdipcomp(k)-1);
                SIR = sqrtm(pinv(R));
                Kernel(spoint:spoint+numdipcomp(k)-1,:) = SIR * Kernel(spoint:spoint+numdipcomp(k)-1,:);
            end
        end
    end

    % ===== LOOSE ORIENTATION: RE-ORIENT COMPONENTS =====
    % WARNING: Changing the orientations of the dipoles from MNE to Brainstorm
    % => Make "Unconstrained" and "loose"/"truncated" models directly comparable
    if ~isempty(Q_Cortex{k})
        % Creating a block diagonal matrix
        N = endd - start + 1;
        Nout = N / numdipcomp(k) * 3;
        iRow = reshape(repmat(reshape(1:Nout,3,[]), numdipcomp(k), 1), 1, []);
        iCol = reshape(repmat(1:N,3,1), 1, []);
        Q_Cortex{k} = sparse(iRow, iCol, Q_Cortex{k}(:));
        % Applying orientations
        Kernel(start:endd,:) = Q_Cortex{k} * Kernel(start:endd,:);
    end

    start=endd;
end
disp(' ');

%% ===== ASSIGN IMAGING KERNEL =====
% Multiply inverse operator and whitening matrix, so no need to whiten data.
Kernel = Kernel * W;
% Return results structure
Results.ImagingKernel = Kernel;
Results.ImageGridAmp  = [];
Results.Whitener      = W;
if (length(numdipcomp) > 1)
    Results.nComponents = 0;
else
    Results.nComponents = numdipcomp;
end

end

%% ==============================================================================
%  ===== HELPER FUNCTIONS =======================================================
%  ==============================================================================
function W = CalculateWhitener(Modality, C_noise, iChannel, rnkC_noise, isPca)
    [V,D] = eig(C_noise(iChannel,iChannel));
    % In some case, eig returns complex values
    if ~isreal(D)
        error('Complex values in the eigvalues...');
    end
    D = diag(D);
    [D,I] = sort(D,'descend');
    V = V(:,I);
    % No PCA case.
    if ~isPca
        display(['wMNE> Not doing PCA for ' Modality '.'])
        D = 1 ./ D;
        W = diag(sqrt(D)) * V';
    % Rey's approach. MNE has been changed to implement this.
    else
        display(['wMNE> Setting small ' Modality ' eigenvalues to zero.'])
        D = 1 ./ D;
        D(rnkC_noise+1:end) = 0;
        W = diag(sqrt(D)) * V';
        W = W(1:rnkC_noise,:); % This line will reduce the actual number of variables in data and leadfield to the true rank. This was not done in the original MNE C code.
    end
end


function [CSF, seg]=mtv_mrQ_Seg_kmeans_simple(T1,BM,M0,mmPerVox,CSFVal)
% function [CSF, seg]=mtv_mrQ_Seg_kmeans_simple(T1,BM,M0)
%
% Clips brain to ventricles area, then segments using k-means with three
% clusters.
%
% This function uses FSL to segment into three tissues. It takes the CSF
% tissue and restricts it by the T1 values. The CSF is also restricted to
% be in the center of the brain in a box of approximately 60mm x 80mm x
% 40mm. Assuming that the brain is in AC-PC space, this is where the
% ventricles should be.
% Assumption:
%       CSFVal: R1 for water at body temperature (in 1/sec).
%                  Default is 0.35.
%
% ~INPUTS~
%       T1: VFA T1.
%       BM: Brain Mask.
%       M0: VFA M0.
%       CSFVal: Threshold on the R1 value of CSF (default=0.35s-1)
%
% ~OUTPUTS~
%       CSF: Refined mask of ventricules (box in the center + smooth)
%       seg: brain clustering (GM=1;DeepGray=2;WM=3;CSF=4)
%
% Adapted from:
% (C) Mezer lab, the Hebrew University of Jerusalem, Israel
%   2015
%
%

%% I. Loading and definitions
if ~exist('CSFVal','var')
    CSFVal=1/0.35; % R1 (in 1/sec) of water at body temp  (minimum value)
end
R1=1./T1;
BM = BM & ~isinf(R1);
%% II. Perform k-means in a "while" loop

    fprintf('\n Performing segmentation for CSF file ...              \n');

seg=zeros(size(R1));

mask= R1>CSFVal & BM;
notdone=0;

while notdone==0

    [IDX,C] =kmeans(R1(mask),3);
    seg(mask)=IDX;

    notdone=1;

    % Check we don't get a strange cluster that is very small in size.
    %(if so, this might be just noise)
    if  length(find(IDX==1))/length(IDX)<0.05
        mask=mask & seg~=1;
        notdone=0;
    end

    if  length(find(IDX==2))/length(IDX)<0.05
        mask=mask & seg~=2;
        notdone=0;
    end

    if  length(find(IDX==3))/length(IDX)<0.05
        mask=mask & seg~=3;
        notdone=0;
    end
end

% check if the clusters' means are too similar
if abs(1-C(1)/C(2))<0.1  &&  abs(1-C(1)/C(3))<0.1
    [IDX,C] =kmeans(R1(mask),1);
elseif abs(1-C(1)/C(2))<0.1
    [IDX,C] =kmeans(R1(mask),2);
elseif abs(1-C(1)/C(3))<0.1
    [IDX,C] =kmeans(R1(mask),2);
elseif abs(1-C(2)/C(3))<0.1
    [IDX,C] =kmeans(R1(mask),2);
end

%% III. Create segmentation file of the clipped brain
seg=zeros(size(R1));
seg(mask)=IDX;

% The tissue with the highest value is white matter (WM), the tissue with
% the lowest value is gray matter (GM), and the tissue with the
% intermediate value is the deep nuclei and the tissue between the WM and
% the GM.
%
% In some segmentations the lowest value is is air, intermediate is
% GM, and highest is WM.
%
% Either way, the order is maintained and we get a segmentation of GM, WM
% and CSF.

[val,idx]=sort(C);
GMclass=1;DEEPclass=2;WMclass=3; CSFclass=4;

seg(seg==idx(1))=4; seg(seg==idx(2))=5;seg(seg==idx(3))=6;
seg(seg==4)=GMclass; seg(seg==5)=DEEPclass; seg(seg==6)=WMclass;

CSF= R1<CSFVal & BM;
seg(CSF)=CSFclass; % any region that is mostly water.

% Clip mask size
boxsize = [30 40 20];
sz=size(CSF); szH=round(sz./2);
XX=boxsize(1)./round(mmPerVox(1));
YY=boxsize(2)./round(mmPerVox(2));
ZZ=boxsize(3)./round(mmPerVox(3));

CSF(szH(1)+XX:end,:,:)=0;
CSF(1:szH(1)-XX,:,:)=0;

CSF(:,1:szH(2)-YY,:)=0;
CSF(:,szH(2)+YY:end,:,:)=0;
if numel(szH)>2
    CSF(:,:,1:max(1,szH(3)-ZZ))=0;
    CSF(:,:,min(end,szH(3)+ZZ):end)=0;
end
%% IV. Smoothe in space
[CSF1] = ordfilt3D(CSF,6);
CSF1=CSF &  CSF1;

CSF2= CSF1 & R1<0.25 & R1>0.2 & M0<prctile(M0(BM),99);

%% V. Some issues

if length(find(CSF2))<200
           fprintf(['\n Warning: We could find only ' num2str(length(find(CSF2))) ' csf voxels. This makes the CSF WF estimation noisy. Consider reducing the CSF T1 threshold.              \n']);
end

% Larger ventricles region.
% This mask is good for cases when CSF ventricles are hard to identify.
% It may reduce the accuracy.
CSF= CSF & R1<0.25 & R1>0.2 & M0<prctile(M0(BM),99);


function [eigvector, eigvalue, bSuccess] = OLPP(W, options, data)
% OLPP: Orthogonal Locality Preserving Projections
%
%       [eigvector, eigvalue, bSuccess] = OLPP(W, options, data)
%
%             Input:
%               data       - Data matrix. Each row vector of fea is a data point.
%               W       - Affinity matrix. You can either call "constructW"
%                         to construct the W, or construct it by yourself.
%               options - Struct value in Matlab. The fields in options
%                         that can be set:
%
%                         Please see OLGE.m for other options.
%
%             Output:
%               eigvector - Each column is an embedding function, for a new
%                           data point (row vector) x,  y = x*eigvector
%                           will be the embedding result of x.
%               eigvalue  - The sorted eigvalue of OLPP eigen-problem.
%
%               bSuccess  - 0 or 1. Indicates whether the OLPP calcuation
%                           is successful. (OLPP needs matrix inverse,
%                           which will lead to eigen-decompose a
%                           non-symmetrical matrix. The caculation precsion
%                           of malab sometimes will cause imaginary numbers
%                           in eigenvectors. It seems that the caculation
%                           precsion of matlab is a little bit random, you
%                           can try again if not successful. More robust
%                           and efficient algorithms are welcome!)
%
%                         Please see OLGE.m for other options.
%
%
%
%
%    Examples:
%
%       fea = rand(50,70);
%       options = [];
%       options.Metric = 'Euclidean';
%       options.NeighborMode = 'KNN';
%       options.k = 5;
%       options.WeightMode = 'HeatKernel';
%       options.t = 1;
%       W = constructW(fea,options);
%       options.PCARatio = 0.99
%       options.ReducedDim = 5;
%       bSuccess = 0
%       while ~bSuccess
%           [eigvector, eigvalue, bSuccess] = OLPP(W, options, fea);
%       end
%       Y = fea*eigvector;
%
%       fea = rand(50,70);
%       gnd = [ones(10,1);ones(15,1)*2;ones(10,1)*3;ones(15,1)*4];
%       options = [];
%       options.Metric = 'Euclidean';
%       options.NeighborMode = 'Supervised';
%       options.gnd = gnd;
%       options.bLDA = 1;
%       W = constructW(fea,options);
%       options.PCARatio = 1;
%       options.ReducedDim = 5;
%       while ~bSuccess
%           [eigvector, eigvalue, bSuccess] = OLPP(W, options, fea);
%       end
%       Y = fea*eigvector;
%
%
%
% See also constructW, LPP, LGE, OLGE.
%
%Reference:
%
%   Deng Cai and Xiaofei He, "Orthogonal Locality Preserving Indexing"
%   The 28th Annual International ACM SIGIR Conference (SIGIR'2005),
%   Salvador, Brazil, Aug. 2005.
%
%   Deng Cai, Xiaofei He, Jiawei Han and Hong-Jiang Zhang, "Orthogonal
%   Laplacianfaces for Face Recognition". IEEE Transactions on Image
%   Processing, vol. 15, no. 11, pp. 3608-3614, November, 2006.
%
%    Written by Deng Cai (dengcai2 AT cs.uiuc.edu), August/2004, Feb/2006,
%                                                   Mar/2007, May/2007

if (~exist('options','var'))
   options = [];
end

[nSmp] = size(data,1);
if size(W,1) ~= nSmp
    error('W and data mismatch!');
end

D = full(sum(W,2));

if isfield(options,'Regu') && options.Regu
    options.ReguAlpha = options.ReguAlpha*sum(D)/length(D);
end

D = sparse(1:nSmp,1:nSmp,D,nSmp,nSmp);

%==========================
% If data is too large, the following centering codes can be commented
%==========================
if isfield(options,'keepMean') && options.keepMean
else
    if issparse(data)
        data = full(data);
    end
    sampleMean = mean(data);
    data = (data - repmat(sampleMean,nSmp,1));
end
%==========================

[eigvector, eigvalue, bSuccess] = OLGE(W, D, options, data);

eigIdx = find(eigvalue < 1e-3);
eigvalue (eigIdx) = [];
eigvector(:,eigIdx) = [];


function [GS, pp, info] = slgmm(X, varargin)
%SLGMM Learns Gaussian Mixture model from samples
%
% $ Syntax $
%   - GS = slgmm(X, ...)
%   - [GS, pp] = slgmm(X, ...)
%   - [GS, pp, info] = slgmm(X, ...)
%
% $ Arguments $
%   - GS:       The Gaussian model struct with mixture weights
%   - pp:       the posteriori of samples to component models
%   - info:     the information of learning process
%   - X:        the sample matrix
%
% $ Description $
%   - GS = slgmm(X, ...) learns Gaussian mixture model from samples
%     in matrix X. The following properties can be specified:
%     \*
%     \t    Table. Properties of Learning Gausssian Mixture Model
%     \h    name         &      description                      \\
%          'method'      & The method using for GMM learning. Currently,
%                          there are only one method available: 'EM',
%                          default = 'EM'.
%          'K'           & The number of initial components in the mixture
%                          model (default = 3)                     \\
%          'update'      & The way of updating (default = 'pass'):
%                          1. 'pass'     Pass-wise update;
%                          2. 'comp'     Component-wise update     \\
%          'cyclecn'     & The ratio of components to be updated in each
%                          cycle for 'comp' update scheme. default = 1.
%          'maxiter'     & The maximum number of iterations
%                          (default = 100) \\
%          'tol'         & The maximum tolerance of posteriori error when
%                          the iteration terminates, (default = 1e-6) \\
%          'verbose'     & whether to display information while iteration,
%                          default = true                              \\
%          'annthres'    & The threshold of annealing in FJ algorithm,
%                          default = 0 (unit = average mixture weight) \\
%          'weights'     & The weights of the samples, default = [],
%                          indicating non-weighted. Weights should be given
%                          by a 1 x n row vector. \\
%          'varform'     & The form of variance: 'univar'|'diagvar'|'covar',
%                          default = 'covar';          \\
%          'sharevar'    & Whether the variance is shared by all models,
%                          default = false;            \\
%          'invparams'   & The parameters for computing inverse of variance
%                          default = {}                \\
%
% $ Remarks $
%   - In current version, the component-wise update scheme is only
%     supported when sharevar is false.
%
% $ History $
%   - Created by Dahua Lin, on Aug 28, 2006
%

%% parse and verify input arguments

if ~isnumeric(X) || ndims(X) ~= 2
    error('sltoolbox:invalidarg', ...
        'X should be a 2D numeric matrix');
end

opts.method = 'EM';
opts.K = 3;
opts.update = 'pass';
opts.cyclecn = 1;
opts.maxiter = 100;
opts.tol = 1e-6;
opts.verbose = true;
opts.annthres = 0;
opts.weights = [];
opts.varform = 'covar';
opts.sharevar = false;
opts.invparams = {};
opts = slparseprops(opts, varargin{:});

n = size(X, 2);

%% initialization by random clustering

if opts.K > n
    error('sltoolbox:sizoverflow', ...
        'The K is larger than the number of samples');
end

init_centerinds = randsample(n, opts.K);
init_centers = X(:, init_centerinds);
initc = slclassify_eucnn(init_centers, X);

%% perform learning based on FMM

estfunctor = {@gmm_estfunc, opts.varform, opts.sharevar, opts.invparams};
evalfunctor = {@gmm_evalfunc};

[GS, cw, pp, info] = slfmm(X, n, estfunctor, evalfunctor, ...
    'method', opts.method, ...
    'update', opts.update, ...
    'cyclecn', opts.cyclecn, ...
    'iter', {'maxiter', opts.maxiter}, ...
    'tol', opts.tol, ...
    'verbose', opts.verbose, ...
    'initc', initc, ...
    'annthres', opts.annthres, ...
    'weights', opts.weights, ...
    'estmode', 'innermul', ...
    'condpmode', 'log', ...
    'manifunc', @gmm_manifunc);

%% post actions

GS.mixweights = cw;

%% Core slot functions

function GS = gmm_estfunc(GS, X, n, W, selinds, varform, sharevar, invparams)

slignorevars(n);

if isempty(selinds)
    GS = slgaussest(X, 'weights', W, ...
        'varform', varform, 'sharevar', sharevar, ...
        'compinv', true, 'invparams', invparams);
else
    if sharevar
        error('sltoolbox:rterror', ...
            'The selective updating scheme is only supported when sharevar is false');
    end

    Wsel = W(selinds, :);
    GSu = slgaussest(X, 'weights', Wsel, ...
        'varform', varform, 'sharevar', sharevar, ...
        'compinv', true, 'invparams', invparams);

    nu = length(selinds);
    switch varform
        case 'univar'
            for idx = 1 : nu
                iu = selinds(idx);
                GS.means(:, iu) = GSu.means(:,idx);
                GS.vars(iu) = GSu.vars(idx);
                GS.invvars(iu) = GSu.invvars(idx);
            end
        case 'diagvar'
            for idx = 1 : nu
                iu = selinds(idx);
                GS.means(:, iu) = GSu.means(:, idx);
                GS.vars(:, iu) = GSu.vars(:, idx);
                GS.invvars(:, iu) = GSu.invvars(:, idx);
            end
        case 'covar'
            for idx = 1 : nu
                iu = selinds(idx);
                GS.means(:, iu) = GSu.means(:, idx);
                GS.covs(:,:, iu) = GSu.covs(:,:, idx);
                GS.invcovs(:,:, iu) = GSu.invcovs(:,:, idx);
            end
    end

end

function condp = gmm_evalfunc(GS, X, n)

slignorevars(n);
condp = slgausspdf(GS, X, 'output', 'log');

function GS = gmm_manifunc(GS0, op, varargin)

switch op
    case 'select'

        selinds = varargin{1};

        tyi = slgausstype(GS0);

        GS.dim = GS0.dim;
        GS.nmodels = length(selinds);
        GS.means = GS0.means(:, selinds);

        switch tyi.varform
            case {'univar', 'diagvar'}
                if tyi.sharevar
                    GS.vars = GS0.vars;
                    GS.invvars = GS0.invvars;
                else
                    GS.vars = GS0.vars(:, selinds);
                    GS.invvars = GS0.invvars(:, selinds);
                end
            case 'covar'
                if tyi.sharevar
                    GS.covs = GS0.covs;
                    GS.invcovs = GS0.invcovs;
                else
                    GS.covs = GS0.covs(:,:,selinds);
                    GS.invcovs = GS0.invcovs(:,:,selinds);
                end
        end

    otherwise
        error('sltoolbox:invalidarg', ...
            'Unsupported manipulation option for GMM: %s', op);
end