function [vp,elbo,elbo_sd,exitflag,output,samples,optimState,stats,vp_train] = ...
    vbmc(fun,x0,LB,UB,PLB,PUB,options,varargin)
%VBMC Posterior and model inference via Variational Bayesian Monte Carlo (v1.0.12).
%   VBMC computes a variational approximation of the full posterior and a
%   lower bound on the normalization constant (marginal likelhood or model
%   evidence) for a provided unnormalized log posterior. As of v1.0, VBMC
%   also supports noisy evaluations of the log posterior (see below).
%
%   VP = VBMC(FUN,X0,LB,UB) initializes the variational posterior in the
%   proximity of X0 (ideally, a posterior mode) and iteratively computes
%   a variational approximation for a given target log posterior FUN.
%   FUN accepts input X and returns the value of the target log-joint, that
%   is the unnormalized log-posterior density, at X. LB and UB define a set
%   of strict lower and upper bounds coordinate vector, X, so that the
%   posterior has support on LB < X < UB. LB and UB can be scalars or
%   vectors. If scalars, the bound is replicated in each dimension. Use
%   empty matrices for LB and UB if no bounds exist. Set LB(i) = -Inf and
%   UB(i) = Inf if the i-th coordinate is unbounded (while other coordinates
%   may be bounded). Note that if LB and UB contain unbounded variables,
%   the respective values of PLB and PUB need to be specified (see below).
%   VBMC returns a variational posterior solution VP, which can then be
%   manipulated via other functions in the VBMC toolbox (see examples below).
%
%   VP = VBMC(FUN,X0,LB,UB,PLB,PUB) specifies a set of plausible lower and
%   upper bounds such that LB < PLB < PUB < UB. Both PLB and PUB
%   need to be finite. PLB and PUB represent a "plausible" range, which
%   should denote a region of high posterior probability mass. Among other
%   things, the plausible box is used to draw initial samples and to set
%   priors over hyperparameters of the algorithm. When in doubt, we found
%   that setting PLB and PUB using the topmost ~68% percentile range of the
%   prior (e.g, mean +/- 1 SD for a Gaussian prior) works well in many
%   cases (but note that additional information might afford a better guess).
%
%   VP = VBMC(FUN,X0,LB,UB,PLB,PUB,OPTIONS) performs variational inference
%   with the default parameters replaced by values in the structure OPTIONS.
%   VBMC('defaults') returns the default OPTIONS struct.
%
%   VP = VBMC(FUN,X0,LB,UB,PLB,PUB,OPTIONS,...) passes additional arguments
%   to FUN.
%
%   VBMC also supports noisy/stochastic estimates of the log-posterior,
%   obtained through techniques such as Inverse Binomial Sampling (see
%   examples and references below). For noisy evaluations, FUN should
%   return as second agument the estimated SD (standard deviation) of the
%   log-likelihood noise at X.
%   Set OPTIONS.SpecifyTargetNoise = 1 to activate support for noisy
%   inference (this is not automatic).
%
%   VP = VBMC(FUN,VP0,...) uses variational posterior VP0 (from a previous
%   run of VBMC) to initialize the current run. You can leave PLB and PUB
%   empty, in which case they will be set using VP0 (recommended).
%
%   [VP,ELBO] = VBMC(...) returns an estimate of the ELBO, the variational
%   expected lower bound on the log marginal likelihood (log model evidence).
%   This estimate is computed via Bayesian quadrature.
%
%   [VP,ELBO,ELBO_SD] = VBMC(...) returns the standard deviation of the
%   estimate of the ELBO, as computed via Bayesian quadrature. Note that
%   this standard deviation is *not* representative of the error between the
%   ELBO and the true log marginal likelihood.
%
%   [VP,ELBO,ELBO_SD,EXITFLAG] = VBMC(...) returns an EXITFLAG that describes
%   the exit condition. Possible values of EXITFLAG and the corresponding
%   exit conditions are
%
%    1  Change in the variational posterior, in the ELBO and its uncertainty
%       have reached a satisfactory level of stability across recent
%       iterations, suggesting convergence of the variational solution.
%    0  Maximum number of function evaluations or iterations reached. Note
%       that the returned solution has *not* stabilized.
%
%   [VP,ELBO,ELBO_SD,EXITFLAG,OUTPUT] = VBMC(...) returns a structure OUTPUT
%   with the following information:
%          function: <Target probability density function name>
%        iterations: <Total iterations>
%         funccount: <Total function evaluations>
%          bestiter: <Iteration of returned solution>
%      trainsetsize: <Size of training set for returned solution>
%        components: <Number of mixture components of returned solution>
%            rindex: <Reliability index (< 1 is good)>
% convergencestatus: <"probable" or "no" convergence>
%          overhead: <Fractional overhead (total runtime / total fcn time - 1)>
%          rngstate: <Status of random number generator>
%         algorithm: <Variational Bayesian Monte Carlo>
%           message: <VBMC termination message>
%              elbo: <Estimated ELBO for returned solution>
%           elbo_sd: <Estimated standard deviation of ELBO at returned solution>
%           retried: <"yes", "no", or "failed" if a retry run was performed>
%
%   OPTIONS = VBMC('defaults') returns a basic default OPTIONS structure.
%
%   EXITFLAG = VBMC('test') runs a battery of tests. Here EXITFLAG is 0 if
%   everything works correctly.
%
%   Examples:
%     FUN can be a function handle (using @)
%       vp = vbmc(@rosenbrock_test, ...)
%     In this case, F = rosenbrock_test(X) returns the scalar log pdf F of
%     the target pdf evaluated at X.
%
%     An example with no hard bounds, only plausible bounds
%       plb = [-5 -5]; pub = [5 5]; options.Plot = 'on';
%       [vp,elbo,elbo_sd] = vbmc(@rosenbrock_test,[0 0],[],[],plb,pub,options);
%
%     FUN can also be an anonymous function:
%        lb = [0 0]; ub = [pi 5]; plb = [0.1 0.1]; pub = [3 4]; options.Plot = 'on';
%        vp = vbmc(@(x) 3*sin(x(1))*exp(-x(2)),[1 1],lb,ub,plb,pub,options)
%
%   See VBMC_EXAMPLES for an extended tutorial with more examples.
%   The most recent version of the algorithm and additional documentation
%   can be found here: https://github.com/acerbilab/vbmc
%   Also, check out the FAQ: https://github.com/acerbilab/vbmc/wiki
%
%   References (please cite both):
%
%   1) Acerbi, L. (2018). "Variational Bayesian Monte Carlo". In Advances
%      in Neural Information Processing Systems 31 (NeurIPS 2018), pp. 8213-8223.
%   2) Acerbi, L. (2020). "Variational Bayesian Monte Carlo with Noisy
%      Likelihoods". In Advances in Neural Information Processing Systems 33
%      (NeurIPS 2020).
%
%   Additional references:
%
%   3) Acerbi, L. (2019). "An Exploration of Acquisition and Mean Functions
%      in Variational Bayesian Monte Carlo". In Proc. Machine Learning
%      Research 96: 1-10. 1st Symposium on Advances in Approximate Bayesian
%      Inference, Montréal, Canada.
%   4) van Opheusden, B.*, Acerbi, L.* & Ma, W. J. (2020). "Unbiased and
%      Efficient Log-Likelihood Estimation with Inverse Binomial Sampling".
%      PLoS Computational Biology 16(12): e1008483. (* equal contribution)
%
%   See also VBMC_EXAMPLES, VBMC_KLDIV, VBMC_MODE, VBMC_MOMENTS, VBMC_MTV,
%   VBMC_PDF, VBMC_RND, VBMC_DIAGNOSTICS, @.

%--------------------------------------------------------------------------
% VBMC: Variational Bayesian Monte Carlo for posterior and model inference.
% To be used under the terms of the GNU General Public License
% (http://www.gnu.org/copyleft/gpl.html).
%
%   Author (copyright): Luigi Acerbi, 2018-2022
%   e-mail: luigi.acerbi@{helsinki.fi,gmail.com}
%   URL: http://luigiacerbi.com
%   Version: 1.0.12
%   Release date: Oct 26, 2022
%   Code repository: https://github.com/acerbilab/vbmc
%--------------------------------------------------------------------------

%% Start timer

t0 = tic;
vbmc_version = '1.0.12';

% Check that all VBMC subfolders are on the MATLAB path
add2path();

%% Basic default options
defopts.Display                 = 'iter         % Level of display ("iter", "notify", "final", or "off")';
defopts.Plot                    = 'off          % Plot marginals of variational posterior at each iteration';
defopts.MaxIter                 = '50*(2+nvars) % Max number of iterations';
defopts.MaxFunEvals             = '50*(2+nvars) % Max number of target fcn evals';
defopts.FunEvalsPerIter         = '5            % Number of target fcn evals per iteration';
defopts.TolStableCount          = '60           % Required stable fcn evals for termination';
defopts.RetryMaxFunEvals        = '0            % Max number of target fcn evals on retry (0 = no retry)';
defopts.MinFinalComponents      = '50           % Number of variational components to refine posterior at termination';
defopts.SpecifyTargetNoise      = 'no           % Target log joint function returns noise estimate (SD) as second output';

%% If called with no arguments or with 'defaults', return default options
if nargout <= 1 && (nargin == 0 || (nargin == 1 && ischar(fun) && strcmpi(fun,'defaults')))
    if nargin < 1
        fprintf('Basic default options returned (type "help vbmc" for help).\n');
    end
    vp = defopts;
    return;
end

%% If called with one argument which is 'test', run test
if nargout <= 1 && (nargin == 1 || nargin == 2) && ischar(fun) && strcmpi(fun,'test')
    % Can run a test with a specific OPTIONS struct (otherwise use defaults)
    if nargin == 2; options = x0; else; options = []; end
    vp = runtest_vbmc(options);
    return;
end

%% If called with one argument which is 'version', return version
if nargout <= 1 && nargin == 1 && ischar(fun) && strcmpi(fun,'version')
    vp = vbmc_version;
    return;
end

%% Advanced options (do not modify unless you *know* what you are doing)

defopts.UncertaintyHandling     = '[]           % Explicit noise handling';
defopts.IntegerVars             = '[]           % Array with indices of integer variables';
defopts.NoiseSize               = '[]           % Base observation noise magnitude (standard deviation)';
defopts.MaxRepeatedObservations = '0            % Max number of consecutive repeated measurements for noisy inputs';
defopts.RepeatedAcqDiscount     = '1            % Multiplicative discount on acquisition fcn to repeat measurement at the same location';
defopts.FunEvalStart            = '10*ceil((D+1)/10) % Number of initial target fcn evals';
defopts.SGDStepSize             = '0.005        % Base step size for stochastic gradient descent';
defopts.SkipActiveSamplingAfterWarmup  = 'no    % Skip active sampling the first iteration after warmup';
defopts.RankCriterion           = 'yes          % Use ranking criterion to pick best non-converged solution';
defopts.TolStableEntropyIters   = '6            % Required stable iterations to switch entropy approximation';
defopts.VariableMeans           = 'yes          % Use variable component means for variational posterior';
defopts.VariableWeights         = 'yes          % Use variable mixture weight for variational posterior';
defopts.WeightPenalty           = '0.1          % Penalty multiplier for small mixture weights';
defopts.Diagnostics             = 'off          % Run in diagnostics mode, get additional info';
defopts.OutputFcn               = '[]           % Output function';
defopts.TolStableExcptFrac      = '0.2          % Fraction of allowed exceptions when computing iteration stability';
defopts.Fvals                   = '[]           % Evaluated fcn values at X0';
defopts.OptimToolbox            = '[]           % Use Optimization Toolbox (if empty, determine at runtime)';
defopts.ProposalFcn             = '[]           % Weighted proposal fcn for uncertainty search';
defopts.NonlinearScaling   = 'on                % Automatic nonlinear rescaling of variables';
defopts.SearchAcqFcn       = '@acqf_vbmc        % Fast search acquisition fcn(s)';
defopts.NSsearch           = '2^13              % Samples for fast acquisition fcn eval per new point';
defopts.NSent              = '@(K) 100*K.^(2/3) % Total samples for Monte Carlo approx. of the entropy';
defopts.NSentFast          = '0                 % Total samples for preliminary Monte Carlo approx. of the entropy';
defopts.NSentFine          = '@(K) 2^12*K       % Total samples for refined Monte Carlo approx. of the entropy';
defopts.NSentBoost         = '@(K) 200*K.^(2/3) % Total samples for Monte Carlo approx. of the entropy (final boost)';
defopts.NSentFastBoost     = '[]                % Total samples for preliminary Monte Carlo approx. of the entropy (final boost)';
defopts.NSentFineBoost     = '[]                % Total samples for refined Monte Carlo approx. of the entropy (final boost)';
defopts.NSentActive        = '@(K) 20*K.^(2/3)  % Total samples for Monte Carlo approx. of the entropy (active sampling)';
defopts.NSentFastActive    = '0                 % Total samples for preliminary Monte Carlo approx. of the entropy (active sampling)';
defopts.NSentFineActive    = '@(K) 200*K        % Total samples for refined Monte Carlo approx. of the entropy (active sampling)';
defopts.NSelbo             = '@(K) 50*K         % Samples for fast approximation of the ELBO';
defopts.NSelboIncr         = '0.1               % Multiplier to samples for fast approx. of ELBO for incremental iterations';
defopts.ElboStarts         = '2                 % Starting points to refine optimization of the ELBO';
defopts.NSgpMax            = '80                % Max GP hyperparameter samples (decreases with training points)';
defopts.NSgpMaxWarmup      = '8                 % Max GP hyperparameter samples during warmup';
defopts.NSgpMaxMain        = 'Inf               % Max GP hyperparameter samples during main algorithm';
defopts.WarmupNoImproThreshold = '20 + 5*nvars  % Fcn evals without improvement before stopping warmup';
defopts.WarmupCheckMax     = 'yes               % Also check for max fcn value improvement before stopping warmup';
defopts.StableGPSampling   = '200 + 10*nvars    % Force stable GP hyperparameter sampling (reduce samples or start optimizing)';
defopts.StableGPvpK        = 'Inf               % Force stable GP hyperparameter sampling after reaching this number of components';
defopts.StableGPSamples    = '0                 % Number of GP samples when GP is stable (0 = optimize)';
defopts.GPSampleThin       = '5                 % Thinning for GP hyperparameter sampling';
defopts.GPTrainNinit       = '1024              % Initial design points for GP hyperparameter training';
defopts.GPTrainNinitFinal  = '64                % Final design points for GP hyperparameter training';
defopts.GPTrainInitMethod  = 'rand              % Initial design method for GP hyperparameter training';
defopts.GPTolOpt           = '1e-5              % Tolerance for optimization of GP hyperparameters';
defopts.GPTolOptMCMC       = '1e-2              % Tolerance for optimization of GP hyperparameters preliminary to MCMC';
defopts.GPTolOptActive     = '1e-4              % Tolerance for optimization of GP hyperparameters during active sampling';
defopts.GPTolOptMCMCActive = '1e-2              % Tolerance for optimization of GP hyperparameters preliminary to MCMC during active sampling';
defopts.TolGPVar           = '1e-4              % Threshold on GP variance used by regulatized acquisition fcns';
defopts.TolGPVarMCMC       = '1e-4              % Threshold on GP variance, used to stabilize sampling';
defopts.gpMeanFun          = 'negquad           % GP mean function';
defopts.gpIntMeanFun       = '0                 % GP integrated mean function';
defopts.KfunMax            = '@(N) N.^(2/3)     % Max variational components as a function of training points';
defopts.Kwarmup            = '2                 % Variational components during warmup';
defopts.AdaptiveK          = '2                 % Added variational components for stable solution';
defopts.HPDFrac            = '0.8               % High Posterior Density region (fraction of training inputs)';
defopts.ELCBOImproWeight   = '3                 % Uncertainty weight on ELCBO for computing lower bound improvement';
defopts.TolLength          = '1e-6              % Minimum fractional length scale';
defopts.CacheSize          = '500               % Size of cache for storing fcn evaluations';
defopts.CacheFrac          = '0.5               % Fraction of search points from starting cache (if nonempty)';
defopts.StochasticOptimizer = 'adam             % Stochastic optimizer for varational parameters';
defopts.TolFunStochastic   = '1e-3              % Stopping threshold for stochastic optimization';
defopts.MaxIterStochastic  = '100*(2+nvars)     % Max iterations for stochastic optimization';
defopts.TolSD              = '0.1               % Tolerance on ELBO uncertainty for stopping (iff variational posterior is stable)';
defopts.TolsKL             = '0.01*sqrt(nvars)  % Stopping threshold on change of variational posterior per training point';
defopts.TolStableWarmup    = '15                % Number of stable fcn evals for stopping warmup';
defopts.VariationalSampler = 'malasample        % MCMC sampler for variational posteriors';
defopts.TolImprovement     = '0.01              % Required ELCBO improvement per fcn eval before termination';
defopts.KLgauss            = 'yes               % Use Gaussian approximation for symmetrized KL-divergence b\w iters';
defopts.TrueMean           = '[]                % True mean of the target density (for debugging)';
defopts.TrueCov            = '[]                % True covariance of the target density (for debugging)';
defopts.MinFunEvals        = '5*nvars           % Min number of fcn evals';
defopts.MinIter            = 'nvars             % Min number of iterations';
defopts.HeavyTailSearchFrac = '0.25               % Fraction of search points from heavy-tailed variational posterior';
defopts.MVNSearchFrac      = '0.25              % Fraction of search points from multivariate normal';
defopts.HPDSearchFrac      = '0                 % Fraction of search points from multivariate normal fitted to HPD points';
defopts.BoxSearchFrac      = '0.25              % Fraction of search points from uniform random box based on training inputs';
defopts.SearchCacheFrac    = '0                 % Fraction of search points from previous iterations';
defopts.AlwaysRefitVarPost = 'no                % Always fully refit variational posterior';
defopts.Warmup             = 'on                % Perform warm-up stage';
defopts.WarmupOptions      = '[]                % Special OPTIONS struct for warmup stage';
defopts.StopWarmupThresh   = '0.2               % Stop warm-up when ELCBO increase below threshold (per fcn eval)';
defopts.WarmupKeepThreshold = '10*nvars         % Max log-likelihood difference for points kept after warmup';
defopts.WarmupKeepThresholdFalseAlarm = '100*(nvars+2) % Max log-likelihood difference for points kept after a false-alarm warmup stop';
defopts.StopWarmupReliability = '100            % Reliability index required to stop warmup';
defopts.SearchOptimizer    = 'cmaes             % Optimization method for active sampling';
defopts.SearchCMAESVPInit  = 'yes               % Initialize CMA-ES search SIGMA from variational posterior';
defopts.SearchCMAESbest    = 'no                % Take bestever solution from CMA-ES search';
defopts.SearchMaxFunEvals  = '500*(nvars+2)     % Max number of acquisition fcn evaluations during search';
defopts.MomentsRunWeight   = '0.9               % Weight of previous trials (per trial) for running avg of variational posterior moments';
defopts.GPRetrainThreshold = '1                 % Upper threshold on reliability index for full retraining of GP hyperparameters';
defopts.ELCBOmidpoint      = 'on                % Compute full ELCBO also at best midpoint';
defopts.GPSampleWidths     = '5                 % Multiplier to widths from previous posterior for GP sampling (Inf = do not use previous widths)';
defopts.HypRunWeight       = '0.9               % Weight of previous trials (per trial) for running avg of GP hyperparameter covariance';
defopts.WeightedHypCov     = 'on                % Use weighted hyperparameter posterior covariance';
defopts.TolCovWeight       = '0                 % Minimum weight for weighted hyperparameter posterior covariance';
defopts.GPHypSampler       = 'slicesample       % MCMC sampler for GP hyperparameters';
defopts.CovSampleThresh    = '10                % Switch to covariance sampling below this threshold of stability index';
defopts.DetEntTolOpt       = '1e-3              % Optimality tolerance for optimization of deterministic entropy';
defopts.EntropySwitch      = 'off               % Switch from deterministic entropy to stochastic entropy when reaching stability';
defopts.EntropyForceSwitch = '0.8               % Force switch to stochastic entropy at this fraction of total fcn evals';
defopts.DetEntropyAlpha    = '0                 % Alpha value for lower/upper deterministic entropy interpolation';
defopts.UpdateRandomAlpha  = 'no                % Randomize deterministic entropy alpha during active sample updates';
defopts.AdaptiveEntropyAlpha = 'no              % Online adaptation of alpha value for lower/upper deterministic entropy interpolation';
defopts.DetEntropyMinD     = '5                 % Start with deterministic entropy only with this number of vars or more';
defopts.TolConLoss         = '0.01              % Fractional tolerance for constraint violation of variational parameters';
defopts.BestSafeSD         = '5                 % SD multiplier of ELCBO for computing best variational solution';
defopts.BestFracBack       = '0.25              % When computing best solution, lacking stability go back up to this fraction of iterations';
defopts.TolWeight          = '1e-2              % Threshold mixture component weight for pruning';
defopts.PruningThresholdMultiplier = '@(K) 1/sqrt(K)   % Multiplier to threshold for pruning mixture weights';
defopts.AnnealedGPMean     = '@(N,NMAX) 0       % Annealing for hyperprior width of GP negative quadratic mean';
defopts.ConstrainedGPMean  = 'no                % Strict hyperprior for GP negative quadratic mean';
defopts.TolGPNoise         = 'sqrt(1e-5)        % Minimum GP observation noise';
defopts.GPLengthPriorMean  = 'sqrt(D/6)         % Prior mean over GP input length scale (in plausible units)';
defopts.GPLengthPriorStd   = '0.5*log(1e3)      % Prior std over GP input length scale (in plausible units)';
defopts.UpperGPLengthFactor = '0                % Upper bound on GP input lengths based on plausible box (0 = ignore)';
defopts.InitDesign         = 'plausible         % Initial samples ("plausible" is uniform in the plausible box)';
defopts.gpQuadraticMeanBound = 'yes             % Stricter upper bound on GP negative quadratic mean function';
defopts.Bandwidth          = '0                 % Bandwidth parameter for GP smoothing (in units of plausible box)';
defopts.FitnessShaping     = 'no                % Heuristic output warping ("fitness shaping")';
defopts.OutwarpThreshBase  = '10*nvars          % Output warping starting threshold';
defopts.OutwarpThreshMult  = '1.25              % Output warping threshold multiplier when failed sub-threshold check';
defopts.OutwarpThreshTol   = '0.8               % Output warping base threshold tolerance (fraction of current threshold)';
defopts.Temperature        = '1                 % Temperature for posterior tempering (allowed values T = 1,2,3,4)';
defopts.SeparateSearchGP   = 'no                % Use separate GP with constant mean for active search';
defopts.NoiseShaping       = 'no                % Discount observations from from extremely low-density regions';
defopts.NoiseShapingThreshold = '10*nvars       % Threshold from max observed value to start discounting';
defopts.NoiseShapingFactor = '0.05              % Proportionality factor of added noise wrt distance from threshold';
defopts.AcqHedge           = 'no                % Hedge on multiple acquisition functions';
defopts.AcqHedgeIterWindow = '4                 % Past iterations window to judge acquisition fcn improvement';
defopts.AcqHedgeDecay      = '0.9               % Portfolio value decay per function evaluation';
defopts.ActiveVariationalSamples = '0           % MCMC variational steps before each active sampling';
defopts.ScaleLowerBound    = 'yes               % Apply lower bound on variational components scale during variational sampling';
defopts.ActiveSampleVPUpdate = 'no              % Perform variational optimization after each active sample';
defopts.ActiveSampleGPUpdate = 'no              % Perform GP training after each active sample';
defopts.ActiveSampleFullUpdatePastWarmup = '2   % # iters past warmup to continue update after each active sample';
defopts.ActiveSampleFullUpdateThreshold = '3    % Perform full update during active sampling if stability above threshold';
defopts.VariationalInitRepo = 'no               % Use previous variational posteriors to initialize optimization';
defopts.SampleExtraVPMeans = '0                 % Extra variational components sampled from GP profile';
defopts.OptimisticVariationalBound = '0         % Uncertainty weight on ELCBO during active sampling';
defopts.ActiveImportanceSamplingVPSamples   = '100 % # importance samples from smoothed variational posterior';
defopts.ActiveImportanceSamplingBoxSamples  = '100 % # importance samples from box-uniform centered on training inputs';
defopts.ActiveImportanceSamplingMCMCSamples = '100 % # importance samples through MCMC';
defopts.ActiveImportanceSamplingMCMCThin    = '1   % Thinning for importance sampling MCMC';
defopts.ActiveSamplefESSThresh  = '1            % fractional ESS threhsold to update GP and VP';
defopts.ActiveImportanceSamplingfESSThresh = '0.9 % % fractional ESS threhsold to do MCMC while active importance sampling';
defopts.ActiveSearchBound  = '2                  % Active search bound multiplier';
defopts.TolBoundX          = '1e-5              % Tolerance on closeness to bound constraints (fraction of total range)';
defopts.RecomputeLCBmax    = 'yes              % Recompute LCB max for each iteration based on current GP estimate';
defopts.BoundedTransform   = 'logit            % Input transform for bounded variables';
defopts.WarpEveryIters     = '5                 % Warp every this number of iterations';
defopts.IncrementalWarpDelay = 'yes             % Increase delay between warpings';
defopts.WarpTolReliability = '3                 % Threshold on reliability index to perform warp';
defopts.WarpRotoScaling    = 'yes               % Rotate and scale input';
defopts.WarpCovReg         = '0                 % Regularization weight towards diagonal covariance matrix for N training inputs';
defopts.WarpRotoCorrThresh = '0.05              % Threshold on correlation matrix for roto-scaling';
defopts.WarpMinK           = '5                 % Min number of variational components to perform warp';
defopts.WarpUndoCheck      = 'yes               % Immediately undo warp if not improving ELBO';
defopts.WarpTolImprovement = '0.1               % Improvement of ELBO required to keep a warp proposal';
defopts.WarpTolSDMultiplier = '2                % Multiplier tolerance of ELBO SD after warp proposal';
defopts.WarpTolSDBase      = '1                 % Base tolerance on ELBO SD after warp proposal';

%% Advanced options for unsupported/untested features (do *not* modify)
defopts.WarpNonlinear      = 'off               % Nonlinear input warping';
defopts.ELCBOWeight        = '0                 % Uncertainty weight during ELCBO optimization';
defopts.VarParamsBack      = '0                 % Check variational posteriors back to these previous iterations';
defopts.AltMCEntropy       = 'no                % Use alternative Monte Carlo computation for the entropy';
defopts.VarActiveSample    = 'no                % Variational active sampling';
defopts.FeatureTest        = 'no                % Test a new experimental feature';
defopts.BOWarmup           = 'no                % Bayesian-optimization-like warmup stage';
defopts.gpOutwarpFun       = '[]                % GP default output warping function';

%% If called with 'all', return all default options
if strcmpi(fun,'all')
    vp = defopts;
    return;
end

%% Input arguments

if nargin < 3 || isempty(LB); LB = -Inf; end
if nargin < 4 || isempty(UB); UB = Inf; end
if nargin < 5; PLB = []; end
if nargin < 6; PUB = []; end
if nargin < 7; options = []; end

%% Initialize display printing options

if ~isfield(options,'Display') || isempty(options.Display)
    options.Display = defopts.Display;
end

switch lower(options.Display(1:min(end,3)))
    case {'not'}                        % notify
        prnt = 1;
    case {'no','non','off'}             % none
        prnt = 0;
    case {'ite','all','on','yes'}       % iter
        prnt = 3;
    case {'fin','end'}                  % final
        prnt = 2;
    otherwise
        prnt = 3;
end

%% Initialize variables and algorithm structures

if isempty(x0)
    if prnt > 2
        fprintf('X0 not specified. Taking the number of dimensions from PLB and PUB...');
    end
    if isempty(PLB) || isempty(PUB)
        error('vbmc:UnknownDims', ...
            'If no starting point is provided, PLB and PUB need to be specified.');
    end
    x0 = NaN(size(PLB));
    if prnt > 2
        fprintf(' D = %d.\n', numel(x0));
    end
end

% Initialize from variational posterior
if vbmc_isavp(x0)
    init_from_vp_flag = true;
    vp0 = x0;
    [x0,LB,UB,PLB,PUB,Xvp] = initFromVP(vp0,LB,UB,PLB,PUB,prnt);
else
    init_from_vp_flag = false;
end

D = size(x0,2);     % Number of variables
optimState = [];

% Setup algorithm options
options = setupoptions_vbmc(D,defopts,options);
if options.Warmup
    options_main = options;
    % Use special options during Warmup
    if isfield(options,'WarmupOptions')
        WarmupOptions = options.WarmupOptions;
        % Copy these fields to avoid re-update in SETUPOPTIONS_VBMC
        copyfields = {'MaxFunEvals','TolStableCount','ActiveSampleGPUpdate','ActiveSampleVPUpdate','SearchAcqFcn'};
        for f = copyfields
            if ~isfield(WarmupOptions,f{:})
                WarmupOptions.(f{:}) = options.(f{:});
            end
        end
        options = setupoptions_vbmc(D,options,WarmupOptions);
    end
end

if init_from_vp_flag    % Finish initialization from variational posterior
    x0 = [x0; robustSampleFromVP(vp0,options.FunEvalStart-1,Xvp)];
    clear Xvp vp0;
end

% Check/fix boundaries and starting points
[x0,LB,UB,PLB,PUB] = boundscheck_vbmc(x0,LB,UB,PLB,PUB,prnt);

% Convert from char to function handles
if ischar(fun); fun = str2func(fun); end

% Setup and transform variables, prepare OPTIMSTATE settings struct
K = options.Kwarmup;
[vp,optimState] = ...
    setupvars_vbmc(x0,LB,UB,PLB,PUB,K,optimState,options,prnt);

% Store target density function
optimState.fun = fun;
funwrapper = @(u_) fun(u_,varargin{:});

% Get information from acquisition function(s)
optimState.acqInfo = getAcqInfo(options.SearchAcqFcn);

% GP struct and GP hyperparameters
gp = [];        hypstruct = [];     hypstruct_search = [];

% Initialize function logger
[~,optimState] = funlogger_vbmc([],D,optimState,'init',options.CacheSize);

if optimState.Cache.active
    displayFormat = ' %5.0f     %5.0f  /%5.0f   %12.2f  %12.2f  %12.2f     %4.0f %10.3g       %s\n';
    displayFormat_warmup = ' %5.0f     %5.0f  /%5.0f   %s\n';
elseif optimState.UncertaintyHandlingLevel > 0 && options.MaxRepeatedObservations > 0
    displayFormat = ' %5.0f       %5.0f %5.0f %12.2f  %12.2f  %12.2f     %4.0f %10.3g     %s\n';
    displayFormat_warmup = ' %5.0f       %5.0f    %12.2f  %s\n';
else
    displayFormat = ' %5.0f      %5.0f   %12.2f %12.2f %12.2f     %4.0f %10.3g     %s\n';
    displayFormat_warmup = ' %5.0f       %5.0f    %12.2f  %s\n';
end
if prnt > 2
    if optimState.UncertaintyHandlingLevel > 0
        fprintf('Beginning variational optimization assuming NOISY observations of the log-joint.\n');
    else
        fprintf('Beginning variational optimization assuming EXACT observations of the log-joint.\n');
    end

    if optimState.Cache.active
        fprintf(' Iteration f-count/f-cache    Mean[ELBO]     Std[ELBO]     sKL-iter[q]   K[q]  Convergence    Action\n');
    else
        if options.BOWarmup
            fprintf(' Iteration   f-count     Max[f]     Action\n');
        elseif optimState.UncertaintyHandlingLevel > 0 && options.MaxRepeatedObservations > 0
            fprintf(' Iteration   f-count (x-count)   Mean[ELBO]     Std[ELBO]     sKL-iter[q]   K[q]  Convergence  Action\n');
        else
            fprintf(' Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action\n');
        end
    end
end

%% Variational optimization loop
iter = 0;
isFinished_flag = false;
exitflag = 0;   output = [];    stats = [];

while ~isFinished_flag
    t_iter = tic;
    timer = timer_init();   % Initialize iteration timer

    iter = iter + 1;
    optimState.iter = iter;
    vp_old = vp;
    action = '';
    optimState.redoRotoscaling = false;

    if iter == 1 && optimState.Warmup; action = 'start warm-up'; end

    % Switch to stochastic entropy towards the end if still on deterministic
    if optimState.EntropySwitch && ...
            optimState.funccount >= options.EntropyForceSwitch*options.MaxFunEvals
        optimState.EntropySwitch = false;
        if isempty(action); action = 'entropy switch'; else; action = [action ', entropy switch']; end
    end

    %% Input warping / reparameterization
    if options.IncrementalWarpDelay
        WarpDelay = options.WarpEveryIters*max(1,optimState.WarpingCount);
    else
        WarpDelay = options.WarpEveryIters;
    end

    DoWarping = (options.WarpRotoScaling || options.WarpNonlinear) && ...
        iter > 1 && ~optimState.Warmup && ...
        (iter - optimState.LastWarping) > WarpDelay && ...
        vp.K >= options.WarpMinK && stats.rindex(iter-1) < options.WarpTolReliability && ...
        vp.D > 1;
        % (stats.stable(iter-1) || optimState.funccount >= options.MaxFunEvals*2/3);

    if DoWarping
        t = tic;
        [vp_tmp,~,~,idx_best] = ...
            best_vbmc(stats,iter-1,options.BestSafeSD,options.BestFracBack,options.RankCriterion,0);

        % Store variables in case warp needs to be undone
        optimState_old = optimState;
        gp_old = gp;
        hypstruct_old = hypstruct;
        elbo_old = elbo;
        elbo_sd_old = elbo_sd;

        % Compute input warping
        [trinfo_warp,optimState,warp_action] = warp_input_vbmc(vp_tmp,optimState,stats.gp(idx_best),options);

        % Update GP hyperparameters and variational posterior
        [vp,hypstruct.hyp] = warp_gpandvp_vbmc(trinfo_warp,vp,gp);

        if isempty(action); action = warp_action; else; action = [action ', ' warp_action]; end

        timer.warping = timer.warping + toc(t);

        if options.WarpUndoCheck
            % Train GP
            t = tic;
            [gp,hypstruct,~,optimState] = ...
                gptrain_vbmc(hypstruct,optimState,stats,options);
            timer.gpTrain = timer.gpTrain + toc(t);

            % Optimize variational parameters
            t = tic;
            if ~vp.optimize_mu  % Variational components fixed to training inputs
                vp.mu = gp.X';
                Knew = size(vp.mu,2);
            else
                % Update number of variational mixture components
                Knew = vp.K;
            end

            % Decide number of fast/slow optimizations
            Nfastopts = ceil(evaloption_vbmc(options.NSelbo,Knew));
            Nslowopts = options.ElboStarts; % Full optimizations

            % Run optimization of variational parameters
            vp = vpoptimize_vbmc(Nfastopts,Nslowopts,vp,gp,Knew,optimState,options,prnt);
            optimState.vpK = vp.K;
            optimState.H = vp.stats.entropy;   % Save current entropy

            % Compute ELBO from real variational posterior (might differ from training posterior)
            vp_real = vptrain2real(vp,0,options);
            elbo = vp_real.stats.elbo;
            elbo_sd = vp_real.stats.elbo_sd;

            timer.variationalFit = timer.variationalFit + toc(t);

            % Compute symmetrized KL-divergence between old and new posteriors
            %Nkl = 1e5;
            %sKL = max(0,0.5*sum(vbmc_kldiv(vp,vp_old,Nkl,options.KLgauss)))

            %[elbo elbo_old]
            %[elbo_sd elbo_sd_old]

            % Keep warping only if it substantially improves ELBO
            % and uncertainty does not blow up too much
            if (elbo < (elbo_old + options.WarpTolImprovement)) || ...
                (elbo_sd > (elbo_sd_old*options.WarpTolSDMultiplier + options.WarpTolSDBase))

                % Undo input warping
                vp = vp_old;
                gp = gp_old;
                optimState = optimState_old;
                hypstruct = hypstruct_old;

                % Still keep track of failed warping (failed warp counts twice)
                optimState.WarpingCount = optimState.WarpingCount + 2;
                optimState.LastWarping = optimState.iter;

                if isempty(action); action = 'undo'; else; action = [action ', undo']; end
            end

        end
    end

    %% Actively sample new points into the training set
    t = tic;
    optimState.trinfo = vp.trinfo;
    if iter == 1; new_funevals = options.FunEvalStart; else; new_funevals = options.FunEvalsPerIter; end
    if optimState.Xn > 0
        optimState.ymax = max(optimState.y(optimState.X_flag));
    end
    if optimState.SkipActiveSampling
        optimState.SkipActiveSampling = false;
    else
        if ~isempty(gp) && options.SeparateSearchGP && ~options.VarActiveSample
            % Train a distinct GP for active sampling
            if mod(iter,2) == 0
                meantemp = optimState.gpMeanfun;
                optimState.gpMeanfun = 'const';
                [gp_search,hypstruct_search] = gptrain_vbmc(hypstruct_search,optimState,stats,options);
                optimState.gpMeanfun = meantemp;
            else
                gp_search = gp;
            end
        else
            gp_search = gp;
        end
        % Performe active sampling
        if options.VarActiveSample % Unused
            % FIX TIMER HERE IF USING THIS
            [optimState,vp,t_active,t_func] = ...
                variationalactivesample_vbmc(optimState,new_funevals,funwrapper,vp,vp_old,gp_search,options);
        else
            optimState.hypstruct = hypstruct;
            [optimState,vp,gp,timer] = ...
                activesample_vbmc(optimState,new_funevals,funwrapper,vp,vp_old,gp_search,stats,timer,options);
            hypstruct = optimState.hypstruct;
        end
    end
    optimState.N = optimState.Xn;  % Number of training inputs
    optimState.Neff = sum(optimState.nevals(optimState.X_flag));

    %% Train GP
    t = tic;
    [gp,hypstruct,Ns_gp,optimState] = ...
        gptrain_vbmc(hypstruct,optimState,stats,options);
    timer.gpTrain = timer.gpTrain + toc(t);

    % Check if reached stable sampling regime
    if Ns_gp == options.StableGPSamples && optimState.StopSampling == 0
        optimState.StopSampling = optimState.N;
    end

%     if ~exist('wsabi_hyp','var'); wsabi_hyp = zeros(1,D+1); end
%     priorMu = (optimState.PLB + optimState.PUB)/2;
%     priorVar = diag(optimState.PUB - optimState.PLB);
%     kernelVar = diag(exp(wsabi_hyp(2:end)));
%     lambda = exp(wsabi_hyp(1));
%     hypVar = [1e4,4*ones(1,D)];
%     [log_mu,log_Var,~,~,~,wsabi_hyp] = wsabi_oneshot(...
%         'L',priorMu,priorVar,kernelVar,lambda,0.8,gp.X,gp.y,hypVar);
%     log_mu

    %% Optimize variational parameters
    t = tic;

    if ~vp.optimize_mu  % Variational components fixed to training inputs
        vp.mu = gp.X';
        Knew = size(vp.mu,2);
    else
        % Update number of variational mixture components
        Knew = updateK(optimState,stats,options);
    end

    % Decide number of fast/slow optimizations
    Nfastopts = ceil(evaloption_vbmc(options.NSelbo,K));

    if optimState.RecomputeVarPost || options.AlwaysRefitVarPost
        Nslowopts = options.ElboStarts; % Full optimizations
        optimState.RecomputeVarPost = false;
    else
        % Only incremental change from previous iteration
        Nfastopts = ceil(Nfastopts * options.NSelboIncr);
        Nslowopts = 1;
    end

    % Run optimization of variational parameters
    if optimState.Warmup && options.BOWarmup
        vp_fields = {'elbo','elbo_sd','G','H','varG','varH'};
        for i = 1:numel(vp_fields); vp.stats.(vp_fields{i}) = NaN; end
        varss = NaN;
        pruned = 0;
%     elseif Knew == vp.K && ~optimState.Warmup && vp.K >= 10
%         [vp,varss] = vpoptimizeweights_vbmc(vp,gp,optimState,options,prnt);
%         pruned = 0;
    else
        [vp,varss,pruned] =  ...
            vpoptimize_vbmc(Nfastopts,Nslowopts,vp,gp,Knew,optimState,options,prnt);
        optimState.vp_repo{end+1} = get_vptheta(vp);
    end

    optimState.vpK = vp.K;
    optimState.H = vp.stats.entropy;   % Save current entropy

    % Get real variational posterior (might differ from training posterior)
    vp_real = vptrain2real(vp,0,options);
    elbo = vp_real.stats.elbo;
    elbo_sd = vp_real.stats.elbo_sd;

    timer.variationalFit = timer.variationalFit + toc(t);

    %% Plot current iteration (to be improved)
    if options.Plot
        vbmc_iterplot(vp,gp,optimState,stats,elbo);
    end

    %hh = [gp.post.hyp];
    %exp(hh(gp.Ncov+gp.Nnoise+2:end,:))

    %mubar
    %Sigma

    %----------------------------------------------------------------------
    %% Finalize iteration
    t = tic;

    % Compute symmetrized KL-divergence between old and new posteriors
    Nkl = 1e5;
    sKL = max(0,0.5*sum(vbmc_kldiv(vp,vp_old,Nkl,options.KLgauss)));
    % mtv = vbmc_mtv(vp,vp_old,Nkl)

    % Evaluate max LCB of GP prediction on all training inputs
    [~,~,fmu,fs2] = gplite_pred(gp,gp.X,gp.y,gp.s2);
    optimState.lcbmax = max(fmu - options.ELCBOImproWeight*sqrt(fs2));

    if options.AdaptiveEntropyAlpha
        % Evaluate deterministic entropy
        Hl = entlb_vbmc(vp,0,0);
        Hu = entub_vbmc(vp,0,0);
        optimState.entropy_alpha = max(0,min(1,(vp.stats.entropy - Hl)/(Hu - Hl)));
        optimState.entropy_alpha
    end

    % Compare variational posterior's moments with ground truth
    if ~isempty(options.TrueMean) && ~isempty(options.TrueCov) ...
        && all(isfinite(options.TrueMean(:))) ...
        && all(isfinite(options.TrueCov(:)))

        [mubar_orig,Sigma_orig] = vbmc_moments(vp_real,1,1e6);
        [kl(1),kl(2)] = mvnkl(mubar_orig,Sigma_orig,options.TrueMean,options.TrueCov);
        sKL_true = 0.5*sum(kl)
    else
        sKL_true = [];
    end

    % Record moments in transformed space
    [mubar,Sigma] = vbmc_moments(vp,0);
    if isempty(optimState.RunMean) || isempty(optimState.RunCov)
        optimState.RunMean = mubar(:);
        optimState.RunCov = Sigma;
        optimState.LastRunAvg = optimState.N;
        % optimState.RunCorrection = 1;
    else
        Nnew = optimState.N - optimState.LastRunAvg;
        wRun = options.MomentsRunWeight^Nnew;
        optimState.RunMean = wRun*optimState.RunMean + (1-wRun)*mubar(:);
        optimState.RunCov = wRun*optimState.RunCov + (1-wRun)*Sigma;
        optimState.LastRunAvg = optimState.N;
        % optimState.RunT = optimState.RunT + 1;
    end

    % t_fits(iter) = toc(timer_fits);
    % dt = (t_active(iter)+t_fits(iter))/new_funevals;

    timer.finalize = toc(t);
    timer.totalruntime = NaN;   % Update at the end of iteration
    % timer

    % Record all useful stats
    stats = savestats(stats, ...
        optimState,vp,elbo,elbo_sd,varss,sKL,sKL_true,gp,hypstruct.full,...
        Ns_gp,pruned,timer,options.Diagnostics);

    %----------------------------------------------------------------------
    %% Check termination conditions and warmup
    [optimState,stats,isFinished_flag,exitflag,action,msg] = ...
        vbmc_termination(optimState,action,stats,options);
    vp.stats.stable = stats.stable(optimState.iter);    % Save stability

    % Check if we are still warming-up
    if optimState.Warmup && iter > 1
        if options.RecomputeLCBmax
        	optimState.lcbmax_vec = recompute_lcbmax(gp,optimState,stats,options)';
        end
        [optimState,action,trim_flag] = vbmc_warmup(optimState,stats,action,options);
        if trim_flag    % Re-update GP after trimming
            gp = gpreupdate(gp,optimState,options);
        end
        if ~optimState.Warmup
            vp.optimize_mu = logical(options.VariableMeans);
            vp.optimize_weights = logical(options.VariableWeights);
            if options.BOWarmup
                optimState.gpMeanfun = options.gpMeanFun;
                hypstruct.hyp = [];
            end
            % Switch to main algorithm options
            options = options_main;
            hypstruct.runcov = [];    % Reset GP hyperparameter covariance
            optimState.vp_repo = []; % Reset VP repository
            optimState.acqInfo = getAcqInfo(options.SearchAcqFcn);   % Re-get acq info
        end
    end
    stats.warmup(iter) = optimState.Warmup;

    % Check and update fitness shaping / output warping threshold
    if ~isempty(optimState.OutwarpDelta) && optimState.R < options.WarpTolReliability
        Xrnd = vbmc_rnd(vp,2e4,0);
        ymu = gplite_pred(gp,Xrnd,[],[],0,1);
        ydelta = max([0,optimState.ymax-quantile(ymu,1e-3)])
        if (ydelta > optimState.OutwarpDelta*options.OutwarpThreshTol) && (optimState.R < 1)
            optimState.OutwarpDelta = optimState.OutwarpDelta*options.OutwarpThreshMult;
        end
    end

    if options.AcqHedge         % Update hedge values
        optimState.hedge = acqhedge_vbmc('upd',optimState.hedge,stats,options);
    end

    %% Write iteration output

    % vp.w

    % Stopped GP sampling this iteration?
    if Ns_gp == options.StableGPSamples && ...
            stats.gpNsamples(max(1,iter-1)) > options.StableGPSamples
        if Ns_gp == 0
            if isempty(action); action = 'switch to GP opt'; else; action = [action ', switch to GP opt']; end
        else
            if isempty(action); action = 'stable GP sampling'; else; action = [action ', stable GP sampling']; end
        end
    end

    if prnt > 2
        if options.BOWarmup && optimState.Warmup
            fprintf(displayFormat_warmup,iter,optimState.funccount,max(optimState.y_orig),action);
        else
            if optimState.Cache.active
                fprintf(displayFormat,iter,optimState.funccount,optimState.cachecount,elbo,elbo_sd,sKL,vp.K,optimState.R,action);
            elseif optimState.UncertaintyHandlingLevel > 0  && options.MaxRepeatedObservations > 0
                fprintf(displayFormat,iter,optimState.funccount,optimState.N,elbo,elbo_sd,sKL,vp.K,optimState.R,action);
            else
                fprintf(displayFormat,iter,optimState.funccount,elbo,elbo_sd,sKL,vp.K,optimState.R,action);
            end
        end
    end

    stats.timer(iter).totalruntime = toc(t0);

end

vp_old = vp;

% Pick "best" variational solution to return (and real vp, if train vp differs)
[vp,elbo,elbo_sd,idx_best] = ...
    best_vbmc(stats,iter,options.BestSafeSD,options.BestFracBack,options.RankCriterion,0);
new_final_vp_flag = idx_best ~= iter;
gp = stats.gp(idx_best);
vp.gp = gp;     % Add GP to variational posterior

% Last variational optimization with large number of components
[vp,elbo,elbo_sd,changedflag] = finalboost_vbmc(vp,idx_best,optimState,stats,options);
if changedflag; new_final_vp_flag = true; end

if new_final_vp_flag && prnt > 2
    % Recompute symmetrized KL-divergence
    sKL = max(0,0.5*sum(vbmc_kldiv(vp,vp_old,Nkl,options.KLgauss)));
end

% Convert training variational posterior to real variational posterior
vp_train = vp;
vp = vptrain2real(vp_train,1);
elbo = vp.stats.elbo;
elbo_sd = vp.stats.elbo_sd;

if new_final_vp_flag && prnt > 2
    if optimState.UncertaintyHandlingLevel > 0 && options.MaxRepeatedObservations > 0
        fprintf(displayFormat,Inf,optimState.funccount,optimState.N,elbo,elbo_sd,sKL,vp.K,stats.rindex(idx_best),'finalize');
    else
        fprintf(displayFormat,Inf,optimState.funccount,elbo,elbo_sd,sKL,vp.K,stats.rindex(idx_best),'finalize');
    end
end

% Set EXITFLAG based on stability (might check other things in the future)
switch exitflag
    case 0
        if vp.stats.stable; exitflag = 1; end
    case 1
        if ~vp.stats.stable; exitflag = 0; end
end

% Print final message
if prnt > 1
    fprintf('\n%s\n', msg);
    fprintf('Estimated ELBO: %.3f +/- %.3f.\n', elbo, elbo_sd);
    if exitflag < 1
        fprintf('Caution: Returned variational solution may have not converged.\n');
    end
    fprintf('\n');
end

if nargout > 4
    output = vbmc_output(vp,optimState,msg,stats,idx_best,vbmc_version);

    % Compute total running time and fractional overhead
    optimState.totaltime = toc(t0);
    output.overhead = optimState.totaltime / optimState.totalfunevaltime - 1;
end

if nargout > 5
    % Prepare SAMPLES struct
    idx = 1:optimState.Xn;
    samples.X = optimState.X_orig(idx,:);
    if isfield(optimState,'temperature') && ~isempty(optimState.temperature)
        T = optimState.temperature;
    else
        T = 1;
    end
    samples.y = optimState.y_orig(idx,:)*T;
    if isfield(optimState,'S') && ~isempty(optimState.S)
        samples.y_sd = optimState.S(idx,:)*T;
    end
    samples.active_flag = optimState.X_flag(idx,:);
    samples.nevals = optimState.nevals(idx,:);
end

if nargout > 7
    % Remove GP from stats struct unless diagnostic run
    if ~options.Diagnostics
        stats = rmfield(stats,'gp');
        stats = rmfield(stats,'gpHypFull');
        stats.timer(iter).totalruntime = toc(t0);
    end
end

if exitflag < 1 && options.RetryMaxFunEvals > 0
    % Rerun VBMC with better initialization if first try did not work
    if prnt > 0
        fprintf('First attempt did not converge. Trying to rerun variational optimization.\n');
    end

    % Get better VBMC parameters and initialization from current run
    vp0 = stats.vp(idx_best);
    [x0,LB,UB,PLB,PUB,Xvp] = initFromVP(vp0,LB,UB,PLB,PUB,0);
    Ninit = max(options.FunEvalStart,ceil(options.RetryMaxFunEvals/10));
    x0 = [x0; robustSampleFromVP(vp0,Ninit-1,Xvp)];

    options.FunEvalStart = Ninit;
    options.MaxFunEvals = options.RetryMaxFunEvals;
    options.RetryMaxFunEvals = 0;                   % Avoid infinite loop
    options.SGDStepSize = 0.2*options.SGDStepSize;  % Increase stability
    options.ActiveSampleGPUpdate = true;
    options.ActiveSampleVPUpdate = true;

    try
        [vp,elbo,elbo_sd,exitflag,output2,samples2,optimState2,stats] = vbmc(fun,x0,LB,UB,PLB,PUB,options,varargin{:});

        if nargout > 4
            optimState2.totaltime = toc(t0);
            output2.overhead = optimState.totaltime / (optimState.totalfunevaltime + optimState2.totalfunevaltime) - 1;
            output2.iterations = output2.iterations + output.iterations;
            output2.funccount = output2.funccount + output.funccount;
            output2.retried = 'yes';
            samples = samples2;
            output = output2;
            optimState = optimState2;
        end
    catch retryException
        msgText = getReport(retryException);
        warning(msgText);
        if prnt > 0
            fprintf('Attempt of rerunning variational optimization FAILED. Keeping original results.\n');
        end
        if nargout > 4
            output.retried = 'error';
        end
    end

else
    if nargout > 4; output.retried = 'no'; end
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function stats = savestats(stats,optimState,vp,elbo,elbo_sd,varss,sKL,sKL_true,gp,hyp_full,Ns_gp,pruned,timer,debugflag)

iter = optimState.iter;
stats.iter(iter) = iter;
stats.N(iter) = optimState.N;
stats.Neff(iter) = optimState.Neff;
stats.funccount(iter) = optimState.funccount;
stats.cachecount(iter) = optimState.cachecount;
stats.vpK(iter) = vp.K;
stats.warmup(iter) = optimState.Warmup;
stats.pruned(iter) = pruned;
stats.elbo(iter) = elbo;
stats.elbo_sd(iter) = elbo_sd;
stats.sKL(iter) = sKL;
if ~isempty(sKL_true)
    stats.sKL_true(iter) = sKL_true;
end
stats.gpNoise_hpd(iter) = sqrt(optimState.sn2hpd);
stats.gpSampleVar(iter) = varss;
stats.gpNsamples(iter) = Ns_gp;
stats.gpHypFull{iter} = hyp_full;
stats.timer(iter) = timer;
stats.vp(iter) = vp;
stats.gp(iter) = gplite_clean(gp);
if ~isempty(optimState.gpOutwarpfun)
    stats.outwarp_threshold(iter) = optimState.OutwarpDelta;
else
    stats.outwarp_threshold(iter) = NaN;
end
stats.lcbmax(iter) = optimState.lcbmax;
stats.t(iter) = NaN;    % Fill it at the end of the iteration

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function add2path()
%ADD2PATH Adds VBMC subfolders to MATLAB path.

subfolders = {'acq','ent','gplite','misc','shared','test','utils'};
% subfolders = {'acq','ent','gplite','misc','utils','warp'};
pathCell = regexp(path, pathsep, 'split');
baseFolder = fileparts(mfilename('fullpath'));

onPath = true;
for iFolder = 1:numel(subfolders)
    folder = [baseFolder,filesep,subfolders{iFolder}];
    if ispc  % Windows is not case-sensitive
      onPath = onPath & any(strcmpi(folder, pathCell));
    else
      onPath = onPath & any(strcmp(folder, pathCell));
    end
end

% ADDPATH is slow, call it only if folders are not on path
if ~onPath
    addpath(genpath(fileparts(mfilename('fullpath'))));
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [x0,LB,UB,PLB,PUB,Xvp] = initFromVP(vp,LB,UB,PLB,PUB,prnt)

if prnt > 2
    fprintf('Initializing VBMC from variational posterior (D = %d).\n', vp.D);
    if ~isempty(PLB) && ~isempty(PUB)
        fprintf('Using provided plausible bounds. Note that it might be better to leave them empty,\nand allow VBMC to set them using the provided variational posterior.\n');
    end
end

% Find mode in transformed space
x0t = vbmc_mode(vp,[],0);
x0 = warpvars_vbmc(x0t,'inv',vp.trinfo);

% Sample from variational posterior and set plausible bounds accordingly
if isempty(PLB) && isempty(PUB)
    Xvp = vbmc_rnd(vp,1e6,[],1);
    PLB = quantile(Xvp,0.05);
    PUB = quantile(Xvp,0.95);
else
    Xvp = [];
end
if isempty(LB); LB = vp.trinfo.lb_orig; end
if isempty(UB); UB = vp.trinfo.ub_orig; end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function Xrnd = robustSampleFromVP(vp,Ns,Xrnd,quantile_thresh)
%ROBUSTSAMPLEFROMVP Robust sample from variational posterior.

if nargin < 3; Xrnd = []; end
if nargin < 4 || isempty(quantile_thresh); quantile_thresh = 0.01; end

Ns_big = 1e4;
Xrnd = [Xrnd; vbmc_rnd(vp,max(0,Ns_big-size(Xrnd,1)),[],1)];
Xrnd = Xrnd(1:Ns_big,:);

y = vbmc_pdf(vp,Xrnd);
y_thresh = quantile(y,quantile_thresh);

Xrnd = Xrnd(y > y_thresh,:);
Xrnd = Xrnd(1:Ns,:);

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function timer = timer_init()
%TIMER_INIT Initialize iteration timer.

timer.activeSampling = 0;
timer.funEvals = 0;
timer.gpTrain = 0;
timer.variationalFit = 0;
timer.warping = 0;
timer.finalize = 0;

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function acqInfo = getAcqInfo(SearchAcqFcn)
%GETACQINFO Get information from acquisition function(s)

for iAcq = 1:numel(SearchAcqFcn)
    try
        % Called with first empty input should return infos
        acqInfo{iAcq} = SearchAcqFcn{iAcq}([]);
    catch
        acqInfo{iAcq} = [];
    end
end

end


function volROI = flat2volROI(flatROI,flatView,volView)
%
% volROI = flat2volROI(flatROI,flatView,volView)
%
% Creates a volume ROI from a flat ROI by looking up the
% corresponding coords.  The tricky thing about this function
% is that FLAT -> VOL is a one to many mapping, and we need to
% find all the corresponding volume coords.
%
% flatROI and volROI are ROI structures, like those found in
% view.ROIs
%
% volView must be the VOLUME structure.
% flatView must be the FLAT structure.
%
% djh, 8/98.
%
% djh, 8/4/99.  Need to round the coords because we no longer do it in
%               loadGLocs.  Also, remove duplicate flat coords.
% djh, 2/2001. Use intersectCols instead of coords2Indices

global mrSESSION

% check that the ROI fields are properly set
flatROI = roiCheck(flatROI);

% This should be much fastter but its busted because intersect removes duplicates and
% we need to keep the duplicates here because flat->gray is a one-to-many mapping.
%
% coords = [];
% for h = 1:2
%     ROIcoords = flatROI.coords([1:2],find(flatROI.coords(3,:) == h));
%     if ~isempty(ROIcoords)
%         [tmp,ROIIndices,flatIndices] = intersectCols(ROIcoords,round(flatView.coords{h}));
%         coords = [coords, flatView.grayCoords{h}(:,flatIndices)];
%     end
% end

coords = [];
flatImSize = flatView.ui.imSize;
for h = 1:2
  ROIcoords = flatROI.coords([1:2],find(flatROI.coords(3,:) == h));
  if ~isempty(ROIcoords)
    ROIIndices = coords2Indices(ROIcoords,flatImSize);
    flatIndices = coords2Indices(round(flatView.coords{h}),flatImSize);
    bothIndices = intersect(ROIIndices,flatIndices);
    for id = 1:length(bothIndices)
      indices = find(flatIndices == bothIndices(id));
      coords = [coords, flatView.grayCoords{h}(:,indices)];
    end
  end
end

% Remove duplicates (there shouldn't be any, but just to make sure).
coords = intersectCols(coords,coords);

% Set the fields
volROI = flatROI;
volROI.coords = coords;
volROI.viewType = volView.viewType;

return


function [data, xAxis, yAxis, misc] = impload(filename)
% Reads PerkinElmer vis image files files.
% This version is compatible with single image files.
%
% [data, xAxis, yAxis] = visload(filename):
%   data: length(x) * length(y) 2D array
%   xAxis: vector for horizontal axis (e.g. micrometers)
%   yAxis: vector for vertical axis (e.g. micrometers)

% Copyright (C)2007 PerkinElmer Life and Analytical Sciences
% Stephen Westlake, Seer Green
%
% History
% 2007-04-29 SW     Initial version

% Read the bitmap
data = imread(filename, 'bmp');;

fid = fopen(filename,'r');
if fid == -1
    error('Cannot open the file.');
    return
end

% Fixed file header of signature and description
signature = setstr(fread(fid, 2, 'uchar')');
if ~strcmp(signature, 'BM')
    error('This is not a PerkinElmer vis file.');
    return
end
fseek(fid, 0, 'bof');

% Read the trailer
fseek(fid, -4 * 8, 'eof');

% YAxis is reversed to suit image()
xAxis(1) = fread(fid, 1, 'double');     % x1
yAxis(2) = fread(fid, 1, 'double');     % y1
xAxis(2) = fread(fid, 1, 'double');     % x2
yAxis(1) = fread(fid, 1, 'double');     % y2

fclose(fid);


function S = stripR(A)
%------------------------------------------------------------------------------
%
% Design and implementation by:
% Dr. Paul M. de Zeeuw <Paul.de.Zeeuw@cwi.nl>  http://homepages.cwi.nl/~pauldz/
% Last Revision: June 6, 1999.
% (c) 1999-2002 Stichting CWI, Amsterdam
%------------------------------------------------------------------------------
[n, m] = size(A);
Y = A.';
S = Y(1:(m-1),:).';
%------------------------------------------------------------------------------


function flatROI = vol2flatROILevels(volROI,volView,flatView)
%
% flatROI = vol2flatROILevels(volROI,volView,flatView)
%
% Creates a flat ROI from a volume ROI by looking up the
% corresponding coords. Will create coordinates in the
% 'separate gray levels' slices.
%
% flatROI and volROI are ROI structures, like those found in
% view.ROIs
%
% volView must be the VOLUME structure.
% flatView must be the FLAT structure.
%
% ras, 10/04. Version for flat level views, off vol2flatROI.
% Intersect volROI.coords with volView.coords.  When in gray mode,
% this restricts the ROI to the gray matter.
ROIcoords = volROI.coords;
ROIcoords = intersectCols(ROIcoords,volView.coords);

% init coords for flat ROI
coords = [];

% get grayCoords for flat, across all slices
grayCoords = []; flatCoords = [];
for slice = 1:numSlices(flatView)
    grayCoords = [grayCoords flatView.grayCoords{slice}];
    sliceCoords = flatView.coords{slice};
    sliceCoords(3,:) = slice;
    flatCoords = [flatCoords sliceCoords];
end

[bothCoords ia ib] = intersectCols(ROIcoords,grayCoords);
coords = flatCoords(:,ib);

% % error check. ROIcoords should be contained entirely within gray coords.
% if (size(ROIcoords,2) ~= size(coords,2)) & (strcmp(volView.viewType,'Gray'))
%    fprintf(['\nGray nodes loaded from gray classification file are ',...
%          'incompatible with those loaded from the flat.mat file. ',...
%          'Rebuild the Gray and Flat view: rm Gray/*.mat Flat/*.mat']);
% end

% Remove duplicates (there shouldn't be any, but just to make sure).
coords = intersectCols(coords,coords);

% Set the fields
flatROI.coords = coords;
flatROI.name = volROI.name;
flatROI.color = volROI.color;
flatROI.viewType = flatView.viewType;

return

% from the older intersect code:
% flatCoords = round(flatView.coords(:,ib));
% flatCoords = [flatCoords; h*ones([1,size(flatCoords,2)])];
% coords = [coords,flatCoords];


function test_example_sourcemodel_aligned2mni_atlas

% MEM 4gb
% WALLTIME 00:10:00

%
%% Create brain atlas based MNI-aligned grids in individual head-space
%
% When combining the source-level data of multiple subjects this data is typically first interpolated (e.g., using **[ft_sourceinterpolate](https://github.com/fieldtrip/fieldtrip/blob/release/ft_sourceinterpolate.m)**) and then spatially normalized to a template brain (e.g., using **[ft_volumenormalise](https://github.com/fieldtrip/fieldtrip/blob/release/ft_volumenormalise.m)**). However it is also possible to define the source-reconstruction grid for each individual subject in such a way that all these grids are already aligned in MNI-space. The combination or statistic of source-level data across subjects can then directly be computed within the source-structure without the need to interpolate and normalize each volume. In addition, the position of the grid points can be chosen in a way that only locations corresponding to particular brain areas (parcels) are included.
%
% The procedure is as follows. First, a template grid is computed on the basis of the standard head model located in the FieldTrip template directory. Subsequently, a brain atlas is loaded using **[ft_read_atlas](https://github.com/fieldtrip/fieldtrip/blob/release/fileio/ft_read_atlas.m)**. On the basis on the brain atlas and the template grid it is possible to create a binary mask of all locations in the template grid that match atlas locations using **[ft_volumelookup](https://github.com/fieldtrip/fieldtrip/blob/release/ft_volumelookup.m)**. Finally the mask is used to determine which location shall be defined as 'inside' the head model.
%
[ftver, ftpath] = ft_version;
load(fullfile(ftpath, 'template/headmodel/standard_singleshell.mat'));

cfg = [];
cfg.xgrid  = -20:1:20;
cfg.ygrid  = -20:1:20;
cfg.zgrid  = -20:1:20;
cfg.unit   = 'cm';
cfg.tight  = 'yes';
cfg.inwardshift = -1.5;
cfg.headmodel   = vol;
template_grid   = ft_prepare_sourcemodel(cfg);
template_grid.coordsys = 'mni';

template_grid = ft_convert_units(template_grid,'cm');

figure;
hold on
ft_plot_mesh(template_grid.pos(template_grid.inside,:));
ft_plot_headmodel(vol,  'facecolor', 'cortex', 'edgecolor', 'none');
ft_plot_axes(vol);
alpha 0.5
camlight

% Read the atlas, convert to units of cm and create the binary mask.
%
atlas = ft_read_atlas(fullfile(ftpath, 'template/atlas/aal/ROI_MNI_V4.nii'));

atlas = ft_convert_units(atlas,'cm');

cfg = [];
cfg.atlas      = atlas;
cfg.roi        = atlas.tissuelabel;  % here you can also specify a single label, i.e. single ROI
mask           = ft_volumelookup(cfg, template_grid);

% Now we determine all indices of the binary mask to be considered as inside the head model. And plot the result. Note the missing dipole locations for example in the vicinity of the ventricles.
%
template_grid.inside = false(template_grid.dim);
template_grid.inside(mask==1) = true;

figure;
ft_plot_mesh(template_grid.pos(template_grid.inside,:));

%
% Load the subject-specific MRI from [here](ftp://ftp.fieldtriptoolbox.org/pub/fieldtrip/tutorial/salzburg/mri.mat) and inverse-warp the subject specific grid to the template grid.
%
mri = ft_read_mri(dccnpath('/home/common/matlab/fieldtrip/data/ftp/tutorial/salzburg/mri.mat'));

cfg                = [];
cfg.warpmni   = 'yes';
cfg.template  = template_grid;
cfg.nonlinear = 'yes'; % use non-linear normalization
cfg.mri            = mri;
sourcemodel        = ft_prepare_sourcemodel(cfg);

% Finally, you can load the subject-specific headmodel from [here](ftp://ftp.fieldtriptoolbox.org/pub/fieldtrip/tutorial/salzburg/hdm.mat) and check the result with the following code.
%
close all

hdm = ft_read_mri(dccnpath('/home/common/matlab/fieldtrip/data/ftp/tutorial/salzburg/hdm.mat'));

hdm         = ft_convert_units(hdm, 'm');
sourcemodel = ft_convert_units(sourcemodel, 'm');

figure
hold on
ft_plot_headmodel(hdm,  'facecolor', 'cortex', 'edgecolor', 'none');
ft_plot_axes(hdm);
alpha 0.4  % make the surface transparent
ft_plot_mesh(sourcemodel.pos(sourcemodel.inside,:));
view ([0 -90 0])


 function out = subsref(ob, arg)
%function out = subsref(ob, arg)
%	evaluate call of the form ob.ref or ob(ref,:)

%
%	ob.?
%
if arg.type == '.'
	out = getfield(struct(ob), arg.subs);

%
%	ob(?)
%
elseif arg.type == '()'

	if ~isempty(ob.index1) | ~isempty(ob.index2)
		error 'subscripts after subscripts perhaps not done'
	end

	sub = arg.subs;

	%
	%	G(:,:)
	%
	if length(sub) == 2 & sub{1} == ':' & sub{2} == ':'
		out = ob.G;

	%
	%	G(:,j)
	%
	elseif length(sub) == 2 & sub{1} == ':'
		jj = col(sub{2});
		if islogical(jj)
			if length(jj) ~= ob.dims(2)
				error 'bad column logical length'
			end
		else
			bad = jj < 1 | jj > ob.dims(2);
			if any(bad)
				printf('bad column indeces:')
				jj(bad)'
				error 'subsref problem'
			end
		end

		%
		%	handle G(:,mask(:)) as a special case
		%
		if length(jj) == ob.nx*ob.ny & all(jj == ob.mask(:)) ...
				& ~ob.is_masked & ~ob.is_transpose
			ob.is_masked = true;
			ob.dims(2) = sum(ob.mask(:));
			ob.G = ob.G(:,ob.mask(:));
			out = ob;

		else
%			G = ob.G;
			out = subsref(ob.G, arg);
		end

	%
	%	G(i,:)
	%
	elseif length(sub) == 2 & sub{2} == ':'
		out = subsref(ob.G, arg);

	else
		arg.subs
		error todo
	end

else
	error(sprintf('type %s notdone', arg.type))
end


function dataOut=plotMultipleVoxAboveThreshPerCondition(view)
%
% dataOut=plotMultipleVoxAboveThreshPerCondition
%
% Bar plot of the number of superthreshold voxels (based on the co) in a particular condition
% All y-axes are made the same. The bar heights
% and total voxels in each ROI can be obtained from the userdata
%
% arw 04/01/05

mrGlobals;

% Select ROIs
nROIs=size(view.ROIs,2);
roiList=cell(1,nROIs);
for r=1:nROIs
    roiList{r}=view.ROIs(r).name;
end
selectedROIs = find(buttondlg('ROIs to Plot',roiList));
nROIs=length(selectedROIs);
if (nROIs==0)
    error('No ROIs selected');
end

% Plot it
selectGraphWin
clf
fontSize = 8;
headerStr = ['Vox above thresh'];
set(gcf,'Name',headerStr);

minylim=0;
maxylim=0;
nrows=0;
ncols=0;

nscans = numScans(view);
ROIcos=zeros(nscans,nROIs);
thisCoThresh=get(view.ui.cothresh.sliderHandle,'value');
totalVoxels=zeros(nscans,nROIs);
voxAboveThresh=zeros(nscans,nROIs);

for r=1:nROIs

    n=selectedROIs(r);
    view = selectROI(view,n); % is there another way?
    ROIcoords=getCurROIcoords(view);

    for scanNum=1:nscans

        subCo = getCurDataROI(view,'co',scanNum,ROIcoords);
        totalVoxels(scanNum,r)=size(subCo,2);
        voxAboveThresh(scanNum,r)=length(find(subCo>=thisCoThresh));
    end

    xstr{r}=int2str(r);
    roiName{r}=view.ROIs(selectedROIs(r)).name;
    fprintf(['\n#%d :',roiName{r}],r);

end
dataOut.totalVoxels=totalVoxels;
dataOut.voxAboveThresh=voxAboveThresh;
dataOut.ROIname=roiName;

% Now do the plotting

if nscans<=3
    nrows=1;
    ncols=nscans;
    fontSize = 9;
elseif nscans<=8
    nrows=2;
    ncols=ceil(nscans/nrows);
    fontSize = 8;
else
    nrows=ceil(sqrt(nscans));
    ncols=ceil(nscans/nrows);
    fontSize = 6;
end

for r=1:nscans

    subplot(nrows,ncols,r);

    %plot the bar graph
    h=bar(voxAboveThresh(r,:));% ,ROIseZ(r,:)',xstr);
    xlabel('ROI','FontSize',fontSize);
    ylabel('Number above threshold','FontSize',fontSize);
    set(gca,'FontSize',ceil(fontSize*1.2));
    conditionName{r}=dataTYPES(view.curDataType).scanParams(r).annotation;
    fprintf(['\nCondition #%d :',conditionName{r}],r);

    fprintf('\nTotalVoxels: %d',totalVoxels(r,:));

    title(['Condition #: ' int2str(r)]);
    yl=ylim;

    if yl(1)< minylim
        minylim=yl(1);
    end
    if yl(2)> maxylim
        maxylim=yl(2);
    end
end

set(gca,'UserData',dataOut);

% give all plots same y-axis

for r=1:nscans
    subplot(nrows,ncols,r);
    ylim([minylim maxylim]);
end

return;


function [ vX, vV, mX ] = ProjectedGdFista( vX, vV, hG, hP, numIterations, stepSize )
% ----------------------------------------------------------------------------------------------- %
% [ vX, vV, mX ] = ProjectedGdFista( vX, vV, hG, hP, numIterations, stepSize )
% Solves an objective function by the Projected Gradient Descent /
% Projected Gradient Method (PGM) with FISTA Acceleration.
% Input:
%   - vX                -   Input Vector.
%                           The starting point for the gradient iterations.
%                           Structure: Vector (numElements x 1).
%                           Type: 'Single' / 'Double'.
%                           Range: (-inf, inf).
%   - vV                -   Acceleration Vector.
%                           The buffer to hold the accelerated direction
%                           vector.
%                           Structure: Vector (numElements x 1).
%                           Type: 'Single' / 'Double'.
%                           Range: (-inf, inf).
%   - hG                -   Gradient Function Handler.
%                           A function handler which accepts a vector of
%                           size (numElements x 1) and returns the gradient
%                           in the form of a vector (numElements x 1).
%                           Structure: Function Handler.
%                           Type: NA.
%                           Range: NA.
%   - hP                -   Gradient Function Handler.
%                           A function handler which accepts a vector of
%                           size (numElements x 1) and returns the gradient
%                           in the form of a vector (numElements x 1).
%                           Structure: Function Handler.
%                           Type: NA.
%                           Range: NA.
%   - numIterations     -   Number of Iterations.
%                           Sets the number of iterations of the gradient
%                           descent.
%                           Structure: Scalar.
%                           Type: 'Single' / 'Double'.
%                           Range: {1, 2, 3, ...}.
%   - stepSize          -   Step Size.
%                           Sets the step size of the Gradient step.
%                           Structure: Scalar.
%                           Type: 'Single' / 'Double'.
%                           Range: (0, inf).
% Output:
%   - vX                -   Output Vector.
%                           The end point for the gradient iterations.
%                           Structure: Vector (numElements x 1).
%                           Type: 'Single' / 'Double'.
%                           Range: (-inf, inf).
%   - vV                -   Acceleration Vector.
%                           The buffer to hold the accelerated direction
%                           vector of last iteration. May used by a wrapper
%                           for warm start.
%                           Structure: Vector (numElements x 1).
%                           Type: 'Single' / 'Double'.
%                           Range: (-inf, inf).
%   - mX                -   The Path Matrix.
%                           The points of the gradient iterations.
%                           Structure: Matrix (numElements x numIterations).
%                           Type: 'Single' / 'Double'.
%                           Range: (-inf, inf).
% References:
%   1.  Geoff Gordon & Ryan Tibshirani - Accelerated First Order Methods (https://www.cs.cmu.edu/~ggordon/10725-F12/slides/09-acceleration.pdf).
% Remarks:
%   1.  B
% TODO:
%   1.  C
% Release Notes:
%   -   1.0.000     26/12/2020  Royi Avital
%       *   First release version.
% ----------------------------------------------------------------------------------------------- %

FALSE   = 0;
TRUE    = 1;

OFF     = 0;
ON      = 1;

mX = zeros(size(vX, 1), numIterations);

vG          = zeros(size(vX, 1), 1);
mX(:, 1)    = vX;
mX(:, 2)    = vX - stepSize * hG(vX);

for ii = 3:numIterations
    vV(:) = mX(:, ii - 1) + (((ii - 2) / (ii + 1)) * (mX(:, ii - 1) - mX(:, ii - 2)));
    vG(:) = hG(vV); %<! The gradient
    vX(:) = hP(vV - (stepSize * vG)); %<! Projection step

    mX(:, ii) = vX;
end

end


function loss_group_info=gen_network_loss_group(train_opts, net_config, group_input_info)

feat_output_dim=group_input_info.var_dims;
assert(length(feat_output_dim)==1);

loss_group_info=My_net_util.gen_group_info_basic(net_config, 'loss_group');

stageout_info=[];
stageout_info.stage_idx=1;
stageout_info.output_dim=feat_output_dim;
stageout_info.stageout_num=1;
stageout_info.gen_prediction_info=true;

[loss_group_info.net_info, dag_net_lossgroup]=do_gen_net_info(train_opts, stageout_info);

if ~isempty(dag_net_lossgroup)
    loss_group_info.dag_net=dag_net_lossgroup;
    loss_group_info.use_dagnn=true;
end

loss_group_info.forward_evaluate_fn=@cnn_loss_group_evaluate;

loss_group_info.prediction_layer_idxes=loss_group_info.net_info.ref.prediction_layer_idxes;
net_config.ref.group_infos{loss_group_info.group_idx,1}=loss_group_info;

net_config.ref.root_group_info.child_group_idxes(end+1,1)=loss_group_info.group_idx;

end

function [net_info, dag_net]=do_gen_net_info(train_opts, stageout_info)

    loss_config=train_opts.loss_config;
    class_num=loss_config.class_num;
    loss_group_input_dim=stageout_info.output_dim;

    layers=cell(0);

    dag_net=[];
    if loss_config.lossgroup_conv_num>0
        [dag_net, output_dim_crf_conv1, one_layers]=gen_dag_net_lossgroup(loss_config, loss_group_input_dim);
        layers=cat(2, layers, one_layers);
    else
        output_dim_crf_conv1=loss_group_input_dim;
    end

    layers{end+1} = struct('type', 'dropout', 'rate', 0.5) ;

    output_dim_crf_conv2=class_num;
    input_dim_crf_conv2=output_dim_crf_conv1;
    layers{end+1} = struct('type', 'conv', ...
                   'filters', 0.01 * randn(3,3,input_dim_crf_conv2, output_dim_crf_conv2,'single'), ...
                   'biases', zeros(1, output_dim_crf_conv2, 'single'), ...
                   'stride', 1, ...
                   'pad', 0 );

    one_layer=[];
    one_layer.type='my_custom';
    one_layer.custom_type='dense_softmaxloss';

    one_layer.forward_fn=@cnn_layer_dense_softmaxloss_forward;
    one_layer.backward_fn=@cnn_layer_dense_softmaxloss_backward;

    one_layer.stageout_info=stageout_info;

    one_layer.layer_update_fn=[];
    layers{end+1}=one_layer;

layers=gen_padding_keep_size(layers);

net_info=My_net_util.gen_net_info_basic();

net_info.ref.name='final_loss';
net_info.ref.layers=layers;
net_info.ref.prediction_layer_idxes=length(layers);

net_info.ref.lr_multiplier=loss_config.lr_multiplier;

end

function [dag_net, output_dim, layers]=gen_dag_net_lossgroup(loss_config, loss_group_input_dim)

conv_num=loss_config.lossgroup_conv_num;
one_output_dim=loss_config.lossgroup_conv_filter_num;

dag_net=dagnn.DagNN();

start_var_name='data_input';

layer_gen_info=[];
layer_gen_info.one_output_dim=loss_group_input_dim;
layer_gen_info.one_outputs={start_var_name};

if loss_group_input_dim~=one_output_dim
	error('should not come here!');
end

layer_name_prefix=sprintf('lossblock');
layer_gen_info=My_net_util.add_res_conv_block(dag_net, layer_gen_info, one_output_dim, conv_num, layer_name_prefix);

output_dim=layer_gen_info.one_output_dim;

output_var_name=layer_gen_info.one_outputs{1};
output_var_idx=dag_net.getVarIndex(output_var_name);
dag_net.vars(output_var_idx).fanout=0;

% debug:
% dag_net.print('Format', 'dot');

My_net_util.fix_padding_resnet(dag_net);

layers=cell(0);

one_layer=[];
one_layer.type='my_custom';
one_layer.custom_type='dagnn_wrapper';
one_layer.forward_fn=@cnn_layer_dagnn_wrapper_forward;
one_layer.backward_fn=@cnn_layer_dagnn_wrapper_backward;
one_layer.layer_update_fn=[];

one_layer.input_var_names=dag_net.getInputs();
one_layer.output_var_names=dag_net.getOutputs();

assert(length(one_layer.input_var_names)==1);
assert(length(one_layer.output_var_names)==1);

one_layer.use_single_output=true;

layers{end+1}=one_layer;

end


classdef GenericSumOfSinesSteeringModel < AbstractSteeringModel
    %GenericSumOfSinesSteeringModel Summary of this class goes here
    %   Detailed explanation goes here

    properties
        gammaAngleModel(1,1) SumOfSinesModel = SumOfSinesModel(0);
        betaAngleModel(1,1) SumOfSinesModel = SumOfSinesModel(0);
        alphaAngleModel(1,1) SumOfSinesModel = SumOfSinesModel(0);

        gammaContinuity(1,1) logical = false;
        betaContinuity(1,1) logical = false;
        alphaContinuity(1,1) logical = false;

        refFrame AbstractReferenceFrame
        controlFrame AbstractControlFrame
    end

    methods
        function dcm = getBody2InertialDcmAtTime(obj, ut, rVect, vVect, bodyInfo)
            gammaAng = obj.gammaAngleModel.getValueAtTime(ut);
            betaAng = obj.betaAngleModel.getValueAtTime(ut);
            alphaAng = obj.alphaAngleModel.getValueAtTime(ut);

%             elemSet = CartesianElementSet(ut, rVect(:), vVect(:), bodyInfo.getBodyCenteredInertialFrame());
%             if(not(isempty(obj.refFrame.getOriginBody())))
%                 elemSet = elemSet.convertToFrame(obj.refFrame);
%             end

            dcm = real(obj.controlFrame.computeDcmToInertialFrame(ut, rVect, vVect, bodyInfo, gammaAng, betaAng, alphaAng, obj.refFrame));
        end

        function [angleModel, continuity] = getAngleNModel(obj, n)
            angleModel = SumOfSinesModel.empty(1,0);

            switch n
                case 1
                    angleModel = obj.gammaAngleModel;
                    continuity = obj.gammaContinuity;
                case 2
                    angleModel = obj.betaAngleModel;
                    continuity = obj.betaContinuity;
                case 3
                    angleModel = obj.alphaAngleModel;
                    continuity = obj.alphaContinuity;
            end
        end

        function t0 = getT0(obj)
            t0 = obj.alphaAngleModel.getT0();
        end

        function setT0(obj, newT0)
            obj.gammaAngleModel.setT0(newT0);
            obj.betaAngleModel.setT0(newT0);
            obj.alphaAngleModel.setT0(newT0);
        end

        function setTimeOffsets(obj, timeOffset)
            obj.gammaAngleModel.setTimeOffset(timeOffset);
            obj.betaAngleModel.setTimeOffset(timeOffset);
            obj.alphaAngleModel.setTimeOffset(timeOffset);
        end

        function [angle1Cont, angle2Cont, angle3Cont] = getContinuityTerms(obj)
            angle1Cont = obj.gammaContinuity;
            angle2Cont = obj.betaContinuity;
            angle3Cont = obj.alphaContinuity;
        end

        function setContinuityTerms(obj, angle1Cont, angle2Cont, angle3Cont)
            obj.gammaContinuity = angle1Cont;
            obj.betaContinuity = angle2Cont;
            obj.alphaContinuity = angle3Cont;
        end

        function setConstsFromDcmAndContinuitySettings(obj, dcm, ut, rVect, vVect, bodyInfo)
            if(obj.gammaContinuity || obj.betaContinuity || obj.alphaContinuity)
                elemSet = CartesianElementSet(ut, rVect(:), vVect(:), bodyInfo.getBodyCenteredInertialFrame());

%                 if(not(isempty(obj.refFrame.getOriginBody())))
%                     elemSet = elemSet.convertToFrame(obj.refFrame);
%                 end

                [gammaAngle, betaAngle, alphaAngle] = obj.controlFrame.getAnglesFromInertialBodyAxes(LaunchVehicleAttitudeState(dcm), elemSet.time, elemSet.rVect(:), elemSet.vVect(:), elemSet.frame.getOriginBody(), obj.refFrame);

                if(obj.gammaContinuity)
                    obj.gammaAngleModel.setConstValueForContinuity(gammaAngle);
                end

                if(obj.betaContinuity)
                    obj.betaAngleModel.setConstValueForContinuity(betaAngle);
                end

                if(obj.alphaContinuity)
                    obj.alphaAngleModel.setConstValueForContinuity(alphaAngle);
                end
            end
        end

        function setInitialAttitudeFromState(obj, stateLogEntry, tOffsetDelta)
            t0 = stateLogEntry.time;
            obj.setT0(t0);

            obj.gammaAngleModel.setTimeOffset(obj.gammaAngleModel.getTimeOffset() + tOffsetDelta);
            obj.betaAngleModel.setTimeOffset(obj.betaAngleModel.getTimeOffset() + tOffsetDelta);
            obj.alphaAngleModel.setTimeOffset(obj.alphaAngleModel.getTimeOffset() + tOffsetDelta);
        end

        function [angle1Name, angle2Name, angle3Name] = getAngleNames(obj)
            angleNames = obj.controlFrame.getControlFrameEnum().angleNames;
            angle1Name = angleNames{1};
            angle2Name = angleNames{2};
            angle3Name = angleNames{3};
        end

        function newSteeringModel = deepCopy(obj)
            newSteeringModel = GenericSumOfSinesSteeringModel(obj.gammaAngleModel.deepCopy(), obj.betaAngleModel.deepCopy(), obj.alphaAngleModel.deepCopy());

            newSteeringModel.gammaContinuity = obj.gammaContinuity;
            newSteeringModel.betaContinuity = obj.betaContinuity;
            newSteeringModel.alphaContinuity = obj.alphaContinuity;
            newSteeringModel.refFrame = obj.refFrame;
            newSteeringModel.controlFrame = obj.controlFrame;
        end

        function optVar = getNewOptVar(obj)
            optVar = SetGenericSumOfSinesSteeringModelActionOptimVar(obj);
        end

        function optVar = getExistingOptVar(obj)
            optVar = obj.optVar;
        end

        function tf = usesRefFrame(~)
            tf = true;
        end

        function refFrame = getRefFrame(obj)
            refFrame = obj.refFrame;
        end

        function setRefFrame(obj, refFrame)
            obj.refFrame = refFrame;
        end

        function tf = usesControlFrame(~)
            tf = true;
        end

        function cFrame = getControlFrame(obj)
            cFrame = obj.controlFrame;
        end

        function setControlFrame(obj, cFrame)
            obj.controlFrame = cFrame;
        end

        function tf = requiresReInitAfterSoIChange(~)
            tf = false;
        end

        function enum = getSteeringModelTypeEnum(~)
            enum = SteerModelTypeEnum.SumOfSinesAngles;
        end
    end

    methods(Access=private)
        function obj = GenericSumOfSinesSteeringModel(gammaAngleModel, betaAngleModel, alphaAngleModel)
            obj.gammaAngleModel = gammaAngleModel;
            obj.betaAngleModel = betaAngleModel;
            obj.alphaAngleModel = alphaAngleModel;

            obj.controlFrame = InertialControlFrame();
        end
    end

    methods(Static)
        function model = getDefaultSteeringModel()
            gammaAngleModel = SumOfSinesModel(0);
            betaAngleModel = SumOfSinesModel(0);
            alphaAngleModel = SumOfSinesModel(0);

            model = GenericSumOfSinesSteeringModel(gammaAngleModel, betaAngleModel, alphaAngleModel);
        end

        function typeStr = getTypeNameStr()
            typeStr = SteeringModelEnum.GenericSumOfSines.nameStr;
        end
    end
end