function [n_best P_best V C]=fast_oopsi(F,V,P)
% this function solves the following optimization problem:
% (*) n_best = argmax_{n >= 0} P(n | F)
% which is a MAP estimate for the most likely spike train given the
% fluorescence signal.  given the model:
%
% <latex>
% \begin{align}
% C_t &= \gamma*C_{t-1} + n_t, \qquad & n_t & \sim \text{Poisson}(n_t; \lamda_t \Delta)
% F_t &= \alpha(C_t + \beta) + \sigma \varepsilon_t, &\varepsilon_t &\sim \mathcal{N}(0,1)
% \end{align}
% </latex>
%
% if F_t is a vector, then 'a' is a vector as well
% we approx the Poisson with an Exponential (which means we don't require integer numbers of spikes).
% we take an "interior-point" approach to impose the nonnegative contraint on (*).
% each step is solved in O(T)
% time by utilizing gaussian elimination on the tridiagonal hessian, as
% opposed to the O(T^3) time typically required for non-negative
% deconvolution.
%
% Input---- only F is REQUIRED.  the others are optional
% F:        fluorescence time series (can be a vector (1 x T) or a matrix (Np x T)
%
% V.        structure of algorithm Variables
%   Ncells: # of cells within ROI
%   T:      # of time steps
%   Npixels:# of pixels in ROI
%   dt:     time step size, ie, frame duration, ie, 1/(imaging rate)
%   n:      if true spike train is known, and we are plotting, plot it (only required is est_a==1)
%   h:      height of ROI (assumes square ROI) (# of pixels) (only required if est_a==1 and we are plotting)
%   w:      width of ROI (assumes square ROI) (# of pixels) (only required if est_a==1 and we are plotting)
%
%   THE FOLLOWING FIELDS CORRESPOND TO CHOICES THAT THE USER MAKE
%
%   fast_poiss:     1 if F_t ~ Poisson, 0 if F_t ~ Gaussian
%   fast_nonlin:    1 if F_t is a nonlinear f(C_t), and 0 if F_t is a linear f(C_t)
%   fast_plot:      1 to plot results after each pseudo-EM iteration, 0 otherwise
%   fast_thr:       1 if thresholding inferred spike train before estiamting {a,b}
%   fast_iter_max:  max # of iterations of pseudo-EM  (1 to use default initial parameters)
%   fast_ignore_post: 1 to keep iterating pseudo-EM even if posterior is not increasing, 0 otherwise
%
%   THE BELOW FIELDS INDICATE WHETHER ONE WANTS TO ESTIMATE EACH OF THE
%   PARAMETERS. IF ANY IS SET TO ZERO, THEN WE DO NOT TRY TO UPDATE THE
%   ORIGINAL ESTIMATE, GIVEN EITHER BY THE USER, OR THE INITIAL ESTIMATE
%   FROM THE CODE
%
%   est_sig:    1 to estimate sig
%   est_lam:    1 to estimate lam
%   est_gam:    1 to estimate gam
%   est_b:      1 to estimate b
%   est_a:      1 to estimate a
%
% P.        structure of neuron model Parameters
%
%   a:      spatial filter
%   b:      background fluorescence
%   sig:    standard deviation of observation noise
%   gam:    decayish, ie, tau=dt/(1-gam)
%   lam:    firing rate-ish, ie, expected # of spikes per frame
%
% Output---
% n_best:   inferred spike train
% P_best:   inferred parameter structure
% V:        structure of Variables for algorithm to run

%% initialize algorithm Variables
starttime   = cputime;
siz         = size(F);      if siz(2)==1, F=F'; siz=size(F); end
j=0;
% variables determined by the data
if nargin < 2,              V   = struct;       end
if ~isfield(V,'Ncells'),    V.Ncells = 1;       end     % # of cells in image
if ~isfield(V,'T'),         V.T = siz(2);       end     % # of time steps
if ~isfield(V,'Npixels'),   V.Npixels = siz(1); end     % # of pixels in ROI
if ~isfield(V,'dt'),                                    % frame duration
    fr = input('\nwhat was the frame rate for this movie (in Hz)?: ');
    V.dt = 1/fr;
end

% variables determined by the user
if ~isfield(V,'fast_poiss'),V.fast_poiss = 0;   end     % whether observations are Poisson
if ~isfield(V,'fast_nonlin'),   V.fast_nonlin   = 0; end
if V.fast_poiss && V.fast_nonlin,
    reply = input('\ncan be nonlinear observations and poisson, \ntype 1 for nonlin, 2 for poisson, anything else for neither: ');
    if reply==1,        V.fast_poiss = 0;   V.fast_nonlin = 1;
    elseif reply==2,    V.fast_poiss = 1;   V.fast_nonlin = 0;
    else                V.fast_poiss = 0;   V.fast_nonlin = 0;
    end
end
if ~isfield(V,'fast_iter_max'), V.fast_iter_max=1; end % max # of iterations before convergence

% things that matter if we are iterating to estimate parameters
if V.fast_iter_max>1;
    if V.fast_poiss || V.fast_nonlin,
        disp('\ncode does not currrently support estimating parameters for \npoisson or nonlinear observations');
        V.fast_iter_max=1;
    end

    if ~isfield(V,'fast_plot'), V.fast_plot = 0; end
    if V.fast_plot==1
        FigNum = 400;
        if V.Npixels>1, figure(FigNum), clf, end        % figure showing estimated spatial filter
        figure(FigNum+1), clf                           % figure showing estimated spike trains
        if isfield(V,'n'), siz=size(V.n); V.n(V.n==0)=NaN; if siz(1)<siz(2), V.n=V.n'; end; end
    end

    if ~isfield(V,'est_sig'),   V.est_sig   = 0; end    % whether to estimate sig
    if ~isfield(V,'est_lam'),   V.est_lam   = 0; end    % whether to estimate lam
    if ~isfield(V,'est_gam'),   V.est_gam   = 0; end    % whether to estimate gam
    if ~isfield(V,'est_a'),     V.est_a     = 0; end    % whether to estimate a
    if ~isfield(V,'est_b'),     V.est_b     = 1; end    % whether to estimate b
    if ~isfield(V,'fast_plot'), V.fast_plot = 1; end    % whether to plot results from each iteration
    if ~isfield(V,'fast_thr'),  V.fast_thr  = 0; end    % whether to threshold spike train before estimating 'a' and 'b'
    if ~isfield(V,'fast_ignore_post'), V.fast_ignore_post=0; end % whether to ignore the posterior, and just keep the last iteration
end

% normalize F if it is only a trace
if V.Npixels==1
    F=detrend(F);
    F=F-min(F)+eps;
end

%% set default model Parameters

if nargin < 3,          P       = struct;                       end
if ~isfield(P,'sig'),   P.sig   = mean(mad(F',1)*1.4826);       end
if ~isfield(P,'gam'),   P.gam   = (1-V.dt/1)*ones(V.Ncells,1);  end
if ~isfield(P,'lam'),   P.lam   = 10*ones(V.Ncells,1);          end
if ~isfield(P,'a'),     P.a     = median(F,2);                  end

if ~isfield(P,'b'),
    if V.Npixels==1, P.b = quantile(F,0.05);
    else P.b=median(F,2);
    end
end

%% define some stuff needed for est_MAP function

% for brevity and expediency
Z   = zeros(V.Ncells*V.T,1);                    % zero vector
M   = spdiags([repmat(-P.gam,V.T,1) repmat(Z,1,V.Ncells-1) (1+Z)], -V.Ncells:0,V.Ncells*V.T,V.Ncells*V.T);  % matrix transforming calcium into spikes, ie n=M*C
I   = speye(V.Ncells*V.T);                      % create out here cuz it must be reused
H1  = I;                                        % initialize memory for Hessian matrix
H2  = I;                                        % initialize memory for other part of Hessian matrix
d0  = 1:V.Ncells*V.T+1:(V.Ncells*V.T)^2;        % index of diagonal elements of TxT matrices
d1  = 1+V.Ncells:V.Ncells*V.T+1:(V.Ncells*V.T)*(V.Ncells*(V.T-1)); % index of off-diagonal elements of TxT matrices
posts = Z(1:V.fast_iter_max);                   % initialize likelihood
if numel(P.lam)==V.Ncells
    lam = V.dt*repmat(P.lam,V.T,1);             % for lik
elseif numel(P.lam)==V.Ncells*V.T
    lam = V.dt*P.lam;
else
    error('lam must either be length V.T or 1');
end

if V.fast_poiss==1
    H       = I;                                % initialize memory for Hessian matrix
    gamlnF  = gammaln(F+1);                     % for lik
    sumF    = sum(F,1)';                        % for grad & Hess
end

%% if not iterating to estimate parameters, only this is necessary
[n C posts(1)] = est_MAP(F,P);
n_best = n;
P_best = P;
V.fast_iter_tot = 1;
V.post = posts(1);
post_max = posts(1);

if V.fast_iter_max>1
    options = optimset('Display','off');        % don't show warnings for parameter estimation
    i       = 1;                                % iteration #
    i_best  = i;                                % iteration with highest likelihood
    conv    = 0;                                % whether algorithm has converged yet
else
    conv    = 1;
end

%%  if parameters are unknown, do pseudo-EM iterations
while conv == 0
    if V.fast_plot == 1, MakePlot(n,F,P,V); end % plot results from previous iteration
    i               = i+1;                      % update iteratation number
    V.fast_iter_tot = i;                        % record of total # of iterations
    P               = est_params(n,C,F,P,b);    % update parameters based on previous iteration
    [n C posts(i)]  = est_MAP(F,P);             % update inferred spike train based on new parameters

    if posts(i)>post_max || V.fast_ignore_post==1% if this is the best one, keep n and P
        n_best  = n;                            % keep n
        P_best  = P;                            % keep P
        i_best  = i;                            % keep track of which was best
        post_max= posts(i);                     % keep max posterior
    end

    % if lik doesn't change much (relatively), or returns to some previous state, stop iterating
    if  i>=V.fast_iter_max || (abs((posts(i)-posts(i-1))/posts(i))<1e-3 || any(posts(1:i-1)-posts(i))<1e-5)% abs((posts(i)-posts(i-1))/posts(i))<1e-5 || posts(i-1)-posts(i)>1e5;
        MakePlot(n,F,P,V);
        disp('convergence criteria met')
        V.post  = posts(1:i);
        conv    = 1;
    end
    sound(3*sin(linspace(0,90*pi,2000)))        % play sound to indicate iteration is over
end

V.fast_time = cputime-starttime;                % time to run code
V           = orderfields(V);                   % order fields alphabetically to they are easier to read
P_best      = orderfields(P_best);
% n_best      = n_best./repmat(max(n_best),V.T,1);

P_best.j=j;

%% fast filter function
    function [n C post] = est_MAP(F,P)

        % initialize n and C
        z = 1;                                  % weight on barrier function
        llam = reshape(1./lam',1,V.Ncells*V.T)';
        if V.fast_nonlin==1
            n = V.gauss_n;
        else
            n = 0.01+0*llam;                    % initialize spike train
        end
        C = 0*n;                                % initialize calcium
        for j=1:V.Ncells
            C(j:V.Ncells:end) = filter(1,[1, -P.gam(j)],n(j:V.Ncells:end)); %(1-P.gam(j))*P.b(j);
        end

        % precompute parameters required for evaluating and maximizing likelihood
        b           = repmat(P.b,1,V.T);       % for lik
        if V.fast_poiss==1
            suma    = sum(P.a);                 % for grad
        else
            M(d1)   = -repmat(P.gam,V.T-1,1);   % matrix transforming calcium into spikes, ie n=M*C
            ba      = P.a'*b; ba=ba(:);         % for grad
            aa      = repmat(diag(P.a'*P.a),V.T,1);% for grad
            aF      = P.a'*F; aF=aF(:);         % for grad
            e       = 1/(2*P.sig^2);            % scale of variance
            H1(d0)  = -2*e*aa;                   % for Hess
        end
        grad_lnprior  = M'*llam;                  % for grad

        % find C = argmin_{C_z} lik + prior + barrier_z
        while z>1e-13                           % this is an arbitrary threshold

            if V.fast_poiss==1
                Fexpect = P.a*(C+b')';          % expected poisson observation rate
                lik = -sum(sum(-Fexpect+ F.*log(Fexpect) - gamlnF)); % lik
            else
                if V.fast_nonlin==1
                    S = C./(C+P.k_d);
                else
                    S = C;
                end
                D = F-P.a*(reshape(S,V.Ncells,V.T))-b; % difference vector to be used in likelihood computation
                lik = e*D(:)'*D(:);             % lik
            end
            post = lik + llam'*n - z*sum(log(n));
            s    = 1;                           % step size
            d    = 1;                           % direction
            while norm(d)>5e-2 && s > 1e-3      % converge for this z (again, these thresholds are arbitrary)
                if V.fast_poiss==1
                    glik    = suma - sumF./(C+b');
                    H1(d0)  = sumF.*(C+b').^(-2); % lik contribution to Hessian
                elseif V.fast_nonlin==1
                    glik    = -2*P.a*P.k_d*D'.*(C+P.k_d).^-2;
                    H1diag  = (-P.a*P.k_d-2*(C+P.k_d).*D').*((C+P.k_d).^-4);
                    H1(d0)  = H1diag;
                else
                    glik    = -2*e*(aF-aa.*C-ba);  % gradient
                end
                g       = glik + grad_lnprior - z*M'*(n.^-1);
                H2(d0)  = n.^-2;                % log barrier part of the Hessian
                H       = H1 - z*(M'*H2*M);     % Hessian
                d   = H\g;                     % direction to step using newton-raphson
                hit = -n./(M*d);                % step within constraint boundaries
                hit=hit(hit>0);
                if any(hit<1)
                    s = min(1,0.99*min(hit));
                else
                    s = 1;
                end
                post1 = post+1;
                while post1>=post+1e-7          % make sure newton step doesn't increase objective
                    C1  = C+s*d;
                    n   = M*C1;
                    if V.fast_poiss==1
                        Fexpect = P.a*(C1+b')';
                        lik1    = -sum(sum(-Fexpect+ F.*log(Fexpect) - gamlnF));
                    else
                        if V.fast_nonlin==1
                            S1 = C1./(C1+P.k_d);
                        else
                            S1 = C1;
                        end
                        D = F-P.a*(reshape(S1,V.Ncells,V.T))-b; % difference vector to be used in likelihood computation
                        lik1 = e*D(:)'*D(:);             % lik
                    end
                    post1 = lik1 + llam'*n - z*sum(log(n));
                    s   = s/5;                  % if step increases objective function, decrease step size

                    if s<1e-20; disp('reducing s further did not increase likelihood'), break; end      % if decreasing step size just doesn't do it
                end
                C    = C1;                      % update C
                post = post1;                   % update post
            end
            z=z/10;                             % reduce z (sequence of z reductions is arbitrary)
        end

        % reshape things in the case of multiple neurons within the ROI
        n=reshape(n,V.Ncells,V.T)';
        C=reshape(C,V.Ncells,V.T)';
    end

%% Parameter Update
    function P = est_params(n,C,F,P,b)

        % generate regressor for spatial filter
        if V.est_a==1 || V.est_b==1
            if V.fast_thr==1
                CC=0*C;
                for j=1:V.Ncells
                    nsort   = sort(n(:,j));
                    nthr    = nsort(round(0.98*V.T));
                    nn      = Z(1:V.T);
                    nn(n(:,j)<=nthr)=0;
                    nn(n(:,j)>nthr)=1;
                    CC(:,j) = filter(1,[1 -P.gam(j)],nn) + (1-P.gam(j))*P.b(j);
                end
            else
                CC      = C;
            end

            if V.est_b==1
                A = [CC -1+Z(1:V.T)];
            else
                A=CC;
            end
            X = A\F';

            P.a = X(1:V.Ncells,:)';
            if V.est_b==1
                P.b = X(end,:)';
                b   = repmat(P.b,1,V.T);
            end

            D   = F-P.a*(reshape(C,V.Ncells,V.T)) - b;

            mse = D(:)'*D(:);
        end

        if V.est_a==0 && V.est_b==0 && (V.est_sig==1 || V.est_lam==1),
            D   = F-P.a*(reshape(C,V.Ncells,V.T)+b);
            mse = D(:)'*D(:);
        end

        % estimate other parameters
        if V.est_sig==1,
            P.sig = sqrt(mse)/V.T;
        end
        if V.est_lam==1,
            nnorm   = n./repmat(max(n),V.T,1);
            if numel(P.lam)==V.Ncells
                P.lam   = sum(nnorm)'/(V.T*V.dt);
                lam     = repmat(P.lam,V.T,1)*V.dt;
            else
                P.lam   = nnorm/(V.T*V.dt);
                lam     = P.lam*V.dt;
            end

        end
    end

%% MakePlot
    function MakePlot(n,F,P,V)
        if V.fast_plot == 1
            if V.Npixels>1                                     % plot spatial filter
                figure(FigNum), nrows=V.Ncells;
                for j=1:V.Ncells, subplot(1,nrows,j),
                    imagesc(reshape(P.a(:,j),V.w,V.h)),
                    title('a')
                end
            end

            figure(FigNum+1),  ncols=V.Ncells; nrows=3; END=V.T; h=zeros(V.Ncells,2);
            for j=1:V.Ncells                                  % plot inferred spike train
                h(j,1)=subplot(nrows,ncols,(j-1)*ncols+1); cla
                if V.Npixels>1, Ftemp=mean(F); else Ftemp=F; end
                plot(z1(Ftemp(2:END))+1), hold on,
                bar(z1(n_best(2:END,j)))
                title(['best iteration ' num2str(i_best)]),
                axis('tight')
                set(gca,'XTickLabel',[],'YTickLabel',[])

                h(j,2)=subplot(nrows,ncols,(j-1)*ncols+2); cla
                bar(z1(n(2:END,j)))
                if isfield(V,'n'), hold on,
                    for k=1:V.Ncells
                        stem(V.n(2:END,k)+k/10,'LineStyle','none','Marker','v','MarkerEdgeColor','k','MarkerFaceColor','k','MarkerSize',2)
                    end
                end
                set(gca,'XTickLabel',[],'YTickLabel',[])
                title(['current iteration ' num2str(i)]),
                axis('tight')
            end

            subplot(nrows,ncols,j*nrows),
            plot(1:i,posts(1:i))    % plot record of likelihoods
            title(['max lik ' num2str(post_max,4), ',   lik ' num2str(posts(i),4)])
            set(gca,'XTick',2:i,'XTickLabel',2:i)
            drawnow
        end
    end
end


function [u] = spm_uc_RF(a,df,STAT,R,n)
% Corrected critical height threshold at a specified significance level
% FORMAT [u] = spm_uc_RF(a,df,STAT,R,n)
% a     - critical probability - {alpha}
% df    - [df{interest} df{residuals}]
% STAT  - Statistical field
%         'Z' - Gaussian field
%         'T' - T field
%         'X' - Chi-squared field
%         'F' - F field
% R     - RESEL Count {defining search volume}
% n     - number of conjoint SPMs
%
% u     - critical height {corrected}
%
%__________________________________________________________________________
%
% spm_uc returns the corrected critical threshold at a specified significance
% level (a). If n > 1 a conjunction probability over the n values of the
% statistic is returned.
%__________________________________________________________________________
% Copyright (C) 1999-2012 Wellcome Trust Centre for Neuroimaging

% Karl Friston
% $Id: spm_uc_RF.m 4634 2012-02-01 19:01:09Z guillaume $

%-Find approximate value
%--------------------------------------------------------------------------
u  = spm_u((a/max(R))^(1/n),df,STAT);
du = 1e-6;

%-Approximate estimate using E{m}
%--------------------------------------------------------------------------
d  = 1;
while abs(d) > 1e-6
    [P, P, p] = spm_P_RF(1,0,u,df,STAT,R,n);
    [P, P, q] = spm_P_RF(1,0,u + du,df,STAT,R,n);
    d         = (a - p)/((q - p)/du);
    u         = u + d;
    if isinf(u), u=+Inf; return; end
end

%-Refined estimate using 1 - exp(-E{m})
%--------------------------------------------------------------------------
d  = 1;
while abs(d) > 1e-6
    p         = spm_P_RF(1,0,u,df,STAT,R,n);
    q         = spm_P_RF(1,0,u + du,df,STAT,R,n);
    d         = (a - p)/((q - p)/du);
    u         = u + d;
    if isinf(u), u=+Inf; return; end
end


function code = tree_rooted_code ( nnode, parent )

%*****************************************************************************80
%
%% TREE_ROOTED_CODE returns the code of a rooted tree.
%
%  Discussion:
%
%    This code for a rooted tree depends on the node ordering, so it's actually
%    the code for a labeled rooted tree.  To eliminate the effects of node
%    labeling, one could choose as the code for a tree the maximum of all
%    the codes associated with the different possible labelings of the tree.
%    There are more effective ways of arriving at this code than simply
%    generating all possible codes and comparing them.
%
%    For a tree with NNODES, the code is a list of 2*NNODE 0's and 1's,
%    describing a traversal of the tree starting at an imaginary node 0,
%    moving "down" to the root (a code entry of 1), and then moving
%    "down" (1) or "up" (0) as the tree is traversed in a depth first
%    manner.  The final move must be from the root up to the imaginary
%    node 0.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    28 June 2013
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer NNODE, the number of nodes.
%
%    Input, integer PARENT(NNODE), is the parent node of each node.
%    The node with parent 0 is the root.
%
%    Output, integer CODE(2*NNODE), the code for the tree.
%

%
%  Find the root.
%
  father = 0;
  for i = 1 : nnode
    if ( parent(i) == 0 )
      k = 1;
      code(1) = 1;
      father = i;
      break
    end
  end

  if ( father == 0 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'TREE_ROOTED_CODE - Fatal error!\n' );
    fprintf ( 1, '  Could not find the root.\n' );
    error ( 'TREE_ROOTED_CODE - Fatal error!' );
  end

  while ( father ~= 0 )

    k = k + 1;
    code(k) = 0;

    for son = 1 : nnode
      if ( parent(son) == father )
        code(k) = 1;
        father = son;
        break
      end
    end

    if ( code(k) == 0 )
      parent(father) = - parent(father);
      father = - parent(father);
    end

  end

  parent(1:nnode) = - parent(1:nnode);

  return
end


function [ less, equal, more ] = r8r8vec_index_search ( n, x, y, indx, ...
  xval, yval )

%*****************************************************************************80
%
%% R8R8VEC_INDEX_SEARCH searches for an R8R8 value in an indexed sorted list.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    14 May 2005
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer N, the size of the current list.
%
%    Input, real X(N), Y(N), the list.
%
%    Input, integer INDX(N), the sort index of the list.
%
%    Input, real XVAL, YVAL, the value to be sought.
%
%    Output, integer LESS, EQUAL, MORE, the indexes in INDX of the
%    entries of X that are just less than, equal to, and just greater
%    than XVAL.  If XVAL does not occur in X, then EQUAL is zero.
%    If XVAL is the minimum entry of X, then LESS is 0.  If XVAL
%    is the greatest entry of X, then MORE is N+1.
%
  if ( n <= 0 )
    less = 0;
    equal = 0;
    more = 0;
    return
  end

  lo = 1;
  hi = n;

  xlo = x(indx(lo));
  ylo = y(indx(lo));

  xhi = x(indx(hi));
  yhi = y(indx(hi));

  compare = r8r8_compare ( xval, yval, xlo, ylo );

  if ( compare == -1 )
    less = 0;
    equal = 0;
    more = 1;
    return
  elseif ( compare == 0 )
    less = 0;
    equal = 1;
    more = 2;
    return
  end

  compare = r8r8_compare ( xval, yval, xhi, yhi );

  if ( compare == 1 )
    less = n;
    equal = 0;
    more = n + 1;
    return
  elseif ( compare == 0 )
    less = n - 1;
    equal = n;
    more = n + 1;
    return
  end

  while ( 1 )

    if ( lo + 1 == hi )
      less = lo;
      equal = 0;
      more = hi;
      return
    end

    mid = round ( ( lo + hi ) / 2 );
    xmid = x(indx(mid));
    ymid = y(indx(mid));

    compare = r8r8_compare ( xval, yval, xmid, ymid );

    if ( compare == 0 )
      equal = mid;
      less = equal - 1;
      more = equal + 1;
      return
    elseif ( compare == -1 )
      hi = mid;
    elseif ( compare == +1 )
      lo = mid;
    end

  end

  return
end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%	SurfBox-MATLAB (c)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%
%%	Yue M. Lu and Minh N. Do
%%
%%	Department of Electrical and Computer Engineering
%%	Coordinated Science Laboratory
%%	University of Illinois at Urbana-Champaign
%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%
%%	surf_vec2coeff.m
%%
%%	First created: 04-06-07
%%	Last modified: 04-06-07
%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function Y = surf_coeff2vec(c, size_info)

%% Convert the surfacelet coefficient stored in a linear vector to a nested
%% cell array.
%%
%% Input:
%%
%% c: a linear vector storing all the coefficients.
%%
%% size_info: a book keeping cell array storing the dimension information
%% of Y. This is useful when we want to convert c back to Y. See
%% surf_vec2coeff.m for details.
%%
%% Output:
%%
%% Y: an L+1 by 1 cell array containing the surfacelet coefficients. See
%% Surfdec.m for details
%%

%% we can inherit the nested structure
Y = size_info;

%% number of multiscale levels
L = length(Y) - 1;

Ncoeff = 1;

for n = 1 : L
    for m = 1 : length(Y{n})
        for k = 1 : length(Y{n}{m})
            sz = size_info{n}{m}{k};
            d = prod(sz);
            Y{n}{m}{k} = reshape(c(Ncoeff : Ncoeff + d - 1), sz);
            Ncoeff = Ncoeff + d;
        end
    end
end

%% the lowpass subband
sz = size_info{end};
Y{end} = reshape(c(Ncoeff : end), sz);


function [ r, s1, s2, s3 ] = r8_random ( s1, s2, s3 )

%*****************************************************************************80
%
%% R8_RANDOM returns a pseudorandom number between 0 and 1.
%
%  Discussion:
%
%    This function returns a pseudo-random number rectangularly distributed
%    between 0 and 1.   The cycle length is 6.95E+12.  (See page 123
%    of Applied Statistics (1984) volume 33), not as claimed in the
%    original article.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    08 July 2008
%
%  Author:
%
%    Original FORTRAN77 version by Brian Wichman, David Hill.
%    MATLAB version by John Burkardt.
%
%  Reference:
%
%    Brian Wichman, David Hill,
%    Algorithm AS 183: An Efficient and Portable Pseudo-Random
%    Number Generator,
%    Applied Statistics,
%    Volume 31, Number 2, 1982, pages 188-190.
%
%  Parameters:
%
%    Input, integer S1, S2, S3, three values used as the
%    seed for the sequence.  These values should be positive
%    integers between 1 and 30,000.
%
%    Output, real R, the next value in the sequence.
%
%    Output, integer S1, S2, S3, updated seed values.
%
  s1 = mod ( 171 * s1, 30269 );
  s2 = mod ( 172 * s2, 30307 );
  s3 = mod ( 170 * s3, 30323 );

  r = mod ( s1 / 30269.0 ...
          + s2 / 30307.0 ...
          + s3 / 30323.0, 1.0 );

  return
end


function [ t, rank ] = subset_colex_successor ( n, t, rank )

%*****************************************************************************80
%
%% SUBSET_COLEX_SUCCESSOR computes the subset colexicographic successor.
%
%  Discussion:
%
%    In the original code, there is a last element with no successor.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    22 August 2011
%
%  Author:
%
%    John Burkardt
%
%  Reference:
%
%    Donald Kreher, Douglas Simpson,
%    Combinatorial Algorithms,
%    CRC Press, 1998,
%    ISBN: 0-8493-3988-X,
%    LC: QA164.K73.
%
%  Parameters:
%
%    Input, integer N, the number of elements in the master set.
%    N must be positive.
%
%    Input/output, integer T(N), describes a subset.  T(I) is 0 if
%    the I-th element of the master set is not in the subset, and is
%    1 if the I-th element is part of the subset.
%    On input, T describes a subset.
%    On output, T describes the next subset in the ordering.
%    If the input T was the last in the ordering, then the output T
%    will be the first.
%
%    Input/output, integer RANK, the rank.
%    If RANK = -1 on input, then the routine understands that this is
%    the first call, and that the user wishes the routine to supply
%    the first element in the ordering, which has RANK = 0.
%    In general, the input value of RANK is increased by 1 for output,
%    unless the very last element of the ordering was input, in which
%    case the output value of RANK is 0.
%

%
%  Return the first element.
%
  if ( rank == -1 )
    t(1:n) = 0;
    rank = 0;
    return
  end
%
%  Check.
%
  subset_check ( n, t );

  for i = 1 : n

    if ( t(i) == 0 )
      t(i) = 1;
      rank = rank + 1;
      return
    else
      t(i) = 0;
    end

  end

  rank = 0;

  return
end


function c = tapas_hgf_ar1_binary_config
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% Contains the configuration for the Hierarchical Gaussian Filter (HGF) for AR(1) processes
% for binary inputs in the absence of perceptual uncertainty.
%
% The HGF is the model introduced in
%
% Mathys C, Daunizeau J, Friston, KJ, and Stephan KE. (2011). A Bayesian foundation
% for individual learning under uncertainty. Frontiers in Human Neuroscience, 5:39.
%
% The binary HGF model has since been augmented with a positive factor kappa1 which
% scales the second level with respect to the first, i.e., the relation between the
% first and second level is
%
% p(x1=1|x2) = s(kappa1*x2), where s(.) is the logistic sigmoid.
%
% By default, kappa1 is fixed to 1, leading (apart from the AR(1) process) to the
% model introduced in Mathys et al. (2011).
%
% This file refers to BINARY inputs (Eqs 1-3 in Mathys et al., (2011));
% for continuous inputs, refer to tapas_hgf_config.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% The HGF configuration consists of the priors of parameters and initial values. All priors are
% Gaussian in the space where the quantity they refer to is estimated. They are specified by their
% sufficient statistics: mean and variance (NOT standard deviation).
%
% Quantities are estimated in their native space if they are unbounded (e.g., the omegas). They are
% estimated in log-space if they have a natural lower bound at zero (e.g., the sigmas).
%
% The phis are estimated in 'logit space' because they are confined to the interval from 0 to 1.
% 'Logit-space' is a logistic sigmoid transformation of native space with a variable upper bound
% a>0:
%
% tapas_logit(x) = ln(x/(a-x)); x = a/(1+exp(-tapas_logit(x)))
%
% Parameters can be fixed (i.e., set to a fixed value) by setting the variance of their prior to
% zero. Aside from being useful for model comparison, the need for this arises whenever the scale
% and origin at the j-th level are arbitrary. This is the case if the observation model does not
% contain the representations mu_j and sigma_j. A choice of scale and origin is then implied by
% fixing the initial value mu_j_0 of mu_j and either kappa_j-1 or omega_j-1.
%
% Fitted trajectories can be plotted by using the command
%
% >> tapas_hgf_binary_plotTraj(est)
%
% where est is the stucture returned by tapas_fitModel. This structure contains the estimated
% perceptual parameters in est.p_prc and the estimated trajectories of the agent's
% representations (cf. Mathys et al., 2011). Their meanings are:
%
%         est.p_prc.mu_0       row vector of initial values of mu (in ascending order of levels)
%         est.p_prc.sa_0       row vector of initial values of sigma (in ascending order of levels)
%         est.p_prc.phi        row vector of phis (representing reversion slope to attractor; in ascending order of levels)
%         est.p_prc.m        row vector of ms (representing attractors; in ascending order of levels)
%         est.p_prc.ka         row vector of kappas (in ascending order of levels)
%         est.p_prc.om         row vector of omegas (in ascending order of levels)
%
% Note that the first entry in all of the row vectors will be NaN because, at the first level,
% these parameters are either determined by the second level (mu_0 and sa_0) or undefined (rho,
% kappa, and omega).
%
%         est.traj.mu          mu (rows: trials, columns: levels)
%         est.traj.sa          sigma (rows: trials, columns: levels)
%         est.traj.muhat       prediction of mu (rows: trials, columns: levels)
%         est.traj.sahat       precisions of predictions (rows: trials, columns: levels)
%         est.traj.v           inferred variance of random walk (rows: trials, columns: levels)
%         est.traj.w           weighting factors (rows: trials, columns: levels)
%         est.traj.da          volatility prediction errors  (rows: trials, columns: levels)
%         est.traj.ud          updates with respect to prediction  (rows: trials, columns: levels)
%         est.traj.psi         precision weights on prediction errors  (rows: trials, columns: levels)
%         est.traj.epsi        precision-weighted prediction errors  (rows: trials, columns: levels)
%         est.traj.wt          full weights on prediction errors (at the first level,
%                                  this is the learning rate) (rows: trials, columns: levels)
%
% Note that in the absence of sensory uncertainty (which is the assumption here), the first
% column of mu, corresponding to the first level, will be equal to the inputs. Likewise, the
% first column of sa will be 0 always.
%
% Tips:
% - When analyzing a new dataset, take your inputs u and responses y and use
%
%   >> est = tapas_fitModel(y, u, 'tapas_hgf_binary_config', 'tapas_bayes_optimal_binary_config');
%
%   to determine the Bayes optimal perceptual parameters (given your current priors as defined in
%   this file here, so choose them wide and loose to let the inputs influence the result). You can
%   then use the optimal parameters as your new prior means for the perceptual parameters.
%
% - If you get an error saying that the prior means are in a region where model assumptions are
%   violated, lower the prior means of the omegas, starting with the highest level and proceeding
%   downwards.
%
% - Alternatives are lowering the prior means of the kappas, if they are not fixed, or adjusting
%   the values of the kappas or omegas, if any of them are fixed.
%
% - If the log-model evidence cannot be calculated because the Hessian poses problems, look at
%   est.optim.H and fix the parameters that lead to NaNs.
%
% - Your guide to all these adjustments is the log-model evidence (LME). Whenever the LME increases
%   by at least 3 across datasets, the adjustment was a good idea and can be justified by just this:
%   the LME increased, so you had a better model.
%
% --------------------------------------------------------------------------------------------------
% Copyright (C) 2012-2017 Christoph Mathys, TNU, UZH & ETHZ
%
% This file is part of the HGF toolbox, which is released under the terms of the GNU General Public
% Licence (GPL), version 3. You can redistribute it and/or modify it under the terms of the GPL
% (either version 3 or, at your option, any later version). For further details, see the file
% COPYING or <http://www.gnu.org/licenses/>.

% Config structure
c = struct;

% Model name
c.model = 'tapas_hgf_ar1_binary';

% Number of levels (minimum: 3)
c.n_levels = 3;

% Input intervals
% If input intervals are irregular, the last column of the input
% matrix u has to contain the interval between inputs k-1 and k
% in the k-th row, and this flag has to be set to true
c.irregular_intervals = false;

% Sufficient statistics of Gaussian parameter priors

% Initial mus and sigmas
% Format: row vectors of length n_levels
% For all but the first two levels, this is usually best
% kept fixed to 1 (determines origin on x_i-scale). The
% first level is NaN because it is determined by the second,
% and the second implies neutrality between outcomes when it
% is centered at 0.
c.mu_0mu = [NaN, 0, 1];
c.mu_0sa = [NaN, 0, 0];

c.logsa_0mu = [NaN, log(0.006), log(4)];
c.logsa_0sa = [NaN,          0,      0];

% Phis
% Format: row vector of length n_levels.
% Undefined (therefore NaN) at the first level.
% Fix this to zero (-Inf in logit space) to set to zero.
c.logitphimu = [NaN, -Inf, tapas_logit(0.1,1)];
c.logitphisa = [NaN,    0,                  2];

% ms
% Format: row vector of length n_levels.
% This should be fixed for all levels where the omega of
% the next lowest level is not fixed because that offers
% an alternative parametrization of the same model.
c.mmu = [NaN, c.mu_0mu(2), c.mu_0mu(3)];
c.msa = [NaN,           0,           1];

% Kappas
% Format: row vector of length n_levels-1.
% Fixing log(kappa1) to log(1) leads to the original HGF model.
% Higher log(kappas) should be fixed (preferably to log(1)) if the
% observation model does not use mu_i+1 (kappa then determines the
% scaling of x_i+1).
c.logkamu = [log(1), log(1)];
c.logkasa = [     0,      0];

% Omegas
% Format: row vector of length n_levels.
% Undefined (therefore NaN) at the first level.
c.ommu = [NaN,  -2,  -6];
c.omsa = [NaN, 4^2, 4^2];

% Gather prior settings in vectors
c.priormus = [
    c.mu_0mu,...
    c.logsa_0mu,...
    c.logitphimu,...
    c.mmu,...
    c.logkamu,...
    c.ommu,...
         ];

c.priorsas = [
    c.mu_0sa,...
    c.logsa_0sa,...
    c.logitphisa,...
    c.msa,...
    c.logkasa,...
    c.omsa,...
         ];

% Check whether we have the right number of priors
expectedLength = 5*c.n_levels+(c.n_levels-1);
if length([c.priormus, c.priorsas]) ~= 2*expectedLength;
    error('tapas:hgf:PriorDefNotMatchingLevels', 'Prior definition does not match number of levels.')
end

% Model function handle
c.prc_fun = @tapas_hgf_ar1_binary;

% Handle to function that transforms perceptual parameters to their native space
% from the space they are estimated in
c.transp_prc_fun = @tapas_hgf_ar1_binary_transp;

return;


function [prob,sol,fmin] = nls_prob(varargin)
%NLS_PROB  Return an OPTI NLS
%
%   prob = nls_prob(no) return a pre-built optiprob of a saved NLS.
%
%   [prob,sol,fmin] = NLS_prob(no) returns the optimum solution and function
%   eval at the optimum
%
%   no = nls_prob() returns the number of problems available for testing.

%   (C) 2012 Jonathan Currie (IPL)

% Functions are taken from:
% More, J. J., B. S. Garbow, and K. E. Hillstrom. "Testing Unconstrained
% Optimization Software." ACM Transactions on Mathematical Software 7,
% no. 1 (1981): 17-41.

%Check if just returning no problems
if(nargin < 1)
    prob = 20; sol = []; fmin = [];
    return;
else
    no = varargin{1};
end

%Big switch yard
switch(no)
    case 1
        fun = @(x) [10*(x(2)-x(1)^2); 1 - x(1)];
        ydata = zeros(2,1);
        x0 = [-1.2;1];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1,1];
        fmin = 0;

    case 2
        fun = @(x) [-13 + x(1) + ((5 - x(2))*x(2) - 2)*x(2);
                    -29 + x(1) + ((x(2) + 1)*x(2) - 14)*x(2)];
        ydata = zeros(2,1);
        x0 = [0.5;-2];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [5,4];
        fmin = 0;

    case 3
        fun = @(x) [1e4*x(1)*x(2) - 1;
                    exp(-x(1)) + exp(-x(2)) - 1.0001];
        ydata = zeros(2,1);
        x0 = [0;1];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1.098e-5;9.106];
        fmin = 0;

    case 4
        fun = @(x) [x(1) - 1e6; x(2) - 2e-6; x(1)*x(2) - 2];
        ydata = zeros(3,1);
        x0 = [1;1];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1e6,2e-6];
        fmin = 0;

    case 5
        fun = @(x) [1.5 - x(1)*(1 - x(2)); 2.25 - x(1)*(1 - x(2)^2); 2.625 - x(1)*(1 - x(2)^3) ];
        ydata = zeros(3,1);
        x0 = [1;1];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [3;0.5];
        fmin = 0;

    case 6
        i = (1:10)';
        fun = @(x) 2 + 2*i - (exp(x(1)*i) + exp(x(2)*i));
        ydata = zeros(10,1);
        x0 = [0.3;0.4];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [0.2578;0.2578];
        fmin = 124.36226865;

    case 7
        fun = @more7;
        ydata = zeros(3,1);
        x0 = [-1;0;0];
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1;0;0];
        fmin = 0;

    case 8
        u = (1:15)';
        v = 16 - u;
        w = zeros(15,1);
        for i = 1:15
            w(i) = min(u(i),v(i));
        end
        y = [0.14;0.18;0.22;0.25;0.29;0.32;0.35;0.39;0.37;0.58;0.73;0.96;1.34;2.10;4.39];
        fun = @(x) (x(1) + u ./ (v*x(2) + w*x(3)));
        x0 = [1;1;1];
        prob = optiprob('fun',fun,'ydata',y,'x0',x0);
        sol = [];
        fmin = 8.21487e-3;

    case 9
        i = (1:15)';
        y = [0.0009;0.0044;0.0175;0.0540;0.1295;0.2420;0.3521;0.3989;0.3521;0.2420;0.1295;0.0540;0.0175;0.0044;0.0009];
        t = (8 - i)./2;
        fun = @(x) x(1)*exp((-x(2)*(t - x(3)).^2)./2);
        x0 = [0.4;1;0];
        prob = optiprob('fun',fun,'ydata',y,'x0',x0);
        sol = [];
        fmin = 1.12793e-8;

    case {10,11}
        i = (1:16)';
        y = [34780;28610;23650;19630;16370;13720;11540;9744;8261;7030;6005;5147;4427;3820;3307;2872];
        t = 45 + 5*i;
        fun = @(x) x(1)*exp(x(2)./(t + x(3)));
        x0 = [0.02;4000;250];
        prob = optiprob('fun',fun,'ydata',y,'x0',x0);
        sol = [];
        fmin = 87.9458;

%     case 11 %not working??
%         m = 100;
%         i = [1:m]';
%         t = i./100;
%         y = 25 + (-50 * log(t)).^(2/3);
%         fun = @(x) exp(-(abs(y*100.*i*x(2)).^x(3))./x(1)) - t;
%         x0 = [5;2.5;0.15];
%         ydata = zeros(m,1);
%         prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
%         sol = [50;25;1.5];
%         fmin = 0;

    case 12
        m = 3;
        i = (1:m)';
        t = 0.1*i;
        fun = @(x) exp(-t*x(1)) - exp(-t*x(2)) - x(3)*exp(-t) - exp(-10*t);
        x0 = [0;10;20];
        ydata = zeros(m,1);
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1;10;1];
        fmin = 0;

    case 13
        fun = @(x) [x(1) + 10*x(2);
                    sqrt(5)*(x(3) - x(4));
                    (x(3) - 2*x(3))^2;
                    sqrt(10)*(x(1)-x(4))^2];
        x0 = [3;-1;0;1];
        ydata = zeros(4,1);
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [0;0;0;0];
        fmin = 0;

    case 14
        fun = @(x) [10*(x(2)-x(1)^2);
                    1 - x(1);
                    sqrt(90)*(x(4)-x(3)^2);
                    1 - x(3);
                    sqrt(10)*(x(2) + x(4) - 2);
                    10^(-0.5)*(x(2)-x(4))];
        x0 = [3;-1;-3;-1];
        ydata = zeros(6,1);
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [1;1;1;1];
        fmin = 0;

    case 15
        y = [0.1957;0.1947;0.1735;0.16;0.0844;0.0627;0.0456;0.0342;0.0323;0.0235;0.0246];
        u = [4;2;1;0.5;0.25;0.167;0.125;0.1;0.0833;0.0714;0.0625];
        fun = @(x) (x(1)*(u.^2 + u*x(2))) ./ (u.^2 + u*x(3) + x(4));
        x0 = [0.25;0.39;0.415;0.39];
        prob = optiprob('fun',fun,'ydata',-y,'x0',x0);
        sol = [];
        fmin = 3.07505e-4;

    case 16
        m = 20;
        i = (1:m)';
        t = i/5;
        fun = @(x) (x(1) + t*x(2) - exp(t)).^2 + (x(3) + x(4)*sin(t) - cos(t)).^2;
        x0 = [25;5;-5;-1];
        ydata = zeros(m,1);
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [];
        fmin = 85822.2016;

    case 17
        m = 33;
        i = (1:m)';
        t = 10*(i - 1);
        y = [0.844;0.908;0.932;0.936;0.925;0.908;0.881;0.850;0.818;0.784;0.751;0.718;0.685;0.658;0.628;0.603;...
             0.58;0.558;0.538;0.522;0.506;0.490;0.478;0.467;0.457;0.448;0.438;0.431;0.424;0.42;0.414;0.411;0.406];
        fun = @(x) (x(1) + x(2)*exp(-t*x(4)) + x(3)*exp(-t*x(5)));
        x0 = [0.5;1.5;-1;0.01;0.02];
        prob = optiprob('fun',fun,'ydata',-y,'x0',x0);
        sol = [];
        fmin = 5.46489e-5;

    case 18
        m = 13;
        i = (1:m)';
        t = 0.1*i;
        y = exp(-t) - 5*exp(-10*t) + 3*exp(-4*t);
        fun = @(x) x(3)*exp(-t*x(1)) - x(4)*exp(-t*x(2)) + x(6)*exp(-t*x(5));
        x0 = [1;2;1;1;1;1];
        prob = optiprob('fun',fun,'ydata',y,'x0',x0);
        sol = [];
        fmin = 0; %5.65565e-3; seems wrong?

    case 19
        m = 65;
        i = (1:m)';
        t = (i-1)/10;
        y = [1.366;1.191;1.112;1.013;0.991;0.885;0.831;0.847;0.786;0.725;0.746;0.679;0.608;0.655;0.616;0.606;0.602;0.626;0.651;0.724;0.649;0.649;...
            0.694;0.644;0.624;0.661;0.612;0.558;0.533;0.495;0.5;0.423;0.395;0.375;0.372;0.391;0.396;0.405;0.428;0.429;0.523;0.562;0.607;0.653;...
            0.672;0.708;0.633;0.668;0.645;0.632;0.591;0.559;0.597;0.625;0.739;0.710;0.729;0.720;0.636;0.581;0.428;0.292;0.162;0.098;0.054];
        fun = @(x) (x(1)*exp(-t*x(5)) + x(2)*exp(-(t-x(9)).^2*x(6)) + x(3)*exp(-(t-x(10)).^2*x(7)) + x(4)*exp(-(t-x(11)).^2*x(8)));
        x0 = [1.3;0.65;0.65;0.7;0.6;3;5;7;2;4.5;5.5];
        prob = optiprob('fun',fun,'ydata',y,'x0',x0);
        sol = [];
        fmin = 4.01377e-2;

    case 20
        m = 31; n = 12;
        fun = @more20;
        x0 = zeros(n,1);
        ydata = zeros(m,1);
        prob = optiprob('fun',fun,'ydata',ydata,'x0',x0);
        sol = [];
        fmin = 4.72238e-10;

    otherwise
        error('Problem not available or not implemented yet');
end

function j = more7(x)

if(x(1) > 0)
    theta = 1/(2*pi)*atan(x(2)/x(1));
else
    theta = 1/(2*pi)*atan(x(2)/x(1)) + 0.5;
end

j(1,1) = 10*(x(3) - 10*theta);
j(2,1) = 10*(x(1)^2 + x(2)^2)^0.5 - 1;
j(3,1) = x(3);

function f = more20(x)
n = length(x);
j = (2:n)';
j2 = (1:n)';
t = 0.1*(1:29)';
f = zeros(31,1);
for k = 1:29
    f(k) = sum( (j-1).*x(2:end).*t(k).^(j-2) ) - sum( x.*t(k).^(j2-1) ) - 1;
end
f(30) = x(1);
f(31) = x(2) - x(1)^2 - 1;


function [blockVectorX,lambda,varargout] = ...
    lobpcg(blockVectorX,operatorA,varargin)
%LOBPCG solves Hermitian partial eigenproblems using preconditioning
%
% [blockVectorX,lambda]=lobpcg(blockVectorX,operatorA)
%
% outputs the array of algebraic smallest eigenvalues lambda and
% corresponding matrix of orthonormalized eigenvectors blockVectorX of the
% Hermitian (full or sparse) operator operatorA using input matrix
% blockVectorX as an initial guess, without preconditioning, somewhat
% similar to
%
% opts.issym=1;opts.isreal=1;K=size(blockVectorX,2);
% [blockVectorX,lambda]=eigs(operatorA,K,'SR',opts);
%
% for real symmetric operator operatorA, or
%
% K=size(blockVectorX,2);[blockVectorX,lambda]=eigs(operatorA,K,'SR');
% for Hermitian operator operatorA.
%
% [blockVectorX,lambda,failureFlag]=lobpcg(blockVectorX,operatorA)
% also returns a convergence flag.
% If failureFlag is 0 then all the eigenvalues converged; otherwise not all
% converged.
%
% [blockVectorX,lambda,failureFlag,lambdaHistory,residualNormsHistory]=...
% lobpcg(blockVectorX,'operatorA','operatorB','operatorT',blockVectorY,...
% residualTolerance,maxIterations,verbosityLevel);
%
% computes smallest eigenvalues lambda and corresponding eigenvectors
% blockVectorX of the generalized eigenproblem Ax=lambda Bx, where
% Hermitian operators operatorA and operatorB are given as functions, as
% well as a preconditioner, operatorT. The operators operatorB and
% operatorT must be in addition POSITIVE DEFINITE. To compute the largest
% eigenpairs of operatorA, simply apply the code to operatorA multiplied by
% -1. The code does not involve ANY matrix factorizations of operratorA and
% operatorB, thus, e.g., it preserves the sparsity and the structure of
% operatorA and operatorB.
%
% residualTolerance and maxIterations control tolerance and max number of
% steps, and verbosityLevel = 0, 1, or 2 controls the amount of printed
% info. lambdaHistory is a matrix with all iterative lambdas, and
% residualNormsHistory are matrices of the history of 2-norms of residuals
%
% Required input:
%  * blockVectorX (class numeric) - initial approximation to eigenvectors,
%    full or sparse matrix n-by-blockSize. blockVectorX must be full rank.
%  * operatorA (class numeric, char, or function_handle) - the main operator
%    of the eigenproblem, can be a matrix, a function name, or handle
%
% Optional function input:
%   * operatorB (class numeric, char, or function_handle) - the second
%     operator, if solving a generalized eigenproblem, can be a matrix,
%      a function name, or handle; by default if empty, operatorB=I.
%   * operatorT  (class char or function_handle) - the preconditioner,
%     by default operatorT(blockVectorX)=blockVectorX.
%
% Optional constraints input:
%   blockVectorY (class numeric) - a full or sparse n-by-sizeY matrix of
%   constraints, where sizeY < n. blockVectorY must be full rank.
%   The iterations will be performed in the (operatorB-)
%   orthogonal complement of the column-space of blockVectorY.
%
% Optional scalar input parameters:
%   residualTolerance (class numeric) - tolerance, by default,
%   residualTolerance=n*sqrt(eps) maxIterations - max number of iterations,
%   by default, maxIterations = min(n,20) verbosityLevel - either 0 (no
%   info), 1, or 2 (with pictures); by default, verbosityLevel = 0.
%
% Required output: blockVectorX and lambda (both class numeric) are
% computed blockSize eigenpairs, where blockSize=size(blockVectorX,2)
% for the initial guess blockVectorX if it is full rank.
%
% Optional output: failureFlag (class integer), lambdaHistory (class numeric)
% and residualNormsHistory (class numeric) are described above.
%
% Functions operatorA(blockVectorX), operatorB(blockVectorX) and
% operatorT(blockVectorX) must support blockVectorX being a matrix, not
% just a column vector.
%
% Every iteration involves one application of operatorA and operatorB, and
% one of operatorT.
%
% Main memory requirements: 6 (9 if isempty(operatorB)=0) matrices of the
% same size as blockVectorX, 2 matrices of the same size as blockVectorY
% (if present), and two square matrices of the size 3*blockSize.
%
% In all examples below, we use the Laplacian operator in a 20x20 square
% with the mesh size 1 which can be generated in MATLAB by running
% A = delsq(numgrid('S',21)); n=size(A,1);
% or in MATLAB and Octave by
% [~,~,A] = laplacian([19,19]); n=size(A,1);
% see http://www.mathworks.com/matlabcentral/fileexchange/27279
%
% The following Example:
%
% [blockVectorX,lambda,failureFlag]=lobpcg(randn(n,8),A,1e-5,50,2);
%
% attempts to compute 8 first eigenpairs without preconditioning,
% but not all eigenpairs converge after 50 steps, so failureFlag=1.
%
% The next Example:
%
% blockVectorY=[];lambda_all=[];
% for j=1:4
%   [blockVectorX,lambda]=...
%                    lobpcg(randn(n,2),A,blockVectorY,1e-5,200,2);
%   blockVectorY=[blockVectorY,blockVectorX];
%   lambda_all=[lambda_all' lambda']'; pause;
% end
%
% attemps to compute the same 8 eigenpairs by calling the code 4 times
% with blockSize=2 using orthogonalization to the previously founded
% eigenvectors.
%
% The following Example:
%
% R=ichol(A,struct('michol','on')); precfun = @(x)R\(R'\x);
% [blockVectorX,lambda,failureFlag]=lobpcg(randn(n,8),A,[],@(x)precfun(x),1e-5,60,2);
%
% computes the same eigenpairs in less then 25 steps, so that failureFlag=0
% using the preconditioner function "precfun", defined inline. If "precfun"
% is defined as a MATLAB function in a file, the function handle
% @(x)precfun(x) can be equivalently replaced by the function name 'precfun'
% Running
%
% [blockVectorX,lambda,failureFlag]=...
%          lobpcg(randn(n,8),A,speye(n),@(x)precfun(x),1e-5,50,2);
%
% produces similar answers, but is somewhat slower and needs more memory as
% technically a generalized eigenproblem with B=I is solved here.
%
% The following Example for a mostly diagonally dominant sparse matrix A
% demonstrates different types of preconditioning, compared to the standard
% use of the main diagonal of A:
%
% clear all; close all;
% n = 1000; M = spdiags([1:n]',0,n,n); precfun=@(x)M\x;
% A=M+sprandsym(n,.1); Xini=randn(n,5); maxiter=15; tol=1e-5;
% [~,~,~,~,rnp]=lobpcg(Xini,A,tol,maxiter,1);
% [~,~,~,~,r]=lobpcg(Xini,A,[],@(x)precfun(x),tol,maxiter,1);
% subplot(2,2,1), semilogy(r'); hold on; semilogy(rnp',':>');
% title('No preconditioning (top)'); axis tight;
% M(1,2) = 2; precfun=@(x)M\x; % M is no longer symmetric
% [~,~,~,~,rns]=lobpcg(Xini,A,[],@(x)precfun(x),tol,maxiter,1);
% subplot(2,2,2),  semilogy(r'); hold on; semilogy(rns','--s');
% title('Nonsymmetric preconditioning (square)'); axis tight;
% M(1,2) = 0; precfun=@(x)M\(x+10*sin(x)); % nonlinear preconditioning
% [~,~,~,~,rnl]=lobpcg(Xini,A,[],@(x)precfun(x),tol,maxiter,1);
% subplot(2,2,3),  semilogy(r'); hold on; semilogy(rnl','-.*');
% title('Nonlinear preconditioning (star)'); axis tight;
% M=abs(M-3.5*speye(n,n)); precfun=@(x)M\x;
% [~,~,~,~,rs]=lobpcg(Xini,A,[],@(x)precfun(x),tol,maxiter,1);
% subplot(2,2,4),  semilogy(r'); hold on; semilogy(rs','-d');
% title('Selective preconditioning (diamond)'); axis tight;
%
% Revision 4.16 adds support for distributed or codistributed arrays
% available in MATLAB BigData toolbox, e.g.,
%
% A = codistributed(diag(1:100)); B = codistributed(diag(101:200));
% [blockVectorX,lambda]=lobpcg(randn(100,2),A,1e-5,5,2)
%
% Revision 4.17 adds support for single precision, e.g.,
% A = diag(1:100); B = single(diag(101:200));
% [blockVectorX,lambda]=lobpcg(randn(100,2),A,1e-5,15,2);
% A = diag(1:100); B = diag(101:200);
% [blockVectorX,lambda]=lobpcg(randn(100,2,'single'),A,1e-5,15,2);
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% This main function LOBPCG is a version of
% the preconditioned conjugate gradient method (Algorithm 5.1) described in
% A. V. Knyazev, Toward the Optimal Preconditioned Eigensolver:
% Locally Optimal Block Preconditioned Conjugate Gradient Method,
% SIAM Journal on Scientific Computing 23 (2001), no. 2, pp. 517-541.
% http://dx.doi.org/10.1137/S1064827500366124
%
% Known bugs/features:
%
% - an excessively small requested tolerance may result in often restarts
% and instability. The code is not written to produce an eps-level
% accuracy! Use common sense.
%
% - the code may be very sensitive to the number of eigenpairs computed,
% if there is a cluster of eigenvalues not completely included, cf.
%
% operatorA=diag([1 1.99 2:99]);
% [blockVectorX,lambda]=lobpcg(randn(100,1),operatorA,1e-10,80,2);
% [blockVectorX,lambda]=lobpcg(randn(100,2),operatorA,1e-10,80,2);
% [blockVectorX,lambda]=lobpcg(randn(100,3),operatorA,1e-10,80,2);
%
% - using a nonsymmetric preconditioner is possible, but may be unstable
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% The main distribution site:
% https://github.com/lobpcg/blopex
%
% A C-version of this code is a part of the
% https://github.com/lobpcg/blopex
% package and is directly available, e.g., in SLEPc and HYPRE.
%
% A python version of this code is in
% https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.lobpcg.html
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   License:  MIT / Apache-2.0
%   Copyright (c) 2000-2019 A.V. Knyazev, Andrew.Knyazev@ucdenver.edu
%   $Revision: 1.2 $  $Date: 13-June-2019
%   This revision is tested in 9.6.0.1114505 (R2019a) Update 2, but is
%   expected to work on any >R2007b MATLAB.
%   Revision 4.13 tested in MATLAB 6.5-7.13.
%   Revision 4.13 tested and available in Octave 3.2.3-3.4.2, see
%   https://octave.sourceforge.io/linear-algebra/function/lobpcg.html
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Begin
% Function gather defined to be identity if nonexistent, before 2016a
if exist("gather", "file") == 2
    mygather=@(x)gather(x);
else
    mygather=@(x)x;
end
% constants
CONVENTIONAL_CONSTRAINTS = 1;
SYMMETRIC_CONSTRAINTS = 2;
%Initial settings
failureFlag = 1;
if nargin < 2
    error('BLOPEX:lobpcg:NotEnoughInputs',...
        strcat('There must be at least 2 input agruments: ',...
        'blockVectorX and operatorA'));
end
if nargin > 8
    warning('BLOPEX:lobpcg:TooManyInputs',...
        strcat('There must be at most 8 input agruments ',...
        'unless arguments are passed to a function'));
end
if ~isnumeric(blockVectorX)
    error('BLOPEX:lobpcg:FirstInputNotNumeric',...
        'The first input argument blockVectorX must be numeric');
end
[n,blockSize]=size(blockVectorX);
if blockSize > n
    error('BLOPEX:lobpcg:FirstInputFat',...
        'The first input argument blockVectorX must be tall, not fat');
end
if n < 6
    error('BLOPEX:lobpcg:MatrixTooSmall',...
        'The code does not work for matrices of small sizes');
end
if isa(operatorA,'numeric')
    nA = size(operatorA,1);
    if any(size(operatorA) ~= nA)
        error('BLOPEX:lobpcg:MatrixNotSquare',...
            'operatorA must be a square matrix or a string');
    end
    if size(operatorA) ~= n
        error('BLOPEX:lobpcg:MatrixWrongSize',...
            ['The size ' int2str(size(operatorA))...
            ' of operatorA is not the same as ' int2str(n)...
            ' - the number of rows of blockVectorX']);
    end
end
count_string = 0;
operatorT = [];
operatorB = [];
residualTolerance = [];
maxIterations = [];
verbosityLevel = [];
blockVectorY = []; sizeY = 0;
for j = 1:nargin-2
    if isequal(size(varargin{j}),[n,n])
        if isempty(operatorB)
            operatorB = varargin{j};
        else
            error('BLOPEX:lobpcg:TooManyMatrixInputs',...
                strcat('Too many matrix input arguments. ',...
                'Preconditioner operatorT must be an M-function'));
        end
    elseif isequal(size(varargin{j},1),n) && size(varargin{j},2) < n
        if isempty(blockVectorY)
            blockVectorY = varargin{j};
            sizeY=size(blockVectorY,2);
        else
            error('BLOPEX:lobpcg:WrongConstraintsFormat',...
                'Something wrong with blockVectorY input argument');
        end
    elseif ischar(varargin{j}) || isa(varargin{j},'function_handle')
        if count_string == 0
            if isempty(operatorB)
                operatorB = varargin{j};
                count_string = count_string + 1;
            else
                operatorT = varargin{j};
            end
        elseif count_string == 1
            operatorT = varargin{j};
        else
            warning('BLOPEX:lobpcg:TooManyStringFunctionHandleInputs',...
                'Too many string or FunctionHandle input arguments');
        end
    elseif isequal(size(varargin{j}),[n,n])
        error('BLOPEX:lobpcg:WrongPreconditionerFormat',...
            'Preconditioner operatorT must be an M-function');
    elseif max(size(varargin{j})) == 1
        if isempty(residualTolerance)
            residualTolerance = varargin{j};
        elseif isempty(maxIterations)
            maxIterations = varargin{j};
        elseif isempty(verbosityLevel)
            verbosityLevel = varargin{j};
        else
            warning('BLOPEX:lobpcg:TooManyScalarInputs',...
                'Too many scalar parameters, need only three');
        end
    elseif isempty(varargin{j})
        if isempty(operatorB)
            count_string = count_string + 1;
        elseif ~isempty(operatorT)
            count_string = count_string + 1;
        elseif ~isempty(blockVectorY)
            error('BLOPEX:lobpcg:UnrecognizedEmptyInput',...
                ['Unrecognized empty input argument number ' int2str(j+2)]);
        end
    else
        error('BLOPEX:lobpcg:UnrecognizedInput',...
            ['Input argument number ' int2str(j+2) ' not recognized.']);
    end
end
if verbosityLevel
    if issparse(blockVectorX)
        fprintf(['The sparse initial guess with %i colunms '...
            'and %i raws is detected  \n'],n,blockSize);
    else
        fprintf(['The full initial guess with %i colunms '...
            'and %i raws is detected  \n'],n,blockSize);
    end
    if ischar(operatorA)
        fprintf('The main operator is detected as an M-function %s \n',...
            operatorA);
    elseif isa(operatorA,'function_handle')
        fprintf('The main operator is detected as an M-function %s \n',...
            func2str(operatorA));
    elseif issparse(operatorA)
        fprintf('The main operator is detected as a sparse matrix \n');
    else
        fprintf('The main operator is detected as a full matrix \n');
    end
    if isempty(operatorB)
        fprintf('Solving standard eigenvalue problem, not generalized \n');
    elseif ischar(operatorB)
        fprintf(['The second operator of the generalized eigenproblem \n'...
            'is detected as an M-function %s \n'],operatorB);
    elseif isa(operatorB,'function_handle')
        fprintf(['The second operator of the generalized eigenproblem \n'...
            'is detected as an M-function %s \n'],func2str(operatorB));
    elseif issparse(operatorB)
        fprintf(strcat('The second operator of the generalized',...
            'eigenproblem \n is detected as a sparse matrix \n'));
    else
        fprintf(strcat('The second operator of the generalized',...
            'eigenproblem \n is detected as a full matrix \n'));
    end
    if isempty(operatorT)
        fprintf('No preconditioner is detected \n');
    elseif ischar(operatorT)
        fprintf('The preconditioner is detected as an M-function %s \n',...
            operatorT);
    elseif isa(operatorT,'function_handle')
        fprintf('The preconditioner is detected as an M-function %s \n',...
            func2str(operatorT));
    end
    if isempty(blockVectorY)
        fprintf('No matrix of constraints is detected \n')
    elseif issparse(blockVectorY)
        fprintf('The sparse matrix of %i constraints is detected \n',sizeY);
    else
        fprintf('The full matrix of %i constraints is detected \n',sizeY);
    end
    if issparse(blockVectorY) ~= issparse(blockVectorX)
        warning('BLOPEX:lobpcg:SparsityInconsistent',...
            strcat('The sparsity formats of the initial guess and ',...
            'the constraints are inconsistent'));
    end
end
% Set defaults
if isempty(residualTolerance)
    residualTolerance = sqrt(eps)*n;
end
if isempty(maxIterations)
    maxIterations = min(n,20);
end
if isempty(verbosityLevel)
    verbosityLevel = 0;
end
if verbosityLevel
    fprintf('Tolerance %e and maximum number of iterations %i \n',...
        residualTolerance,maxIterations)
end
%constraints preprocessing
if isempty(blockVectorY)
    constraintStyle = 0;
else
    %    constraintStyle = SYMMETRIC_CONSTRAINTS; % more accurate?
    constraintStyle = CONVENTIONAL_CONSTRAINTS;
end
if constraintStyle == CONVENTIONAL_CONSTRAINTS

    if isempty(operatorB)
        gramY = blockVectorY'*blockVectorY;
    else
        if isnumeric(operatorB)
            blockVectorBY = operatorB*blockVectorY;
        else
            blockVectorBY = feval(operatorB,blockVectorY);
        end
        gramY=blockVectorY'*blockVectorBY;
    end
    gramY=(gramY'+gramY)*0.5;
    if isempty(operatorB)
        blockVectorX = blockVectorX - ...
            blockVectorY*(gramY\(blockVectorY'*blockVectorX));
    else
        blockVectorX =blockVectorX - ...
            blockVectorY*(gramY\(blockVectorBY'*blockVectorX));
    end

elseif constraintStyle == SYMMETRIC_CONSTRAINTS

    if ~isempty(operatorB)
        if isnumeric(operatorB)
            blockVectorY = operatorB*blockVectorY;
        else
            blockVectorY = feval(operatorB,blockVectorY);
        end
    end
    if isempty(operatorT)
        gramY = blockVectorY'*blockVectorY;
    else
        blockVectorTY = feval(operatorT,blockVectorY);
        gramY = blockVectorY'*blockVectorTY;
    end
    gramY=(gramY'+gramY)*0.5;
    if isempty(operatorT)
        blockVectorX = blockVectorX - ...
            blockVectorY*(gramY\(blockVectorY'*blockVectorX));
    else
        blockVectorX = blockVectorX - ...
            blockVectorTY*(gramY\(blockVectorY'*blockVectorX));
    end

end
%Making the initial vectors (operatorB-) orthonormal
if isempty(operatorB)
    %[blockVectorX,gramXBX] = qr(blockVectorX,0);
    gramXBX=mygather(blockVectorX'*blockVectorX);
    if ~isreal(gramXBX)
        gramXBX=(gramXBX+gramXBX')*0.5;
    end
    [gramXBX,cholFlag]=chol(gramXBX);
    if  cholFlag ~= 0
        error('BLOPEX:lobpcg:ConstraintsTooTight',...
            'The initial approximation after constraints is not full rank');
    end
    blockVectorX = blockVectorX/gramXBX;
else
    %[blockVectorX,blockVectorBX] = orth(operatorB,blockVectorX);
    if isnumeric(operatorB)
        blockVectorBX = operatorB*blockVectorX;
    else
        blockVectorBX = feval(operatorB,blockVectorX);
    end
    gramXBX=blockVectorX'*blockVectorBX;
    if ~isreal(gramXBX)
        gramXBX=(gramXBX+gramXBX')*0.5;
    end
    [gramXBX,cholFlag]=chol(gramXBX);
    if  cholFlag ~= 0
        error('BLOPEX:lobpcg:InitialNotFullRank',...
            '%s\n%s', ...
            'The initial approximation after constraints is not ',...
            'full rank or/and operatorB is not positive definite');
    end
    blockVectorX = blockVectorX/gramXBX;
    blockVectorBX = blockVectorBX/gramXBX;
end
% Checking if the problem is big enough for the algorithm,
% i.e. n-sizeY > 5*blockSize
% Theoretically, the algorithm should be able to run if
% n-sizeY > 3*blockSize,
% but the extreme cases might be unstable, so we use 5 instead of 3 here.
if n-sizeY < 5*blockSize
    error('BLOPEX:lobpcg:MatrixTooSmall','%s\n%s', ...
        'The problem size is too small, relative to the block size.',...
        'Try using eig() or eigs() instead.');
end
% Preallocation
residualNormsHistory=zeros(blockSize,maxIterations);
lambdaHistory=zeros(blockSize,maxIterations+1);
condestGhistory=zeros(1,maxIterations+1);
blockVectorAR=zeros(n,blockSize, 'like', blockVectorX);
blockVectorP=zeros(n,blockSize, 'like', blockVectorX);
blockVectorAP=zeros(n,blockSize, 'like', blockVectorX);
if ~isempty(operatorB)
    blockVectorBR=zeros(n,blockSize, 'like', blockVectorX);
    blockVectorBP=zeros(n,blockSize, 'like', blockVectorX);
end
%Initial settings for the loop
if isnumeric(operatorA)
    blockVectorAX = operatorA*blockVectorX;
else
    blockVectorAX = feval(operatorA,blockVectorX);
end
gramXAX = full(blockVectorX'*blockVectorAX);
gramXAX = (gramXAX + gramXAX')*0.5;
% eig(...,'chol') uses only the diagonal and upper triangle -
% not true in MATLAB
% Octave v3.2.3-4, eig() does not support inputting 'chol'
[coordX,gramXAX]=eig(gramXAX,eye(blockSize));
lambda=diag(gramXAX); %eig returns non-ordered eigenvalues on the diagonal
if issparse(blockVectorX)
    coordX=sparse(coordX);
end
blockVectorX  =  blockVectorX*coordX;
blockVectorAX = blockVectorAX*coordX;
if ~isempty(operatorB)
    blockVectorBX = blockVectorBX*coordX;
end
clear coordX
condestGhistory(1)=-log10(eps)/2;  %if too small cause unnecessary restarts
lambdaHistory(1:blockSize,1) = mygather(lambda);
activeMask = true(blockSize,1);
% currentBlockSize = blockSize; %iterate all
%
% restart=1;%steepest descent
%The main part of the method is the loop of the CG method: begin
for iterationNumber=1:maxIterations

    %     %Computing the active residuals
    %     if isempty(operatorB)
    %%         if currentBlockSize > 1
    %%             blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %%                 blockVectorX(:,activeMask)*spdiags(lambda(activeMask),0,currentBlockSize,currentBlockSize);
    %%         else
    %%             blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %%                 blockVectorX(:,activeMask)*lambda(activeMask);
    %%         end
    %         blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %         bsxfun(@times,blockVectorX(:,activeMask),lambda(activeMask)');
    %     else
    %%         if currentBlockSize > 1
    %%             blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %%                 blockVectorBX(:,activeMask)*spdiags(lambda(activeMask),0,currentBlockSize,currentBlockSize);
    %%         else
    %%             blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %%                 blockVectorBX(:,activeMask)*lambda(activeMask);
    %%         end
    %         blockVectorR(:,activeMask)=blockVectorAX(:,activeMask) - ...
    %         bsxfun(@times,blockVectorBX(:,activeMask),lambda(activeMask)');
    %     end

    %Computing all residuals
    if isempty(operatorB)
        %         if blockSize > 1
        %             blockVectorR = blockVectorAX - ...
        %                 blockVectorX*spdiags(lambda,0,blockSize,blockSize);
        %         else
        %             blockVectorR = blockVectorAX - blockVectorX*lambda;
        %             %to make blockVectorR full when lambda is just a scalar
        %         end
        blockVectorR = blockVectorAX - ...
            bsxfun(@times,blockVectorX,lambda');
    else
        %        if blockSize > 1
        %             blockVectorR = blockVectorAX - ...
        %                 blockVectorBX*spdiags(lambda,0,blockSize,blockSize);
        %        else
        %             blockVectorR = blockVectorAX - blockVectorBX*lambda;
        %        end
        blockVectorR = blockVectorAX - ...
            bsxfun(@times,blockVectorBX,lambda');
    end

    %Satisfying the constraints for the active residulas
    if constraintStyle == SYMMETRIC_CONSTRAINTS
        if isempty(operatorT)
            blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
                blockVectorY*(gramY\(blockVectorY'*...
                blockVectorR(:,activeMask)));
        else
            blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
                blockVectorY*(gramY\(blockVectorTY'*...
                blockVectorR(:,activeMask)));
        end
    end

    residualNorms = full(sqrt(sum(conj(blockVectorR).*blockVectorR)'));
    residualNormsHistory(1:blockSize,iterationNumber) = ...
        mygather(residualNorms);

    %index antifreeze
    activeMask = full(residualNorms > residualTolerance) & activeMask;
    %activeMask = full(residualNorms > residualTolerance);
    %above allows vectors back into active, which causes problems with frosen Ps
    %activeMask = full(residualNorms > 0);      %iterate all, ignore freeze

    currentBlockSize = mygather(sum(activeMask));
    if  currentBlockSize == 0
        failureFlag=0; %all eigenpairs converged
        break
    end

    %Applying the preconditioner operatorT to the active residulas
    if ~isempty(operatorT)
        blockVectorR(:,activeMask) = ...
            feval(operatorT,blockVectorR(:,activeMask));
    end

    if constraintStyle == CONVENTIONAL_CONSTRAINTS
        if isempty(operatorB)
            blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
                blockVectorY*(gramY\(blockVectorY'*...
                blockVectorR(:,activeMask)));
        else
            blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
                blockVectorY*(gramY\(blockVectorBY'*...
                blockVectorR(:,activeMask)));
        end
    end

    %Making active (preconditioned) residuals orthogonal to blockVectorX
    if isempty(operatorB)
        blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
            blockVectorX*(blockVectorX'*blockVectorR(:,activeMask));
    else
        blockVectorR(:,activeMask) = blockVectorR(:,activeMask) - ...
            blockVectorX*(blockVectorBX'*blockVectorR(:,activeMask));
    end

    %Making active residuals orthonormal
    if isempty(operatorB)
        %[blockVectorR(:,activeMask),gramRBR]=...
        %qr(blockVectorR(:,activeMask),0); %to increase stability
        gramRBR=blockVectorR(:,activeMask)'*blockVectorR(:,activeMask);
        if ~isreal(gramRBR)
            gramRBR=(gramRBR+gramRBR')*0.5;
        end
        [gramRBR,cholFlag]=chol(gramRBR);
        if  cholFlag == 0
            blockVectorR(:,activeMask) = blockVectorR(:,activeMask)/gramRBR;
        else
            warning('BLOPEX:lobpcg:ResidualNotFullRank',...
                'The residual is not full rank.');
            break
        end
    else
        if isnumeric(operatorB)
            blockVectorBR(:,activeMask) = ...
                operatorB*blockVectorR(:,activeMask);
        else
            blockVectorBR(:,activeMask) = ...
                feval(operatorB,blockVectorR(:,activeMask));
        end
        gramRBR=blockVectorR(:,activeMask)'*blockVectorBR(:,activeMask);
        if ~isreal(gramRBR)
            gramRBR=(gramRBR+gramRBR')*0.5;
        end
        [gramRBR,cholFlag]=chol(gramRBR);
        if  cholFlag == 0
            blockVectorR(:,activeMask) = ...
                blockVectorR(:,activeMask)/gramRBR;
            blockVectorBR(:,activeMask) = ...
                blockVectorBR(:,activeMask)/gramRBR;
        else
            warning('BLOPEX:lobpcg:ResidualNotFullRankOrElse',...
                strcat('The residual is not full rank or/and operatorB ',...
                'is not positive definite.'));
            break
        end
    end
    clear gramRBR;

    if isnumeric(operatorA)
        blockVectorAR(:,activeMask) = ...
            mygather(operatorA*blockVectorR(:,activeMask));
    else
        blockVectorAR(:,activeMask) = ...
            feval(operatorA,blockVectorR(:,activeMask));
    end

    condestGmean = mean(condestGhistory(max(1,iterationNumber-10-...
        round(log(currentBlockSize))):iterationNumber));

    %  restart=1;

    % The Raileight-Ritz method for [blockVectorX blockVectorR blockVectorP]
    if isa(blockVectorAR, 'single')    % single initial
        myeps = 1; % play safe
    elseif isa(blockVectorR, 'single') %single somethings else
        myeps = eps(single(1));
    else                               % double everything
        myeps = eps;
    end
    if  mygather(residualNorms) > myeps^0.6
        explicitGramFlag = 0;
    else
        explicitGramFlag = 1;  %suggested by Garrett Moran, private
    end

    activeRSize=size(blockVectorR(:,activeMask),2);
    if iterationNumber == 1
        activePSize=0;
        restart=1;
    else
        activePSize=size(blockVectorP(:,activeMask),2);
        restart=0;
    end

    gramXAR=full(blockVectorAX'*blockVectorR(:,activeMask));
    gramRAR=full(blockVectorAR(:,activeMask)'*blockVectorR(:,activeMask));
    gramRAR=(gramRAR'+gramRAR)*0.5;

    if explicitGramFlag
        gramXAX=full(blockVectorAX'*blockVectorX);
        gramXAX=(gramXAX'+gramXAX)*0.5;
        if isempty(operatorB)
            gramXBX=full(blockVectorX'*blockVectorX);
            gramRBR=full(blockVectorR(:,activeMask)'*...
                blockVectorR(:,activeMask));
            gramXBR=full(blockVectorX'*blockVectorR(:,activeMask));
        else
            gramXBX=full(blockVectorBX'*blockVectorX);
            gramRBR=full(blockVectorBR(:,activeMask)'*...
                blockVectorR(:,activeMask));
            gramXBR=full(blockVectorBX'*blockVectorR(:,activeMask));
        end
        gramXBX=(gramXBX'+gramXBX)*0.5;
        gramRBR=(gramRBR'+gramRBR)*0.5;

    end

    if iterationNumber > 1
        %Making active conjugate directions orthonormal
        if isempty(operatorB)
            %[blockVectorP(:,activeMask),gramPBP] = qr(blockVectorP(:,activeMask),0);
            gramPBP=blockVectorP(:,activeMask)'*blockVectorP(:,activeMask);
            if ~isreal(gramPBP)
                gramPBP=(gramPBP+gramPBP')*0.5;
            end
            [gramPBP,cholFlag]=chol(gramPBP);
            if  cholFlag == 0
                blockVectorP(:,activeMask) = ...
                    blockVectorP(:,activeMask)/gramPBP;
                blockVectorAP(:,activeMask) = ...
                    blockVectorAP(:,activeMask)/gramPBP;
                restart = 0;
            else
                warning('BLOPEX:lobpcg:DirectionNotFullRank',...
                    'The direction matrix is not full rank.');
                restart = 1;
            end
        else
            gramPBP=blockVectorP(:,activeMask)'*blockVectorBP(:,activeMask);
            if ~isreal(gramPBP)
                gramPBP=(gramPBP+gramPBP')*0.5;
            end
            [gramPBP,cholFlag]=chol(gramPBP);
            if  cholFlag == 0
                blockVectorP(:,activeMask) = ...
                    blockVectorP(:,activeMask)/gramPBP;
                blockVectorAP(:,activeMask) = ...
                    blockVectorAP(:,activeMask)/gramPBP;
                blockVectorBP(:,activeMask) = ...
                    blockVectorBP(:,activeMask)/gramPBP;
                restart = 0;
            else
                warning('BLOPEX:lobpcg:DirectionNotFullRank',...
                    strcat('The direction matrix is not full rank ',...
                    'or/and operatorB is not positive definite.'));
                restart = 1;
            end
        end
        clear gramPBP
    end

    for cond_try=1:2           %cond_try == 2 when restart

        if ~restart
            gramXAP=full(blockVectorAX'*blockVectorP(:,activeMask));
            gramRAP=full(blockVectorAR(:,activeMask)'*...
                blockVectorP(:,activeMask));
            gramPAP=full(blockVectorAP(:,activeMask)'*...
                blockVectorP(:,activeMask));
            gramPAP=(gramPAP'+gramPAP)*0.5;

            if explicitGramFlag
                gramA = [ gramXAX     gramXAR     gramXAP
                    gramXAR'    gramRAR     gramRAP
                    gramXAP'     gramRAP'    gramPAP ];
            else
                gramA = [ diag(lambda)  gramXAR  gramXAP
                    gramXAR'      gramRAR  gramRAP
                    gramXAP'      gramRAP'  gramPAP ];
            end

            clear gramXAP  gramRAP gramPAP

            if isempty(operatorB)
                gramXBP=full(blockVectorX'*blockVectorP(:,activeMask));
                gramRBP=full(blockVectorR(:,activeMask)'*...
                    blockVectorP(:,activeMask));
            else
                gramXBP=full(blockVectorBX'*blockVectorP(:,activeMask));
                gramRBP=full(blockVectorBR(:,activeMask)'*...
                    blockVectorP(:,activeMask));
                %or blockVectorR(:,activeMask)'*blockVectorBP(:,activeMask);
            end

            if explicitGramFlag
                if isempty(operatorB)
                    gramPBP=full(blockVectorP(:,activeMask)'*...
                        blockVectorP(:,activeMask));
                else
                    gramPBP=full(blockVectorBP(:,activeMask)'*...
                        blockVectorP(:,activeMask));
                end
                gramPBP=(gramPBP'+gramPBP)*0.5;
                gramB = [ gramXBX  gramXBR  gramXBP
                    gramXBR' gramRBR  gramRBP
                    gramXBP' gramRBP' gramPBP ];
                clear   gramPBP
            else
                gramB=[eye(blockSize) zeros(blockSize,activeRSize) gramXBP
                    zeros(blockSize,activeRSize)' eye(activeRSize) gramRBP
                    gramXBP' gramRBP' eye(activePSize) ];
            end

            clear gramXBP  gramRBP;

        else

            if explicitGramFlag
                gramA = [ gramXAX   gramXAR
                    gramXAR'    gramRAR  ];
                gramB = [ gramXBX  gramXBR
                    gramXBR' eye(activeRSize)  ];
                clear gramXAX gramXBX gramXBR
            else
                gramA = [ diag(lambda)  gramXAR
                    gramXAR'        gramRAR  ];
                gramB = eye(blockSize+activeRSize);
            end

            clear gramXAR gramRAR;

        end

        condestG = log10(cond(gramB))+1;
        if (condestG/condestGmean > 2 && condestG > 2 )|| condestG > 8
            %black magic - need to guess the restart
            if verbosityLevel
                fprintf('Restart on step %i as condestG %5.4e \n',...
                    iterationNumber,condestG);
            end
            if cond_try == 1 && ~restart
                restart=1; %steepest descent restart for stability
            else
                warning('BLOPEX:lobpcg:IllConditioning',...
                    'Gramm matrix ill-conditioned: results unpredictable');
            end
        else
            break
        end

    end

    [gramA,gramB]=eig(gramA,gramB);
    lambda=diag(gramB(1:blockSize,1:blockSize));
    coordX=gramA(:,1:blockSize);

    clear gramA gramB

    if issparse(blockVectorX)
        coordX=sparse(coordX);
    end

    if ~restart
        blockVectorP =  blockVectorR(:,activeMask)*...
            coordX(blockSize+1:blockSize+activeRSize,:) + ...
            blockVectorP(:,activeMask)*...
            coordX(blockSize+activeRSize+1:blockSize + ...
            activeRSize+activePSize,:);
        blockVectorAP = blockVectorAR(:,activeMask)*...
            coordX(blockSize+1:blockSize+activeRSize,:) + ...
            blockVectorAP(:,activeMask)*...
            coordX(blockSize+activeRSize+1:blockSize + ...
            activeRSize+activePSize,:);
        if ~isempty(operatorB)
            blockVectorBP = blockVectorBR(:,activeMask)*...
                coordX(blockSize+1:blockSize+activeRSize,:) + ...
                blockVectorBP(:,activeMask)*...
                coordX(blockSize+activeRSize+1:blockSize+activeRSize+activePSize,:);
        end
    else %use block steepest descent
        blockVectorP =   blockVectorR(:,activeMask)*...
            coordX(blockSize+1:blockSize+activeRSize,:);
        blockVectorAP = blockVectorAR(:,activeMask)*...
            coordX(blockSize+1:blockSize+activeRSize,:);
        if ~isempty(operatorB)
            blockVectorBP = blockVectorBR(:,activeMask)*...
                coordX(blockSize+1:blockSize+activeRSize,:);
        end
    end

    blockVectorX = blockVectorX*coordX(1:blockSize,:) + blockVectorP;
    blockVectorAX = blockVectorAX*coordX(1:blockSize,:) + blockVectorAP;
    if ~isempty(operatorB)
        blockVectorBX = blockVectorBX*coordX(1:blockSize,:) + blockVectorBP;
    end
    clear coordX
    %%end RR

    lambdaHistory(1:blockSize,iterationNumber+1) = mygather(lambda);
    condestGhistory(iterationNumber+1) = mygather(condestG);

    if verbosityLevel
        fprintf('Iteration %i current block size %i \n',...
            iterationNumber,currentBlockSize);
        fprintf('Eigenvalues lambda %17.16e \n',mygather(lambda));
        fprintf('Residual Norms %e \n',mygather(residualNorms'));
    end
end
%The main step of the method was the CG cycle: end
%Postprocessing
%Making sure blockVectorX's "exactly" satisfy the blockVectorY constrains??
%Making sure blockVectorX's are "exactly" othonormalized by final "exact" RR
if isempty(operatorB)
    gramXBX=full(blockVectorX'*blockVectorX);
else
    if isnumeric(operatorB)
        blockVectorBX = operatorB*blockVectorX;
    else
        blockVectorBX = feval(operatorB,blockVectorX);
    end
    gramXBX = full(blockVectorX'*blockVectorBX);
end
gramXBX=(gramXBX'+gramXBX)*0.5;
if isnumeric(operatorA)
    blockVectorAX = operatorA*blockVectorX;
else
    blockVectorAX = feval(operatorA,blockVectorX);
end
gramXAX = full(blockVectorX'*blockVectorAX);
gramXAX = (gramXAX + gramXAX')*0.5;
%Raileigh-Ritz for blockVectorX, which is already operatorB-orthonormal
[coordX,gramXBX] = eig(gramXAX,gramXBX);
lambda=diag(gramXBX);
if issparse(blockVectorX)
    coordX=sparse(coordX);
end
blockVectorX  =   blockVectorX*coordX;
blockVectorAX  =  blockVectorAX*coordX;
if ~isempty(operatorB)
    blockVectorBX  =  blockVectorBX*coordX;
end
%Computing all residuals
if isempty(operatorB)
    %     if blockSize > 1
    %         blockVectorR = blockVectorAX - ...
    %             blockVectorX*spdiags(lambda,0,blockSize,blockSize);
    %     else
    %         blockVectorR = blockVectorAX - blockVectorX*lambda;
    %     end
    blockVectorR = blockVectorAX - ...
        bsxfun(@times,blockVectorX,lambda');
else
    %     if blockSize > 1
    %         blockVectorR=blockVectorAX - ...
    %             blockVectorBX*spdiags(lambda,0,blockSize,blockSize);
    %     else
    %         blockVectorR = blockVectorAX - blockVectorBX*lambda;
    %     end
    blockVectorR = blockVectorAX - ...
        bsxfun(@times,blockVectorBX,lambda');
end
residualNorms=full(sqrt(sum(conj(blockVectorR).*blockVectorR)'));
residualNormsHistory(1:blockSize,iterationNumber) = ...
    mygather(residualNorms);
if verbosityLevel
    fprintf('Final Eigenvalues lambda %17.16e \n',mygather(lambda));
    fprintf('Final Residual Norms %e \n',mygather(residualNorms'));
end
if verbosityLevel == 2
    whos
    figure(491)
    semilogy((abs(residualNormsHistory(1:blockSize,1:iterationNumber-1)))');
    title('Residuals for Different Eigenpairs','fontsize',16);
    ylabel('Eucledian norm of residuals','fontsize',16);
    xlabel('Iteration number','fontsize',16);
    %axis tight;
    %axis([0 maxIterations+1 1e-15 1e3])
    set(gca,'FontSize',14);
    figure(492);
    semilogy(abs((lambdaHistory(1:blockSize,1:iterationNumber)-...
        repmat(mygather(lambda),1,iterationNumber)))');
    title('Eigenvalue errors for Different Eigenpairs','fontsize',16);
    ylabel('Estimated eigenvalue errors','fontsize',16);
    xlabel('Iteration number','fontsize',16);
    %axis tight;
    %axis([0 maxIterations+1 1e-15 1e3])
    set(gca,'FontSize',14);
    drawnow;
end
varargout(1)={failureFlag};
varargout(2)={lambdaHistory(1:blockSize,1:iterationNumber)};
varargout(3)={residualNormsHistory(1:blockSize,1:iterationNumber-1)};
end


function ierror = part_rsf_check ( n, npart, a )

%*****************************************************************************80
%
%% PART_RSF_CHECK checks a reverse standard form partition of an integer.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    25 January 2011
%
%  Author:
%
%    John Burkardt
%
%  Reference:
%
%    Donald Kreher, Douglas Simpson,
%    Combinatorial Algorithms,
%    CRC Press, 1998,
%    ISBN: 0-8493-3988-X,
%    LC: QA164.K73.
%
%  Parameters:
%
%    Input, integer N, the integer to be partitioned.
%    N must be positive.
%
%    Input, integer NPART, the number of parts of the partition.
%    1 <= NPART <= N.
%
%    Input, integer A(NPART), contains the partition.
%    A(1) through A(NPART) contain the nonzero integers which
%    sum to N.  The entries must be in ASCENDING order.
%
%    Output, integer IERROR, error flag.
%    0, no error.
%    -1, N is illegal.
%    -2, NPART is illegal.
%    -3, the entries do not add up to N.
%    I, the I-th entry of A is illegal.
%
  ierror = 0;

  if ( n < 1 )
    ierror = -1;
    fprintf ( 1, '\n' );
    fprintf ( 1, 'PART_RSF_CHECK - Fatal error!\n' );
    fprintf ( 1, '  N < 1.\n' );
    error ( 'PART_RSF_CHECK - Fatal error!\n' );
  end

  if ( npart < 1 || n < npart )
    ierror = -2;
    return
  end
%
%  Every entry must lie between 1 and N.
%
  for i = 1 : npart
    if ( a(i) < 1 || n < a(i) )
      ierror = i;
      return
    end
  end
%
%  The entries must be in ascending order.
%
  for i = 2 : npart
    if ( a(i) < a(i-1) )
      ierror = i;
      return
    end
  end
%
%  The entries must add up to N.
%
  if ( sum ( a(1:npart) ) ~= n )
    ierror = -3;
  end

  return
end


function segment_length = p05_boundary_segment_length ( segment_index, h )

%*****************************************************************************80
%
%% P05_BOUNDARY_SEGMENT_LENGTH returns boundary segment lengths in problem 05.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    12 April 2009
%
%  Author:
%
%    John Burkardt
%
%  Reference:
%
%    Per-Olof Persson and Gilbert Strang,
%    A Simple Mesh Generator in MATLAB,
%    SIAM Review,
%    Volume 46, Number 2, June 2004, pages 329-345.
%
%  Parameters:
%
%    Input, integer SEGMENT_INDEX, the index of one of the boundary segments.
%
%    Input, real H, the suggested spacing between points.
%
%    Output, integer SEGMENT_LENGTH, the number of points in the segment.
%
  center1 = [  0.0, 0.0 ];
  center2 = [ -0.4, 0.0 ];
  r1 = 1.00;
  r2 = 0.55;

  if ( h <= 0.0 )
    fprintf ( 1, '\n' );
    fprintf ( 1, 'P05_BOUNDARY_SEGMENT_LENGTH - Fatal error!\n' );
    fprintf ( 1, '  Nonpositive H = %f\n', h );
    error ( 'P05_BOUNDARY_SEGMENT_LENGTH - Fatal error!' );
  end

  if ( segment_index == 1 )

    n = round ( ( pi * ( r1  + r2 ) ...
      + ( ( center2(1) - r2 ) - ( center1(1) - r1 ) ) ...
      + ( ( center1(1) + r1 ) - ( center2(1) + r2 ) ) ) / h );
    n = max ( n, 21 );
    segment_length = n;

  else

    fprintf ( 1, '\n' );
    fprintf ( 1, 'P05_BOUNDARY_SEGMENT_LENGTH - Fatal error!\n' );
    fprintf ( 1, '  Illegal SEGMENT_INDEX = %d\n', segment_index );
    error ( 'P05_BOUNDARY_SEGMENT_LENGTH - Fatal error!' );

  end

  return
end


function [rglm,yclean] = spm_rglm (y,X,m,priors,verbose)
% Fit a Robust GLM
% FORMAT [rglm,yclean] = spm_rglm (y,X,m,priors,verbose)
%
% The noise is modelled with a Mixture of Zero-Mean Gaussians
%
% y          [N x 1] data vector
% X          [N x p] design matrix
% m          Number of mixture components
% priors     .alpha      [1 x 1] weight precision (default=0.001)
%            .mean_err   [m x 1] vector of mean error SD
%            .std_err    [m x 1] vector of dev of error SD
% verbose    0/1 to printout inner workings (default=0)
%
% rglm       Returned model
% yclean     'Clean' data
%
% -------------------------------------------------------
% The fields in rglm are:
%
% m                The number of error components
% fm               The negative free energy
% loops            Number of iterations used
%
%                  In the field priors:
%
% lambda_0         Dirichlet parameters for mixing coeffs
% b_0,c_0          Gamma parameters for precisions
%
%                  In the field posts:
%
% lambda           Dirichlet parameters  for mixing coeffs
% b,c              Gamma parameters for precisions
% w_mean           Mean estimated regression coefficients
% w_cov            Covariance of regression coefficients
% pi               mixing coefficients (lambda/sum(lambda))
% variances        variances (1./(b.*c))
% gamma            the responsilities of each noise component
%___________________________________________________________________________
% Copyright (C) 2008 Wellcome Trust Centre for Neuroimaging

% Will Penny
% $Id: spm_rglm.m 1276 2008-03-28 18:29:19Z guillaume $

if nargin < 4 | isempty(priors)
    mean_alpha=0.001;
    % Set variance priors
    b_0=1000*ones(m,1);
    c_0=0.001*ones(m,1);
else
    mean_alpha=priors.alpha;
    for mm=1:m,
        b_0(mm)=(priors.std_err(mm)^2)/priors.mean_err(mm);
        c_0(mm)=priors.mean_err(mm)/b_0(mm);
    end
end

if (m==1)
    rglm=spm_glm(y,X,mean_alpha);
    return
end

if nargin < 5 | isempty(verbose)
    verbose=0;
end

y=y(:);
N=length(y);
p=size(X,2);
spX=issparse(X);

% Initialise regression coefficients to maximum likelihood solution
if spX
    iX=inv(X'*X);
    w_mean=iX*X'*y;
else
    iX=inv(X'*X);
    w_mean = pinv(X)*y;
end
y_pred = full(X*w_mean);
err=y-y_pred;
v=mean((y-y_pred).^2);
w_cov = v*iX;

% Set mixing priors
lambda_0=5;

% Cluster on absolute difference from mean
zmix=spm_kmeans1(abs(err-mean(err)),m);
% Posterior for mixers
lambda=100*zmix.pi+lambda_0;
zmean=[zmix.m].^2;
% Posterior for precisions
var_precision=1/std(zmean)^2;
for s=1:m,
      % Set so that b*c=precision and  b^2*c=var_precision
      precision=1/zmean(s);
      b(s)=var_precision/precision;
      c(s)=(precision^2)/var_precision;
end

if verbose
    disp('Init');
    for s=1:m,
        disp(sprintf('State %d mix=%1.2f var=%1.2f',s,lambda(s)/sum(lambda),1/(b(s)*c(s))));
    end
end

lik=[];
tol=0.0001;
max_loops=100;
WLOOPS=5;
for loops=1:max_loops,

    % E-step
    lambda_tot=sum(lambda);
    ypred=full(X*w_mean);
    ypred2=ypred.^2;
    y2=y.^2;
    xt=X';
    y_err=full(sum(xt.*(w_cov*xt)));

    tv=y2-2*ypred.*y+y_err'+ypred2;
    tv=tv';
    for s=1:m,
        log_tilde_pi(s)=psi(lambda(s))-psi(lambda_tot);
        log_tilde_beta(s)=psi(c(s))+log(b(s));
        tilde_pi(s)=exp(log_tilde_pi(s));
        tilde_beta(s)=exp(log_tilde_beta(s));
        mean_beta(s)=c(s)*b(s);
        tilde_var(s,:)=tv;
        gamma(s,:)=tilde_pi(s)*(tilde_beta(s)^0.5)*exp(-0.5*mean_beta(s)*tv);
    end
    gamma_n=sum(gamma);
    for s=1:m,
        if mean(gamma_n) > eps
            % If component still exists
            gamma(s,:)=gamma(s,:)./gamma_n;
        end
    end

    % M-step
    % Part I
    for s=1:m,
        pi_bar(s)=mean(gamma(s,:));
        N_bar(s)=N*pi_bar(s);
        mean_bar(s)=mean(gamma(s,:)'.*y);
        dg=sparse(1:N,1:N,gamma(s,:));
        x_bar(s,:)=mean(dg*X);
        var_bar_mu(s)=mean(gamma(s,:).*tilde_var(s,:));
    end

    % CALCULATE THE FREE ENERGY
    avg_likelihood=sum(N_bar.*(log_tilde_pi+0.5*log_tilde_beta));
    fit=-0.5*N*sum(mean_beta.*var_bar_mu);
    % Ensure that 0 log 0 = 0
    ent_s=sum(sum(-gamma.*log(gamma+eps)));
    avg_likelihood=avg_likelihood+fit+ent_s;
    lambda_p=lambda_0*ones(1,m);
    kl_dir=spm_kl_dirichlet(lambda,lambda_p,log_tilde_pi);
    kl_gamm=0;
    for s=1:m,
        kl_gamm=kl_gamm+spm_kl_gamma(b(s),c(s),b_0(s),c_0(s));
    end
    kl_weights=spm_kl_normald(w_mean,w_cov,zeros(1,p),(1/mean_alpha)*eye(p));
    fm= avg_likelihood - kl_dir - kl_gamm - kl_weights;

    % Convergence criterion
    oldlik=lik;
    lik=fm;
    if (mod(loops-1,WLOOPS)==0)
        if (loops>1)
            if abs((lik-oldlik)/lik) < tol
                break;
            end
        end
    end

    % M-Step: Part II
    for s=1:m,
        % Mixers
        lambda(s)=N_bar(s)+lambda_0;
        % Precisions
        b(s)=1/(0.5*N*var_bar_mu(s)+1/b_0(s));
        c(s)=0.5*N_bar(s)+c_0(s);
        mean_beta(s)=c(s)*b(s);
    end

    if mod(loops,WLOOPS)==0
        % Regression coefficients
        cc=sparse(p,p);
        cw=zeros(p,1);
        for s=1:m,
            dg=sparse(1:N,1:N,gamma(s,:));
            cc=cc+mean_beta(s)*X'*dg*X;
            cw=cw+mean_beta(s)*X'*dg*y;
        end
        cc=cc+mean_alpha*speye(p);
        w_cov=inv(cc);
        w_mean=w_cov*cw;
    end

    if verbose
        disp(sprintf('It=%d, L_AV =%1.2f, KL Mix=%1.2f, KL Prec=%1.2f, KL-W=%1.2f, Fm=%1.2f',loops,avg_likelihood,kl_dir,kl_gamm,kl_weights,fm));
        disp(' ');
        disp(sprintf('Iteration number=%d',loops));
        for s=1:m,
            var=1/(b(s)*c(s));
            disp(sprintf('State %d pi_bar=%1.2f var=%1.2f',s,pi_bar(s),var));
        end
    end

end

% Put variables into data structure
rglm.m=m;
rglm.fm=fm;
rglm.kl_dir=kl_dir;
rglm.kl_gamm=kl_gamm;
rglm.kl_w=kl_weights;
rglm.mean_alpha=mean_alpha;
rglm.priors.lambda_0=lambda_0;
rglm.priors.c_0=c_0;
rglm.priors.b_0=b_0;
for k=1:m,
  rglm.posts.lambda(k)=lambda(k);
end
for k=1:m,
  rglm.posts.c(k)=c(k);
  rglm.posts.b(k)=b(k);
end
rglm.posts.w_mean=w_mean;
rglm.posts.w_cov=w_cov;

for k=1:m,
  rglm.posts.pi(k)=lambda(k)/sum(lambda);
end
rglm.posts.variances=1./(b.*c);
rglm.posts.gamma=gamma;
rglm.loops=loops;

% Get 'clean' data
if m > 1
    [tmp,outlier_class]=min(rglm.posts.pi);
    e=y-ypred;
    outlier_error=rglm.posts.gamma(outlier_class,:)'.*e;
    yclean=ypred+e-outlier_error;
else
    yclean=y;
end


function [y,ysup] = erf_rnd(x,rnd)
% input x real non-negative column vector
% rnd  -1  y = lower bound for erf(x)
%       1  y = upper bound for erf(x)
%      []  [y,ysup] inclusion of erf(x)
% rounding may be altered after leaving erf_rnd
%

% written  05/30/13     S.M. Rump
%

  x2 = hex2num('4017744f8f74e94b');     % ~5.86: erf(x)>pred(1) for x>=x2

  y = x;
  if isempty(rnd)
    ysup = x;
  end

  index = ( x<=0.5 );                   % first method
  if any(index)                         % x in [0,0.5]
    if isempty(rnd)
      [y(index),ysup(index)] = erf_rnd1(x(index),rnd,6);
    else
      y(index) = erf_rnd1(x(index),rnd,6);
    end
  end
  Index = index;                        % store finished indices

  index = ( ~Index ) & ( x<x2 );        % second method
  if any(index)                         % x in (0.5,x2)
    y_index = 1 - erfc_rnd2(x(index));
    if rnd==-1
      y(index) = y_index.inf;
    elseif rnd==1
      y(index) = y_index.sup;
    else
      y(index) = y_index.inf;
      ysup(index) = y_index.sup;
    end
  end

  index = ( x>=x2 );
  if any(index)                         % x in [x2,inf]
    if isempty(rnd)                     % inclusion [y,ysup] of erf(x)
      y(index) = hex2num('3fefffffffffffff');   % pred(1)
      ysup(index) = 1;
    elseif rnd==1                       % y upper bound for erf(x)
      y(index) = 1;
    else                                % y lower bound for erf(x)
      y(index) = hex2num('3fefffffffffffff');   % pred(1)
    end
  end


function T = RpToTrans(R, p)
% *** CHAPTER 3: RIGID-BODY MOTIONS ***
% Takes rotation matrix R and position p.
% Returns the corresponding homogeneous transformation matrix T in SE(3).
% Example Input:
%
% clear; clc;
% R = [[1, 0, 0]; [0, 0, -1]; [0, 1, 0]];
% p = [1; 2; 5];
% T = RpToTrans(R, p)
%
% Output:
% T =
%     1     0     0     1
%     0     0    -1     2
%     0     1     0     5
%     0     0     0     1

T = [R, p; 0, 0, 0, 1];
end


function [imgOut, exposure_value] = SelectOverexposedTMO(img, percent)
%
%        [imgOut, exposure_value] = SelectOverexposedTMO(img, percent)
%
%
%        A simple TMO that generates an SDR image with a given percentage
%        of overexposed pixels.
%
%        Input:
%           -img: input HDR image
%           -percent: percentage of overexposed pixels.
%
%        Output:
%           -imgOut: a single exposure with the selected percentage of
%           overexposed pixels.
%
%
%
%     Copyright (C) 2022 Francesco Banterle
%
%     This program is free software: you can redistribute it and/or modify
%     it under the terms of the GNU General Public License as published by
%     the Free Software Foundation, either version 3 of the License, or
%     (at your option) any later version.
%
%     This program is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
%
%     You should have received a copy of the GNU General Public License
%     along with this program.  If not, see <http://www.gnu.org/licenses/>.
%

if ~exist('percent', 'var')
    percent = 0.1;
end

gamma_inv = 1.0 / 2.2;

    function err = residual(c)

        if c > 0.0
            tmp = ClampImg((img * c).^gamma_inv, 0.0, 1.0);
            L = lum(tmp);
            tmp_percent = length(find(L >= 0.95)) / numel(L);

            err = abs(tmp_percent - percent);
        else
            err = 1e9;
        end
    end

exposure_start = 1.0;
exposure_value = fminsearch(@residual, exposure_start, getOpts);

imgOut = ClampImg(img * exposure_value, 0.0, 1.0);

end