function [ci,bootstat]  = stat_bootci(nboot,bootfun,varargin)
%BOOTCI Bootstrap Confidence Interval
%   CI = BOOTCI(NBOOT,BOOTFUN,...) computes the 95 percent BCa bootstrap
%   confidence interval of the statistic defined by the function BOOTFUN.
%   NBOOT is a positive integer indicating the number of bootstrap data
%   samples used in the computation. BOOTFUN is a function handle specified
%   with @. The third and later input arguments to BOOTCI are data
%   (scalars, column vectors, or matrices) that are used to create inputs
%   to BOOTFUN. BOOTCI creates each bootstrap sample by sampling with
%   replacement from the rows of the non-scalar data arguments (these must
%   have the same number of rows). Scalar data are passed to BOOTFUN
%   unchanged.
%
%   If BOOTFUN returns a scalar, CI is a vector containing the lower and
%   upper bounds of the confidence interval. If BOOTFUN returns a vector of
%   length M, CI is an array of size 2-by-M, where CI(1,:) are lower bounds
%   and CI(2,:) are upper bounds. If BOOTFUN returns an array of size
%   M-by-N-by-P-by-..., CI is an array of size 2-by-M-by-N-by-P-by-...,
%   where CI(1,:,:,:,...) is an array of lower bounds and CI(2,:,:,:,...)
%   is an array of upper bounds.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},'alpha',ALPHA) computes the 100*(1-ALPHA)
%   percent BCa bootstrap confidence interval of the statistic defined by the
%   function BOOTFUN. ALPHA is a scalar between 0 and 1. The default value of
%   ALPHA is 0.05. A cell array groups BOOTFUN and the arguments used to create
%   inputs to it. ALPHA and any other arguments to BOOTCI appear outside the
%   cell array.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},...,'type',TYPE) computes the bootstrap
%   confidence interval of the statistic defined by the function BOOTFUN.
%   TYPE is the confidence interval type, specifying different methods of
%   computing the confidence interval. TYPE is a string chosen from
%       'norm' or 'normal':               normal approximated interval with
%                                         bootstrapped bias and standard
%                                         error;
%       'per' or 'percentile':            basic percentile method;
%       'cper' or 'corrected percentile': bias corrected percentile method;
%       'bca' :                           bias corrected and accelerated
%                                         percentile method;
%       'stud' or 'student':              studentized confidence interval.
%   The default value of TYPE is 'bca'.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},...,'type','stud','nbootstd',NBOOTSTD)
%   computes the studentized bootstrap confidence interval of the statistic
%   defined by the function BOOTFUN. The standard error of the bootstrap
%   statistics is estimated using bootstrap with NBOOTSTD bootstrap data
%   samples. NBOOTSTD is a positive integer value. The default value of
%   NBOOTSTD is 100.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},...,'type','stud','stderr',STDERR)
%   computes the studentized bootstrap confidence interval of statistics
%   defined by the function BOOTFUN. The standard error of the bootstrap
%   statistics is evaluated by the function STDERR. STDERR is a function
%   handle created using @. STDERR should take the same arguments as
%   BOOTFUN and return the standard error of the statistic computed by
%   BOOTFUN.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},...,'Weights',WEIGHTS) specifies
%   observation weights. WEIGHTS must be a vector of non-negative numbers
%   with at least one positive element. The number of elements in WEIGHTS
%   must be equal to the number of rows in non-scalar input arguments to
%   BOOTFUN. To obtain one bootstrap replicate, BOOTSTRP samples N out of N
%   with replacement using these weights as multinomial sampling
%   probabilities.
%
%   CI = BOOTCI(NBOOT,{BOOTFUN,...},...,'Options',OPTIONS) contains options
%   that specify whether to compute bootstrap iterations in parallel, and how
%   to use random numbers during the bootstrap sampling. OPTIONS is a struct
%   that can be created by a call to STATSET. BOOTCI uses the following fields:
%       'UseParallel'
%       'UseSubstreams'
%       'Streams'
%   For information on these fields see PARALLELSTATS.
%   NOTE: if 'UseParallel' is 'always' and 'UseSubstreams'
%   is 'never', then the length of Streams must equal the number
%   of processors used by BOOTCI. There are two possibilities.
%   If a MATLAB pool is open, then Streams is the same length as
%   the size of the MATLAB pool. If a MATLAB pool is not open,
%   then Streams must supply a single random number stream.
%
%   [CI,BOOTSTAT] = BOOTCI(...) also returns the bootstrapped statistic
%   computed for each of the NBOOT bootstrap replicate samples.  Each row of
%   BOOTSTAT contains the results of applying BOOTFUN to one bootstrap sample.
%   If BOOTFUN returns a matrix or array, then this output is converted to a
%   row vector for storage in BOOTSTAT.
%
%   Example:
%     Compute the confidence interval for the capability index in
%     statistical process control:
%          y = normrnd(1,1,30,1);                  % simulated process data
%          LSL = -3;  USL = 3;                     % process specifications
%          capable = @(x) (USL-LSL)./(6* std(x));  % process capability
%          bootci(2000,capable, y)                 % Bca confidence interval
%          bootci(2000,{capable, y},'type','per')  % basic percentile method
%
%   See also: BOOTSTRP, JACKKNIFE, STATSET, STATGET,
%   RANDSAMPLE, PARFOR, PARALLELSTATS.

% The BCa method is described in the following references:
%
%     T.J. DiCicio and B. Efron (1996), "Bootstrap confidence intervals,"
%     Statistical Science, v. 11, n. 3, pp. 189-228.
%
%     B. Efron and R.J. Tibshirani (1993), An Introduction to the
%     Bootstrap, Chapman & Hall, New York.
%
% Their formula involves a Z0 factor that is computed using the proportion
% of bootstrap values less than the original sample value.  In order to get
% reasonable results when the sample is lumpy, we include half of the
% bootstrap values that are tied with the original sample value when we
% compute Z0.

% Copyright 2005-2010 The MathWorks, Inc.
%   $Revision: 1.1.8.8 $  $Date: 2011/07/20 00:08:12 $

if nargin<2
    error(message('stats:bootci:TooFewInputs'));
end;
if nboot<=0 || nboot~=round(nboot)
    error(message('stats:bootci:BadNboot'))
end;

if ~iscell(bootfun) % default syntax
    type = 'bca';
    alpha = .05;
    fun = bootfun;
    data = varargin;
    weights = [];
    bootstrpOptions = statset('bootstrp');
else % syntax with optional type, alpha, nbootstd, and stderrfun name/value pairs
    fun = bootfun{1};
    data = bootfun(2:end);
    pnames = {'type', 'alpha', 'stderr', 'nbootstd', 'weights', 'options'};
    dflts =  {'bca', .05, [], 100, [], statset('bootstrp')};
    [type,alpha,stderrfun,nbootstd,weights,bootstrpOptions] = ...
                           internal.stats.parseArgs(pnames, dflts, varargin{:});
end

% error check for the bootfun
try
    obsstat = fun(data{:});
catch ME
    m = message('stats:bootci:BadBootfun',func2str(fun));
    throw(addCause(MException(m.Identifier,'%s',getString(m)),ME));
end
if any(~isfinite(obsstat))
        error(message('stats:bootci:NonfiniteBootfun'));
end;

% Save the original size of stat and reshape it
sz = [];
if ~isvector(obsstat)
    sz = size(obsstat);
end
obsstat = obsstat(:)'; % turn into a row-vector

% call subfunctions to compute the intervals
switch (lower(type))
    case {'norm','normal'}
        [ci,bootstat] = bootnorm(obsstat,nboot,fun,alpha,weights,bootstrpOptions,data{:});
    case {'per','percentile'}
        [ci,bootstat] = bootper(nboot,fun,alpha,weights,bootstrpOptions,data{:});
    case {'cper', 'corrected percentile'}
        [ci,bootstat] = bootcper(obsstat,nboot,fun,alpha,weights,bootstrpOptions,data{:});
    case 'bca'
        [ci,bootstat] = bootbca(obsstat,nboot,fun,alpha,weights,bootstrpOptions,data{:});
    case {'stud','student'}
        [ci,bootstat] = bootstud(obsstat,nboot,fun,alpha,nbootstd,stderrfun,...
            weights,bootstrpOptions,data{:});
    otherwise
        error(message('stats:bootci:BadType'))
end;

% Reshape
if ~isempty(sz)
    ci = reshape(ci,[2 sz]);
end

end   % bootci()

%-------------------------------------------------------------------------
function [ci,bstat] = bootnorm(stat,nboot,bootfun,alpha,weights,bootstrpOptions,varargin)
% normal approximation interval
% A.C. Davison and D.V. Hinkley (1996), p198-200

bstat = bootstrp(nboot,bootfun,varargin{:},'weights',weights,'Options',bootstrpOptions);

se = std(bstat,0,1);   % standard deviation estimate
bias = mean(bsxfun(@minus,bstat,stat),1);
za = norminv(alpha/2);   % normal confidence point
lower = stat - bias + se*za; % lower bound
upper = stat - bias - se*za;  % upper bound

% return
ci = [lower;upper];
end   % bootnorm()

%-------------------------------------------------------------------------
function [ci,bstat] = bootper(nboot,bootfun,alpha,weights,bootstrpOptions,varargin)
% percentile bootstrap CI

bstat = bootstrp(nboot,bootfun,varargin{:},'weights',weights,'Options',bootstrpOptions);

pct1 = 100*alpha/2;
pct2 = 100-pct1;
lower = prctile(bstat,pct1,1);
upper = prctile(bstat,pct2,1);

% return
ci =[lower;upper];
end % bootper()

%-------------------------------------------------------------------------
function [ci,bstat] = bootcper(stat,nboot,bootfun,alpha,weights,bootstrpOptions,varargin)
% corrected percentile bootstrap CI
% B. Efron (1982), "The jackknife, the bootstrap and other resampling
% plans", SIAM.

bstat = bootstrp(nboot,bootfun,varargin{:},'weights',weights,'Options',bootstrpOptions);

% stat is transformed to a normal random variable z0.
% z0 = invnormCDF[ECDF(stat)]
z_0 = fz0(bstat,stat);
z_alpha = norminv(alpha/2); % normal confidence point

% transform z0 back using the invECDF[normCDF(2z0-za)] and
% invECDF[normCDF(2z0+za)]
pct1 = 100*normcdf(2*z_0-z_alpha);
pct2 = 100*normcdf(2*z_0+z_alpha);

% inverse ECDF
m = numel(stat);
lower = zeros(1,m);
upper = zeros(1,m);
for i=1:m
    lower(i) = prctile(bstat(:,i),pct2(i),1);
    upper(i) = prctile(bstat(:,i),pct1(i),1);
end

% return
ci = [lower;upper];
end % bootcper()

%-------------------------------------------------------------------------
function [ci,bstat] = bootbca(stat,nboot,bootfun,alpha,weights,bootstrpOptions,varargin)
% corrected and accelerated percentile bootstrap CI
% T.J. DiCiccio and B. Efron (1996), "Bootstrap Confidence Intervals",
% statistical science, 11(3)

bstat = bootstrp(nboot,bootfun,varargin{:},'weights',weights,'Options',bootstrpOptions);

% same as bootcper, this is the bias correction
z_0 = fz0(bstat,stat);

% apply jackknife
try
    jstat = jackknife(bootfun,varargin{:},'Options',bootstrpOptions);
catch ME
    m = message('stats:bootci:JackknifeFailed',func2str(bootfun));
    throw(addCause(MException(m.Identifier,'%s',getString(m)),ME));
end
N = size(jstat,1);
if isempty(weights)
    weights = repmat(1/N,N,1);
else
    weights = weights(:);
    weights = weights/sum(weights);
end

% acceleration finding, see DiCiccio and Efron (1996)
mjstat = sum(bsxfun(@times,jstat,weights),1); % mean along 1st dim.
score = bsxfun(@minus,mjstat,jstat); % score function at stat; ignore (N-1) factor because it cancels out in the skew
iszer = all(score==0,1);
skew = sum(bsxfun(@times,score.^3,weights),1) ./ ...
    (sum(bsxfun(@times,score.^2,weights),1).^1.5) /sqrt(N); % skewness of the score function
skew(iszer) = 0;
acc = skew/6;  % acceleration

% transform back with bias corrected and acceleration
z_alpha1 = norminv(alpha/2);
z_alpha2 = -z_alpha1;
pct1 = 100*normcdf(z_0 +(z_0+z_alpha1)./(1-acc.*(z_0+z_alpha1)));
pct1(z_0==Inf) = 100;
pct1(z_0==-Inf) = 0;
pct2 = 100*normcdf(z_0 +(z_0+z_alpha2)./(1-acc.*(z_0+z_alpha2)));
pct2(z_0==Inf) = 100;
pct2(z_0==-Inf) = 0;

% inverse of ECDF
m = numel(stat);
lower = zeros(1,m);
upper = zeros(1,m);
for i=1:m
    lower(i) = prctile(bstat(:,i),pct2(i),1);
    upper(i) = prctile(bstat(:,i),pct1(i),1);
end

% return
ci = sort([lower;upper],1);
end % bootbca()

%-------------------------------------------------------------------------
function [ci,bstat] = bootstud(stat,nboot,bootfun,alpha,nbootstd,stderrfun,weights,bootstrpOptions,varargin)
% studentized bootstrap CI with bootstrp to estimate the se
% T.J. DiCiccio and B. Efron (1996), "Bootstrap Confidence Intervals",
% statistical science, 11(3)

% Should we use stderrfun to compute st. dev. or an inner bootstrap loop?
if isempty(stderrfun) % studentized with bootstrap error
    if nbootstd<=0 || nbootstd~=round(nbootstd)
        error(message('stats:bootci:BadNbootstd'))
    end
else  % studentized with stderrfun fun
    % error check for stderrfun
    try
        out=stderrfun(varargin{:});
    catch ME
        m = message('stats:bootci:BadStderr',func2str(stderrfun));
        throw(addCause(MException(m.Identifier,'%s',getString(m)),ME));
    end
    if any(~isfinite(out))
        error(message('stats:bootci:NonfiniteStderr'));
    end
    if isvector(out)
        out = out(:)';
    end
    if size(stat)~=size(out)
        error(message('stats:bootci:BadStderrSize'));
    end
end;

% bootstrap
[bstat,bootsam] = bootstrp(nboot,bootfun,varargin{:}, ...
    'weights',weights,'Options',bootstrpOptions); % bootstrap statistics

% find non-scalar data in varargin
la = length(varargin);
scalard = zeros(la,1);
for k = 1:la
   [row,col] = size(varargin{k});
   if max(row,col) == 1
      scalard(k) = 1;
   end
   if row == 1 && col ~= 1
      varargin{k} = varargin{k}(:);
   end
end

% bootstrap generated bootstrap replica to get student errors
N = size(bootsam,1);
sd_t = zeros(nboot,numel(stat));
for b=1:nboot
    db = cell(la,1);
    for k = 1:la
        % store the bootstrap data samples in a cell array
        if scalard(k) == 0
            db{k} = varargin{k}(bootsam(:,b),:);
        else
            db{k} = varargin{k};
        end
    end
    if isempty(weights)
        w = [];
    else
        w = weights(bootsam(:,b));
    end
    if isempty(stderrfun)
        bstatstd = bootstrp(nbootstd,bootfun,db{:},'weights',w,...
            'Options',bootstrpOptions);
        sd_t(b,:) = std(bstatstd,0,1);
    else
        bstatstd = stderrfun(db{:});
        sd_t(b,:) = reshape(bstatstd,1,numel(stat));
    end
end

% studentized statistics
above0 = sd_t>0;
mbstat = bsxfun(@minus,bstat,stat);
tstat = zeros(size(mbstat));
tstat(above0) = mbstat(above0)./sd_t(above0);

% percentiles for the studentized stats are computed.
lower = prctile(tstat,100*alpha/2,1);
upper = prctile(tstat,100*(1-alpha/2),1);

% back to the original stats from the studentized stats
lower = lower.*std(bstat,0,1) + stat;
upper = upper.*std(bstat,0,1) + stat;

% return
ci = [lower; upper];
end

% -------------------------
function z0=fz0(bstat,stat)
% Compute bias-correction constant z0
z0 = norminv(mean(bsxfun(@lt,bstat,stat),1) + mean(bsxfun(@eq,bstat,stat),1)/2);
end   % fz0()


function [out_val,out_ind,nevals]=tt_stat(tt, varargin)
% Compute statistics (min,max) of a tt_tensor via maxvol cross
%   [out_val,out_ind,nevals]=tt_stat(tt, varargin)
%
% Tries to compute values and indices of the sought quantities, listed in
% varargin. Possible choices are (case insensitive):
%   'LM', 'SM': largest, resp. smallest in magnitude entries,
%   'LR', 'SR': entries having largest, resp. smallest real parts, or
%   'LI', 'SI': entries having largest, resp. smallest imaginary parts.
% out_val is a 1 x N vector of output values, listed in the same order as
% in varargin, out_ind is a d x N matrix of indices, where d is the
% dimension of tt. The number of inspected elements is returned in nevals.
%
% Additional parameters can be given in varargin in the form
% 'param_name', param_value:
%       'kickrank': enrich tt by random subspaces of sizes kickrank. Larger
%       kickrank may improve the accuracy in "difficult" cases, for the price
%       of larger neval. Default value is 0.
%       'nswp': number of alternating iterations (default 3)
%
% Examples:
%   Mimic the call to tt_abs_max, tt_max and tt_min simultaneously:
%       [out,ind,neval]=tt_stat(tt, 'lm', 'lr', 'sr');
%   The same with kickrank set to 5:
%       [out,ind,neval]=tt_stat(tt, 'lm', 'lr', 'sr', 'kickrank', 5);
%   Minimize the banana function on a grid 32768 x 32768:
%       n = 15;
%       x = tkron(tt_x(2,n)/2^n, tt_ones(2,n));
%       y = tkron(tt_ones(2,n), tt_x(2,n)/2^n);
%       f = (0.5-x).^2 + 100*(y-x.^2).^2;
%       v = tt_stat(f,'sm','nswp',11)
%   This should give v ~ 2e-13.
%
% The difference with tt_*max* is the usage of dnr^2 entries
% instead of dr^2 (outputs are sought over TT blocks, not density matrices).
% Besides, it can carry out several iterations with random enrichment.

kickrank = 0; % random enrich may actually perturb the values, but sometimes it converges faster
nswp = 3; % these seem to be enough quite often

if (numel(varargin)==0)
    error('at least one sought quantity is needed');
end;

% Distinguish soughts from parameters
soughts = cell(1,numel(varargin));
j = 1;
while (j<=numel(varargin))
    if    ((strcmpi(varargin{j}, 'lm'))||(strcmpi(varargin{j}, 'sm'))||...
           (strcmpi(varargin{j}, 'lr'))||(strcmpi(varargin{j}, 'sr'))||...
           (strcmpi(varargin{j}, 'li'))||(strcmpi(varargin{j}, 'si')))
       soughts{j} = varargin{j};
    end;
    if (strcmpi(varargin{j}, 'kickrank'))
        kickrank = varargin{j+1};
        j = j+1;
    end;
    if (strcmpi(varargin{j}, 'nswp'))
        nswp = varargin{j+1};
        j = j+1;
    end;
    j = j+1;
end;
soughts(cellfun('isempty', soughts))=[];

d = tt.d;
n = tt.n;
r = tt.r;
tt = core2cell(tt);
r0 = r; % r will become different below

ind_global = cell(d+1,1);
ind_local = cell(d+1,1); % for reducing units
out_val = nan*ones(1,numel(soughts));
out_ind = zeros(d,numel(soughts));
% reductions
XX0 = cell(d+1,1);
XX0{1} = 1; XX0{d+1} = 1;

% We need to subtract unitary vectors to exlude already considered indices
% and select fresh ones
% Prepare storages for them
units = cell(d,numel(soughts));
for j=1:numel(soughts)
    for i=1:d
        units{i,j} = zeros(n(i),1);
    end;
end;
XU = cell(d+1,numel(soughts));
for j=1:numel(soughts)
    XU{1,j}=1; XU{d+1,j}=1;
end;
nevals = 0;

% Start iteration
swp = 0; % Zeroth iteration will be just QR and MV
dir = -1;
i = d;
while (swp<=nswp)
    if (swp>0)
        % Compute statistics
        q = reshape(tt{i}, r0(i), n(i)*r0(i+1));
        q = XX0{i}*q;
        q = reshape(q, r(i)*n(i), r0(i+1));
        q = q*XX0{i+1};
        q = reshape(q, r(i)*n(i)*r(i+1), 1);
        nevals = nevals + r(i)*n(i)*r(i+1); % number of evaluations
        ind_new = zeros(numel(soughts), 3);
        % We need to add the found extrema indices to maxvol, since they
        % will likely give good estimates in the next dimensions.
        pos = 0;
        for j=1:numel(soughts)
            out_changed = false;
            if (strcmpi(soughts{j}, 'lm'))
                [val,ind]=max(abs(q));
                out_changed = (val>abs(out_val(j)))||(isnan(out_val(j)));
            end;
            if (strcmpi(soughts{j}, 'lr'))
                [val,ind]=max(real(q));
                out_changed = (val>real(out_val(j)))||(isnan(out_val(j)));
            end;
            if (strcmpi(soughts{j}, 'li'))
                [val,ind]=max(imag(q));
                out_changed = (val>imag(out_val(j)))||(isnan(out_val(j)));
            end;
            if (strcmpi(soughts{j}, 'sm'))
                [val,ind]=min(abs(q));
                out_changed = (val<abs(out_val(j)))||(isnan(out_val(j)));
            end;
            if (strcmpi(soughts{j}, 'sr'))
                [val,ind]=min(real(q));
                out_changed = (val<real(out_val(j)))||(isnan(out_val(j)));
            end;
            if (strcmpi(soughts{j}, 'si'))
                [val,ind]=min(imag(q));
                out_changed = (val<imag(out_val(j)))||(isnan(out_val(j)));
            end;
            if (out_changed)
                out_val(j) = q(ind);
                ind = tt_ind2sub([r(i), n(i), r(i+1)], ind);
                % Add the new index to the expansion set
                ind_new(pos+1,:) = ind;
                pos = pos+1;
                if (i>1)&&(i<d)
                    out_ind(:,j) = indexmerge(ind_global{i}(ind(1),:), ind(2), ind_global{i+1}(ind(3),:));
                elseif (i==1)
                    out_ind(:,j) = indexmerge(ind(2), ind_global{i+1}(ind(3),:));
                else % i==d
                    out_ind(:,j) = indexmerge(ind_global{i}(ind(1),:), ind(2));
                end;
                % Update units vectors
                for k=1:d
                    units{k,j}=zeros(n(k),1);
                    units{k,j}(out_ind(k,j))=1;
                end;
                % and their reductions
                for k=1:i-1
                    XU{k+1,j} = XU{k,j}*units{k,j}.';
                    XU{k+1,j} = reshape(XU{k+1,j}, r(k)*n(k), 1);
                    XU{k+1,j} = XU{k+1,j}(ind_local{k+1}, :);
                end;
                for k=d:-1:i+1
                    XU{k,j} = units{k,j}*XU{k+1,j};
                    XU{k,j} = reshape(XU{k,j}, 1, n(k)*r(k+1));
                    XU{k,j} = XU{k,j}(:, ind_local{k});
                end;
            end;
        end;
        ind_new = ind_new(1:pos, :);

        % Subtract units to get fresh indices
        for j=1:numel(soughts)
            qu = XU{i,j}*reshape(units{i,j}*XU{i+1,j}, 1, n(i)*r(i+1));
            qu = reshape(qu, r(i)*n(i)*r(i+1),1);
            q = q-qu*out_val(j);
        end;
    end;

    if (dir>0)&&(i<d)
        q = reshape(q, r(i)*n(i), r(i+1));
        rr = zeros(r(i)*n(i), pos);
        % Enrich it by the extrema indices (just add corresponding unit
        % vectors)
        if (pos>0)
            for j=1:pos
                rr(tt_sub2ind([r(i), n(i)], ind_new(j,1:2)), j) = 1;
            end;
        end;
        % Enrich it by random
        rr2 = randn(r(i)*n(i), kickrank);
        [q,~]=qr([q, rr, rr2],0);
        rnew = size(q,2);

        % Maxvol
        ind = maxvol2(q);
        qq = q(ind,:);
        q = q/qq;
        % Indices
        ind_local{i+1} = ind; % for units
        ind_global{i+1} = indexmerge(ind_global{i}, (1:n(i))');
        ind_global{i+1} = ind_global{i+1}(ind,:);
        r(i+1) = rnew;
        % Reduction
        XX0{i+1} = reshape(tt{i}, r0(i), n(i)*r0(i+1));
        XX0{i+1} = XX0{i}*XX0{i+1};
        XX0{i+1} = reshape(XX0{i+1}, r(i)*n(i), r0(i+1));
        XX0{i+1} = XX0{i+1}(ind,:);
        % For units
        for j=1:numel(soughts)
            XU{i+1,j} = XU{i,j}*reshape(units{i,j}, 1, n(i));
            XU{i+1,j} = reshape(XU{i+1,j}, r(i)*n(i), 1);
            XU{i+1,j} = XU{i+1,j}(ind, :);
        end;
    elseif (dir<0)&&(i>1)
        % QR and right indices
        if (swp==0)
            % Don't perturb the original solution, just QR
            q = reshape(tt{i}, r(i), n(i)*r(i+1));
            rr = zeros(n(i)*r(i+1), 0);
        else
            q = reshape(q, r(i), n(i)*r(i+1));
            % Enrich by the extrema indices
            rr = zeros(n(i)*r(i+1), pos);
            if (pos>0)
                for j=1:pos
                    rr(tt_sub2ind([n(i) r(i+1)], ind_new(j,2:3)), j) = 1;
                end;
            end;
        end;
        % Enrich by random
        rr2 = randn(n(i)*r(i+1), kickrank*double(swp>0));
        [q,rv]=qr([q.',rr,rr2], 0);
        rv = rv(:,1:r(i)).';
        q = q.';
        rnew = size(q,1);
        % Maxvol
        ind = maxvol2(q.');
        qq = q(:, ind);
        q = qq\q;
        rv = rv*qq;
        % Indices
        ind_local{i} = ind; % for units
        ind_global{i} = indexmerge((1:n(i))', ind_global{i+1});
        ind_global{i} = ind_global{i}(ind,:);
        if (swp==0)
            tt{i} = reshape(q, rnew, n(i), r(i+1));
            % Cast rv to the next block
            q = reshape(tt{i-1}, r(i-1)*n(i-1), r(i));
            tt{i-1} = reshape(q*rv, r(i-1), n(i-1), rnew);
            r0(i) = rnew;
        end;
        r(i) = rnew;
        % Reduction
        XX0{i} = reshape(tt{i}, r0(i)*n(i), r0(i+1));
        XX0{i} = XX0{i}*XX0{i+1};
        XX0{i} = reshape(XX0{i}, r0(i), n(i)*r(i+1));
        XX0{i} = XX0{i}(:,ind);
        for j=1:numel(soughts)
            XU{i,j} = units{i,j}*XU{i+1,j};
            XU{i,j} = reshape(XU{i,j}, 1, n(i)*r(i+1));
            XU{i,j} = XU{i,j}(:, ind);
        end;
    end;

    i = i+dir;
    if ((dir>0)&&(i==d))||((dir<0)&&(i==1))
        dir = -dir;
        swp = swp+1;
    end;
end;

end

function [J]=indexmerge(varargin)
% Merges two or three indices in the little-endian manner
sz1 = max(size(varargin{1},1),1);
sz2 = max(size(varargin{2},1),1);
sz3 = 1;
if (nargin>2) % Currently allows only 3
    sz3 = max(size(varargin{3}, 1), 1);
end;
% J1 goes to the fastest index, just copy it
J1 = repmat(varargin{1}, sz2*sz3, 1);
% J2 goes to the middle
J2 = reshape(varargin{2}, 1, []);
J2 = repmat(J2, sz1, 1); % now sz1 ones will be the fastest
J2 = reshape(J2, sz1*sz2, []);
J2 = repmat(J2, sz3, 1);
J = [J1,J2];
if (nargin>2)
    % J3 goes to the slowest
    J3 = reshape(varargin{3}, 1, []);
    J3 = repmat(J3, sz1*sz2, 1); % now sz1 ones will be the fastest
    J3 = reshape(J3, sz1*sz2*sz3, []);
    J = [J,J3];
end;
end


classdef ParadigmRWSRCSP < ParadigmBase
    % Riemannian Weighted Subjects Regularized CSP (RWSRCSP)
    %
    % This paradigm implements RWSRCSP [1], which is a generalization of the Common Spatial Patterns
    % algorithms to calibration data comprising multiple subjects (or recordings). To train a model
    % for a particular "goal" (or target) subject using auxiliary data from other subjects, this
    % algorithm blends the subject-specific covariance matrices used in standard CSP and LDA with
    % weighted averages of covariance matrices from other subjects. The weighting is determined
    % based on the riemannian distance between the target subject and each respective other subject.
    % Also, the method can optionally average predictions obtained with different values of the
    % regularization parameter instead of doing a (costly) parameter search as suggested in [1].
    %
    % References:
    % [1] Lotte, F.
    %     "Signal processing approaches to minimize or suppress calibration time in oscillatory activity-based Brain-Computer Interfaces",
    %     Proceedings of the IEEE, vol. 103, no. 6, pp. 871-890, 2015
    %
    % Name:
    %   Riemannian Weighted Subjects Regularized Common Spatial Patterns
    %
    %                            Christian Kothe, Syntrogi
    %                            2015-07-24

    methods

        function defaults = preprocessing_defaults(self)
            defaults = {'FIRFilter',{'Frequencies',[6 8 28 32],'Type','minimum-phase'}, 'EpochExtraction',[0.5 3.5], 'Resampling',100};
        end

        function defaults = machine_learning_defaults(self)
            % set up the default parameters for machine learning; not necessary
            defaults = {'lda'};
        end

        function model = calibrate(self,varargin)
            % calibrate an SSRCSP model from a corpus of training sets
            args = arg_define(varargin, ...
                arg_norep({'collection','Collection'}), ...
                arg_norep({'goal_identifier','GoalIdentifier'}), ...
                arg({'patterns','PatternPairs'},3,uint32([1 1 64 10000]),'Number of CSP patterns (times two).'),...
                arg({'lambdas','Lambdas'},0.1:0.1:0.9,[0 1],'Covariance shrinkage. A range of shrinkage parameters to run over (classifiers will be averaged).'),...
                arg({'group_tasks_by','GroupTasksBy'},'subject',{'group','subject','day','montage','session','recording','block'},'Group tasks into. This allows to group the training data into tasks solved in a multi-task manner, e.g., such that data of a given subject forms one task. When hyper-parameters need to be optimized, this would usually be done using a basic blockwise cross-validation within each task.'), ...
                arg({'shrinkage_cov','ShrinkageCovariance','ShrinkageCov'},false,[],'Shrinkage covariance estimator. Whether to use shrinkage to estimate the covariance matrices.'), ...
                arg({'weight_bias','WeightedBias'}, false, [], 'Account for class priors in bias. If you do have unequal probabilities for the different classes, this should be enabled.'), ...
                arg({'weight_cov','WeightedCov'}, false, [], 'Account for class priors in covariance. If you do have unequal probabilities for the different classes, it makes sense to enable this.'), ...
                arg({'apply_to','ApplyTo'},'channels',{'channels','sources','components','full CSD'},'Apply classifier to. Allows to select the type of time series to apply this model to.'), ...
                arg_sub({'flt','SignalProcessing'}, self.preprocessing_defaults(), @flt_pipeline, 'Signal processing stages. These parameters control filter stages that run on the signal level; they can be enabled, disabled and configured for the given paradigm. The prediction operates on the outputs of this stage.'), ...
                arg_sub({'ml','MachineLearning'},{'Learner',self.machine_learning_defaults()},@ml_train,'Machine learning stage of the paradigm. Operates on the feature vectors that are produced by the feature-extraction stage.'),...
                arg({'arg_dialogsel','ConfigLayout'},self.dialog_layout_defaults(),[],'Parameters displayed in the config dialog. Cell array of parameter names to display (dot-notation allowed); blanks are translated into empty rows in the dialog. Referring to a structure argument lists all parameters of that struture, except if it is a switchable structure - in this case, a pulldown menu with switch options is displayed.','type','cellstr','shape','row'));

            % if this is run on a worker, we'll set the cache capacity to zero since no machine
            % has enough RAM to hold multiple workers' copies of the corpus in memory
            on_worker = hlp_iscaller('par_worker');
            if on_worker
                global tracking;
                tracking.cache.capacity = 0;
            end

            % split data into reference data (of goal subject) and remaining data
            [refsets,remaining] = utl_collection_closest(args.collection,args.goal_identifier);

            % recombine and move ref data to the beginning (because this paradigm will later
            % extract the weights learned for the first task) and turn into struct array for
            % convenience
            corpus = [refsets{:} remaining{:}];

            % determine group membership
            group_membership = {corpus.(args.group_tasks_by)};
            if iscellstr(group_membership)
                groups = unique(group_membership);
            else
                groups = num2cell(unique([group_membership{:}]));
            end

            % initialize per-class covariance matrices
            CovCSP = deal(cell(2,length(groups)));
            CovLDA = deal(cell(2,length(groups)));
            % for each group...
            for s=length(groups):-1:1
                matches = find(cellfun(@(x)isequal(x,groups{s}),group_membership));
                matchdata = cell(1,length(matches));

                % collect all matching data sets...
                for p=1:length(matches)
                    matchdata{p} = corpus(matches(p));
                    if length(matchdata{p}.streams) > 1
                    disp_once('Note: ParadigmRWCSP will use only the first data stream of a recording (no support for multi-modal data).'); end
                    matchdata{p} = matchdata{p}.streams{1};
                end

                % concatenate them into a single set
				procdata = set_concat(matchdata{:});
				% and preprocess the result (with further settings/overrides according to args.flt)
                procdata = flt_pipeline('Signal',procdata, args.flt); %#ok<*NODEF>
                if on_worker
                    % if we're running on a worker we don't cache the result due to memory
                    % constraints
                    procdata = exp_eval(procdata);
                else
                    procdata = exp_eval_optimized(procdata);
                end
                % extract data
                switch args.apply_to
                    case 'channels'
                        X = procdata.data;
                    case 'components'
                        X = reshape((procdata.icaweights*procdata.icasphere)*procdata.data(procdata.icachansind,:),[],procdata.pnts,procdata.trials);
                    case 'sources'
                        X = procdata.srcpot;
                    case 'full CSD'
                        X = procdata.srcpot_all;
                end
                X(~isfinite(X(:))) = 0;
                for k=1:2
                    trials{k} = exp_eval(set_picktrials(procdata,'rank',k));
                    % calculate the CSP class covariance matrix
                    if args.shrinkage_cov
                        CovCSP{k,s} = hlp_diskcache('featuremodels',@cov_shrink,reshape(trials{k}.data,size(trials{k}.data,1),[])');
                    else
                        CovCSP{k,s} = cov(reshape(trials{k}.data,size(trials{k}.data,1),[])');
                    end
                    CovCSP{k,s}(~isfinite(CovCSP{k,s})) = 0;
                end
                % solve CSP
                [V,D] = eig(CovCSP{1,s},CovCSP{1,s}+CovCSP{2,s}); %#ok<NASGU>
                featuremodel.filters = V(:,[1:args.patterns end-args.patterns+1:end]);
                features = self.feature_extract(procdata,featuremodel); % nT x nF
                % get the LDA covariance matrix per class
                %targets = [signal.epoch.target];
                for k=1:2
                    CovLDA{k,s} = cov(features([trials{k}.epoch.targnum],:)); end
            end

%             % check for and remove bad data
%             remove = [];
%             for s=1:length(features)
%                 if size(features{s},3) < 2
%                     fprintf('Encountered bad data at subject %s/%i.\n',corpus(s).streams{1}.parts{2}.parts{1:2});
%                     remove(end+1) = s;
%                 end
%             end
%             if ~isempty(remove)
%                 fprintf('Removing bad data...\n');
%                 scales(remove) = [];
%                 features(remove) = [];
%                 targets(remove) = [];
%                 transforms(remove) = [];
%                 % data_weights(remove) = [];
%             end

            % calculate Riemannian weights for each subject
            Covs = vertcat(CovCSP,CovLDA);
            % for each type of covariance matrix...
            for c=size(Covs,1):-1:1
                Ctarg = Covs{c,1};

                % calc distance to all other covs
                for s=size(Covs,2):-1:2
                    Cother = Covs{c,s};
                    [V,D] = eig(Ctarg\Cother);
                    dist(s) = sqrt(sum(log(diag(D)).^2));
                end

                % calc weighting
                sumdist = sum(dist(2:end));
                for s=2:size(Covs,2)
                    weights(s) = 1./(dist(s)/sumdist); end

                % use it to avg the covs
                Avg = zeros(length(Ctarg));
                for s=2:size(Covs,2)
                    Avg = Avg + weights(s)*Covs{c,s}; end
                AvgCovs{c} = Avg;
            end

            % for each reg. param, solve a classifier
            signal = procdata;
            targets = [signal.epoch.target];
            classes = unique(targets);
            for li = length(args.lambdas):-1:1
                lam = args.lambdas(li);
                for c=size(Covs,1):-1:1
                    BlendCovs{c} = lam*Covs{c,1} + (1-lam)*AvgCovs{c}; end
                [V,D] = eig(BlendCovs{1},BlendCovs{1}+BlendCovs{2}); P = inv(V); %#ok<NASGU>
                % train CSP part
                model.featuremodel.filters{li} = V(:,[1:args.patterns end-args.patterns+1:end]);
                model.featuremodel.patterns{li} = P([1:args.patterns end-args.patterns+1:end],:);
                % extract features
                trials = self.feature_extract(signal,model.featuremodel,li);
                % train LDA part
                for c = 1:2
                    X = trials(targets==classes(c),:);
                    n{c} = size(X,1);
                    mu{c} = mean(X,1);
                    sig{c} = BlendCovs{c+2};
                end
                ns = quickif(args.weight_cov,n,{1 1});
                nb = quickif(args.weight_bias,n,{1 1});
                % do the math
                mu_both = (mu{1}*nb{2} + mu{2}*nb{1}) / (nb{1}+nb{2});
                sig_both = (sig{1}*ns{1} + sig{2}*ns{2}) / (ns{1}+ns{2});
                w = (mu{2} - mu{1}) / sig_both;
                w = w / (mu{2}*w' - mu_both*w');
                model.predictivemodel.model{li} = struct('w',{w}, 'b',{mu_both*w'}, 'classes',{classes});
            end

            % set the filter graph based on the reference data
            model.tracking.filter_graph = signal;
            % also store channel locations for model visualization
            model.chanlocs = signal.chanlocs;
            model.classes = classes;
        end

        function predictions = predict(self,bundle,model)
            % for each lambda
            for m=length(model.featuremodel.filters):-1:1
                % extract features
                features = self.feature_extract(bundle.streams{1},model.featuremodel.filters{m});
                % apply classifier
                raw_preds(:,m) = features*model.predictivemodel.model{m}.w' - model.predictivemodel.model{m}.b;
            end
            % average all predictions
            raw_labels = mean(raw_preds,2);
            raw_labels = min(+1,max(-1,raw_labels));
            predictions = {'disc', [(1-raw_labels)/2 1-(1-raw_labels)/2], model.classes};
        end

        function features = feature_extract(self,signal,featuremodel,lam)
            if isstruct(featuremodel)
                featuremodel = featuremodel.filters; end
            if iscell(featuremodel)
                featuremodel = featuremodel{lam}; end
            % extract log-variance features from an epoched and preprocessed recording
            features = zeros(size(signal.data,3),size(featuremodel,2));
            for t=1:size(signal.data,3)
                features(t,:) = sum((signal.data(:,:,t)'*featuremodel).^2,1); end
            features = log(features/size(signal.data,2));
        end

        function visualize(self,varargin) %#ok<*INUSD>
            % visualize an mklCSP model
            args = arg_define(varargin, ...
                arg_norep({'model','Model'},[],[],'BCI Model to visualize.'), ...
                arg({'patterns','PlotPatterns'},true,[],'Plot patterns instead of filters. Whether to plot spatial patterns (forward projections) rather than spatial filters.'), ...
                arg({'paper','PaperFigure'},false,[],'Use paper-style font sizes. Whether to generate a plot with font sizes etc. adjusted for paper.'));

            f = figure;
            % get number of pairs, and index of pattern per subplot
            np = size(args.model.featuremodel.patterns,1)/2;
            idx = [1:np 2*np:-1:np+1];
            % for each CSP pattern...
            for p=1:np*2
                subplot(2,np,p,'Parent',f);
                if args.patterns
                    topoplot(args.model.featuremodel.patterns(idx(p),:),args.model.featuremodel.chanlocs);
                else
                    topoplot(args.model.featuremodel.filters(:,idx(p)),args.model.featuremodel.chanlocs);
                end
                t = title(['CSP Pattern ' num2str(idx(p))]);
                if args.paper
                    set(t,'FontUnits','normalized');
                    set(t,'FontSize',0.1);
                end
            end
        end

        function layout = dialog_layout_defaults(self)
            % define the default configuration dialog layout
            layout = {'SignalProcessing.Resampling.SamplingRate', 'SignalProcessing.FIRFilter.Frequencies', ...
                'SignalProcessing.FIRFilter.Type', 'SignalProcessing.EpochExtraction', '', ...
                'PatternPairs', 'CovarianceShrinkage', '', 'MachineLearning.Learner'};
        end

    end
end

% (turn off a few editor warnings because some actual implementations are missing in this file)
%#ok<*INUSD,*STOUT,*MANU>


function [infixOut, notaVAR] = pref2inf(prefixIn)
%PREF2INF   Convert an expression on prefix form to infix form
%   [INFIXOUT, NOTAVAR] = PREF2INF(PREFIXIN) goes recursively through the
%   expression PREFIXIN, which is on prefix form. The output, INFIXOUT, is a
%   string, representing the expression on infix form. NOTAVAR is a string, that
%   corresponds to the second variable in the kernel for FRED and VOLT
%   operators.

% Copyright 2017 by The University of Oxford and The Chebfun Developers.
% See http://www.chebfun.org/ for Chebfun information.

prefCounter = 1;
NOTAVAROUT = [];
infixOut = getInfix();
notaVAR = NOTAVAROUT;
NOTAVAROUT = [];

function infixOut = getInfix()
%GETINFIX   Return the infix form.

next = char(prefixIn(prefCounter,2));

% Steep through the prefix expression, dealing with each different type of
% tokens separately.
if ( ~isempty(strmatch('OP', next)) )
    % Increase the counter.
    prefCounter = prefCounter + 1;

    % We are dealing with a binary operator, so obtain two infix expressions.
    exp1 = getInfix();
    exp2 = getInfix();

    % We now return different outputs depending on which operator we have.
    switch next
        case 'OP='
            infixOut = [exp1, '=', exp2];

        case 'OP+'
            if ( strcmp(exp1, '0') && strcmp(exp2, '0') )
                % Don't want to introduce unnecessary zeros.
                infixOut = '';
            elseif  ( strcmp(exp1, '0') )
                infixOut = exp2;
            elseif ( strcmp(exp2, '0') )
                infixOut = exp1;
            else
                infixOut = ['(', exp1, '+', exp2, ')'];
            end

        case 'OP-'
            if ( strcmp(exp1, '0') && strcmp(exp2, '0') )
                infixOut = '';
            elseif  ( strcmp(exp1, '0') )
                infixOut = ['-',exp2];
            elseif ( strcmp(exp2, '0') )
                infixOut = exp1;
            else
                infixOut = ['(', exp1, '-', exp2, ')'];
            end

        case 'OP*'
            % Check whether we have some special arguments, which we know we can
            % simplify.
            if ( strcmp(exp1, '1') && strcmp(exp2, '1') )
                infixOut = '1';
            elseif  ( strcmp(exp1, '1') )
                infixOut = exp2;
            elseif ( strcmp(exp2, '1') )
                infixOut = exp1;
            elseif  ( strcmp(exp1, '-1') )
                infixOut = ['-', exp2];
            elseif ( strcmp(exp2, '-1') )
                infixOut = ['-', exp1];
            elseif ( strcmp(exp1, '0') || strcmp(exp2, '0') )
                infixOut = '0';
            elseif ( strcmp(exp1, '-0') || strcmp(exp2, '-0') )
                infixOut = '0';
            else
                infixOut = ['(', exp1, '.*', exp2, ')'];
            end

        case 'OP/'
            if ( strcmp(exp2, '1') )
                % Division by 1 is trivial.
                infixOut = exp1;
            else
                infixOut = ['(', exp1, './', exp2, ')'];
            end

        case 'OP^'
            infixOut = [exp1, '.^(', exp2, ')'];

        case {'OP>', 'OP>=', 'OP<', 'OP<='}
            nextSym = next(3:end);
            infixOut = ['(', exp1, nextSym, exp2, ')'];
    end

elseif ( strcmp(next, 'FUNC1') )
    % A function which has one argument.
    nextFun = char(prefixIn(prefCounter, 1));
    prefCounter = prefCounter + 1;
    funcArg = getInfix();
    infixOut = [nextFun, '(', funcArg , ')'];

elseif ( strcmp(next, 'FUNC2') )
    % A function with two arguments.
    nextFun = char(prefixIn(prefCounter, 1));
    prefCounter = prefCounter + 1;
    % Obtain two infix expressions.
    funcArg1 = getInfix();
    funcArg2 = getInfix();

    % Some methods need a special treatment.
    if ( (strcmp(nextFun, 'diff') || strcmp(nextFun, 'cumsum')) && ...
            strcmp(funcArg2, '1') )
        infixOut = [nextFun, '(', funcArg1, ')'];

    elseif ( any(strcmp(nextFun, {'fred', 'volt'})) )
        % For fred() and volt(), we need to find what arguments denote the
        % kernel variable.
        [ignored1, yFredVar, ignored2, ignored3, xFredVar] = ...
            stringParser.lexer(funcArg1, 'bvp');
        anonStr = ['@(' xFredVar{1} ',' yFredVar{1} ')'];
        infixOut = [nextFun, '(', anonStr, funcArg1 , ',', funcArg2 ,  ')'];
        NOTAVAROUT = [NOTAVAROUT ;  yFredVar];

    else
        infixOut = [nextFun, '(', funcArg1 , ',', funcArg2 ,  ')'];
    end

elseif ( strcmp(next, 'FUNC3') )
    % Methods with three arguments.
    nextFun = char(prefixIn(prefCounter,1));
    prefCounter = prefCounter + 1;
    funcArg1 = getInfix();
    funcArg2 = getInfix();
    funcArg3 = getInfix();
    infixOut = [nextFun, '(', funcArg1 , ',', funcArg2, ',' , funcArg3 ')'];

elseif ( strcmp(next(1:end-1), 'DER') )
    % Differentiation via the ' symbol.
    prefCounter = prefCounter + 1;
    derivArg = getInfix();
    derivOrder = next(4:end);
    if ( ~strcmp(derivOrder, '1') )
        infixOut = ['diff(', derivArg, ',' derivOrder, ')'];
    else
        infixOut = ['diff(', derivArg, ')'];
    end

elseif ( ~isempty(strmatch('UN', next)) )
    % Unary operators
    nextUnary = char(prefixIn(prefCounter, 1));
    prefCounter = prefCounter + 1;
    unaryArg = getInfix();
    % Only care about unary -
    if ( strcmp(nextUnary, '-') )
        infixOut = [nextUnary, unaryArg];
    else
        infixOut = unaryArg;
    end

elseif ( strmatch(next, 'COMMA') )
    % The , symbol.
    prefCounter = prefCounter + 1;
    exp1 = getInfix();
    exp2 = getInfix();
    infixOut = ['(', exp1, ';', exp2, ')'];
else
    % We have arrived at a terminal token. Return it and increase the counter.
    infixOut = char(prefixIn(prefCounter,1));
    prefCounter = prefCounter + 1;
end

end

% Initialization for simplification.
newlen = length(infixOut);
len = inf;

% Do some simplifications on the returned string. In particular, we want to
% change -- to +, a -+ to -, a +- to -, and ++ to +.
while ( newlen ~= len )     % Loop until length doesn't change anymore.
    len = newlen;

    k = 1;
    while ( k < (numel(infixOut) - 1) )
        if ( strcmp(infixOut(k), '-') )
            if ( strcmp(infixOut(k+1), '-') )
                infixOut(k) = '+';
                infixOut(k+1) = [];
            elseif ( strcmp(infixOut(k+1), '+') )
                infixOut(k+1) = [];
            else
                k = k+1;
            end
        elseif ( strcmp(infixOut(k), '+') )
            if ( strcmp(infixOut(k+1), '-') )
                infixOut(k) = '-';
                infixOut(k+1) = [];
            elseif ( strcmp(infixOut(k+1), '+') )
                infixOut(k+1) = [];
            else
                k = k+1;
            end
        else
            k = k+1;
        end
    end

    newlen = length(infixOut);

end

% Clear global variables
prefixIn = [];
prefCounter = [];

end


function b = dspsl ( ap, n, kpvt, b )

%*****************************************************************************80
%
%% DSPSL solves the real symmetric system factored by DSPFA.
%
%  Discussion:
%
%    To compute inverse(A) * C where C is a matrix with P columns:
%
%      call dspfa ( ap, n, kpvt, info )
%
%      if ( info /= 0 ) go to ...
%
%      do j = 1, p
%        call dspsl ( ap, n, kpvt, c(1,j) )
%      end do
%
%    A division by zero may occur if DSPCO has set RCOND == 0.0D+00
%    or DSPFA has set INFO /= 0.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    17 June 2005
%
%  Author:
%
%    MATLAB version by John Burkardt.
%
%  Reference:
%
%    Dongarra, Moler, Bunch and Stewart,
%    LINPACK User's Guide,
%    SIAM, (Society for Industrial and Applied Mathematics),
%    3600 University City Science Center,
%    Philadelphia, PA, 19104-2688.
%    ISBN 0-89871-172-X
%
%  Parameters:
%
%    Input, real ( kind = 8 ) AP(N*(N+1)/2), the output from DSPFA.
%
%    Input, integer N, the order of the matrix.
%
%    Input, integer KPVT(N), the pivot vector from DSPFA.
%
%    Input, real B(N), the right hand side.
%
%    Output, real B(N), the solution.
%

%
%  Loop backward applying the transformations and D inverse to B.
%
  k = n;
  ik = floor ( ( n * ( n - 1 ) ) / 2 );

  while ( 0 < k )

    kk = ik + k;

    if ( 0 <= kpvt(k) )
%
%  1 x 1 pivot block.
%
      if ( k ~= 1 )

        kp = kpvt(k);
%
%  Interchange.
%
        if ( kp ~= k )
          temp = b(k);
          b(k) = b(kp);
          b(kp) = temp;
        end
%
%  Apply the transformation.
%
        b(1:k-1) = daxpy ( k-1, b(k), ap(ik+1:ik+k-1), 1, b(1:k-1), 1 );

      end
%
%  Apply D inverse.
%
      b(k) = b(k) / ap(kk);
      k = k - 1;
      ik = ik - k;

    else
%
%  2 x 2 pivot block.
%
      ikm1 = ik - ( k - 1 );

      if ( k ~= 2 )

        kp = abs ( kpvt(k) );
%
%  Interchange.
%
        if ( kp ~= k-1 )
          temp = b(k-1);
          b(k-1) = b(kp);
          b(kp) = temp;
        end
%
%  Apply the transformation.
%
        b(1:k-2) = daxpy ( k-2, b(k), ap(ik+1:ik+k-2), 1, b(1:k-2), 1 );
        b(1:k-2) = daxpy ( k-2, b(k-1), ap(ikm1+1:ikm1+k-2), 1, b(1:k-2), 1 );

      end
%
%  Apply D inverse.
%
      km1k = ik + k - 1;
      kk = ik + k;
      ak = ap(kk) / ap(km1k);
      km1km1 = ikm1 + k - 1;
      akm1 = ap(km1km1) / ap(km1k);
      bk = b(k) / ap(km1k);
      bkm1 = b(k-1) / ap(km1k);
      denom = ak * akm1 - 1.0;
      b(k) = ( akm1 * bk - bkm1 ) / denom;
      b(k-1) = ( ak * bkm1 - bk ) / denom;
      k = k - 2;
      ik = ik - ( k + 1 ) - k;

    end

  end
%
%  Loop forward applying the transformations.
%
  k = 1;
  ik = 0;

  while ( k <= n )

    if ( 0 <= kpvt(k) )
%
%  1 x 1 pivot block.
%
      if ( k ~= 1 )
%
%  Apply the transformation.
%
        b(k) = b(k) + ddot ( k-1, ap(ik+1:ik+k-1), 1, b(1:k-1), 1 );
        kp = kpvt(k);
%
%  Interchange.
%
        if ( kp ~= k )
          temp = b(k);
          b(k) = b(kp);
          b(kp) = temp;
        end

      end

      ik = ik + k;
      k = k + 1;

    else
%
%  2 x 2 pivot block.
%
      if ( k ~= 1 )
%
%  Apply the transformation.
%
        b(k) = b(k) + ddot ( k-1, ap(ik+1:ik+k-1), 1, b(1:k-1), 1 );
        ikp1 = ik + k;
        b(k+1) = b(k+1) + ddot ( k-1, ap(ikp1+1:ikp1+k-1), 1, b(1:k-1), 1 );
        kp = abs ( kpvt(k) );
%
%  Interchange.
%
        if ( kp ~= k )
          temp = b(k);
          b(k) = b(kp);
          b(kp) = temp;
        end

      end

      ik = ik + k + k + 1;
      k = k + 2;

    end

  end

  return
end


%% Copyright (C) 2022 Alex Vong
%%
%% This file is part of OctSymPy.
%%
%% OctSymPy 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 software 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 software; see the file COPYING.
%% If not, see <https://www.gnu.org/licenses/>.

%% -*- texinfo -*-
%% @documentencoding UTF-8
%% @defmethod  @@sym piecewise (@var{cond1}, @var{val1}, @var{cond2}, @var{val2}, @dots{})
%% @defmethodx @@sym piecewise (@var{cond1}, @var{val1}, @var{cond2}, @var{val2}, @dots{}, @var{else_val})
%% Construct piecewise function.
%%
%% The returned piecewise function evaluates to @var{val1} if @var{cond1}
%% holds, @var{val2} if @var{cond2} holds, @dots{} etc.  In the case where none
%% of the conditions hold, it evaluates to @var{else_val} if provided.  If
%% @var{else_val} is not provided, it evaluates to @code{nan}.
%%
%% Examples:
%% @example
%% @group
%% syms x real
%% f = piecewise (abs (x) < 1, exp (- 1 / (1 - x^2)), abs (x) >= 1, 0)
%%   @result{} f = (sym)
%%       ⎧  -1
%%       ⎪ ──────
%%       ⎪      2
%%       ⎨ 1 - x
%%       ⎪ℯ        for │x│ < 1
%%       ⎪
%%       ⎩   0      otherwise
%% @end group
%% @end example
%%
%% For this piecewise function, we can omit the redundant condition at the end:
%% @example
%% @group
%% syms x real
%% f = piecewise (abs (x) < 1, exp (- 1 / (1 - x^2)), 0)
%%   @result{} f = (sym)
%%       ⎧  -1
%%       ⎪ ──────
%%       ⎪      2
%%       ⎨ 1 - x
%%       ⎪ℯ        for │x│ < 1
%%       ⎪
%%       ⎩   0      otherwise
%% @end group
%% @end example
%%
%% @seealso{if}
%% @end defmethod

function f = piecewise (varargin)
  if nargin < 1
    print_usage ();
  end

  cmd = {'def pack_args_without_else(args):'
         '    return [(val, cond) for cond, val in unflatten(args, 2)]'
         'def pack_args(args):'
         '    if len(args) % 2 == 0:'
         '        return pack_args_without_else(args)'
         '    else:'
         '        return pack_args_without_else(args[:-1]) + [(args[-1], True)]'
         'return Piecewise(*pack_args(_ins))'
        };

  args = cellfun (@sym, varargin, 'UniformOutput', false);
  f = pycall_sympy__ (cmd, args{:});
end

%!test
%! % basic
%! syms x real
%! f = piecewise (abs (x) < 1, 1);
%! assert (isnan (subs (f, -1)));
%! assert (isequal (subs (f, 0), 1));
%! assert (isnan (subs (f, 1)));

%!test
%! % heaviside
%! syms x real
%! f = rewrite (heaviside (x, 1 / sym (2)), 'Piecewise');
%! g = piecewise (x < 0, 0, x == 0, 1 / sym (2), x > 0, 1);
%! assert (logical (simplify (f == g)));

%% FIXME: expand test suite, add SMT compat tests, ...


function K = lfmXlfmKernCompute(lfmKern1, lfmKern2, t1, t2)

% LFMXLFMKERNCOMPUTE Compute a cross kernel between two LFM kernels.
% FORMAT
% DESC computes cross kernel terms between two LFM kernels for
% the multiple output kernel.
% ARG lfmKern1 : the kernel structure associated with the first LFM
% kernel.
% ARG lfmKern2 : the kernel structure associated with the second LFM
% kernel.
% ARG t : inputs for which kernel is to be computed.
% RETURN K : block of values from kernel matrix.
%
% FORMAT
% DESC computes cross kernel terms between two LFM kernels for
% the multiple output kernel.
% ARG lfmKern1 : the kernel structure associated with the first LFM
% kernel.
% ARG lfmKern2 : the kernel structure associated with the second LFM
% kernel.
% ARG t1 : row inputs for which kernel is to be computed.
% ARG t2 : column inputs for which kernel is to be computed.
% RETURN K : block of values from kernel matrix.
%
% SEEALSO : lfmKernParamInit, lfmKernCompute, lfmKernParamInit
%
% COPYRIGHT : David Luengo, 2007, 2008, Mauricio Alvarez, 2008
%
% MODIFICATIONS : Neil D. Lawrence, 2007, 2008,
%

% KERN

if nargin < 4
    t2 = t1;
end
if size(t1, 2) > 1 || size(t2, 2) > 1
    error('Input can only have one column');
end
if lfmKern1.inverseWidth ~= lfmKern2.inverseWidth
    error('Kernels cannot be cross combined if they have different inverse widths.')
end

% Get length scale out.
sigma2 = 2/lfmKern1.inverseWidth;
sigma = sqrt(sigma2);

% Parameters of the kernel
alpha(1) = lfmKern1.damper./(2*lfmKern1.mass);
alpha(2) = lfmKern2.damper./(2*lfmKern2.mass);
omega(1) = sqrt(lfmKern1.spring./lfmKern1.mass - alpha(1)*alpha(1));
omega(2) = sqrt(lfmKern2.spring./lfmKern2.mass - alpha(2)*alpha(2));

% Creation of the time matrices

if isreal(omega)
    % Precomputations to increase speed
    gamma1 = alpha(1) + j*omega(1);
    gamma2 = alpha(2) + j*omega(2);
    preGamma(1) = gamma1 + gamma2;
    preGamma(2) = conj(gamma1) + gamma2;
    preConst = 1./preGamma;
    preExp1 = exp(-gamma1*t1);
    preExp2 = exp(-gamma2*t2);
    % Actual computation of the kernel
    sK = real(lfmComputeH3(gamma1, gamma2, sigma2, t1,t2,preConst, 0, 1) + ...
        lfmComputeH3(gamma2, gamma1, sigma2, t2,t1,preConst(2) - preConst(1), 0, 0).' + ...
        lfmComputeH4(gamma1, gamma2, sigma2, t1, preGamma, preExp2, 0, 1  ) + ...
        lfmComputeH4(gamma2, gamma1, sigma2, t2, preGamma, preExp1,0, 0 ).');
    if lfmKern1.isNormalised
        K0 = (lfmKern1.sensitivity*lfmKern2.sensitivity)/(4*sqrt(2)*lfmKern1.mass*lfmKern2.mass*prod(omega));
    else
        K0 = (sigma*sqrt(pi)*lfmKern1.sensitivity*lfmKern2.sensitivity)/(4*lfmKern1.mass*lfmKern2.mass*prod(omega));
    end
    K = K0*sK;
else
    % Precomputations to increase the speed
    preExp1 = zeros(length(t1),2);
    preExp2 = zeros(length(t2),2);
    gamma1_p = alpha(1) + j*omega(1);
    gamma1_m = alpha(1) - j*omega(1);
    gamma2_p = alpha(2) + j*omega(2);
    gamma2_m = alpha(2) - j*omega(2);
    preGamma(1) = gamma1_p + gamma2_p;
    preGamma(2) = gamma1_p + gamma2_m;
    preGamma(3) = gamma1_m + gamma2_p;
    preGamma(4) = gamma1_m + gamma2_m;
    preConst = 1./preGamma;
    preFactors(1) = preConst(2) - preConst(1);
    preFactors(2) = preConst(3) - preConst(4);
    preFactors(3) = preConst(3) - preConst(1);
    preFactors(4) = preConst(2) - preConst(4);
    preExp1(:,1) = exp(-gamma1_p*t1);
    preExp1(:,2) = exp(-gamma1_m*t1);
    preExp2(:,1) = exp(-gamma2_p*t2);
    preExp2(:,2) = exp(-gamma2_m*t2);
    % Actual computation of the kernel
    sK = (  lfmComputeH3(gamma1_p, gamma1_m, sigma2, t1,t2,preFactors([1 2]), 1) + ...
        lfmComputeH3(gamma2_p, gamma2_m, sigma2, t2,t1,preFactors([3 4]), 1).' + ...
        lfmComputeH4(gamma1_p, gamma1_m, sigma2, t1, preGamma([1 2 4 3]), preExp2, 1 ) + ...
        lfmComputeH4(gamma2_p, gamma2_m, sigma2, t2, preGamma([1 3 4 2]), preExp1, 1 ).');
    if lfmKern1.isNormalised
        K0 = (lfmKern1.sensitivity*lfmKern2.sensitivity)/(8*sqrt(2)*lfmKern1.mass*lfmKern2.mass*prod(omega));
    else
        K0 = (sigma*sqrt(pi)*lfmKern1.sensitivity*lfmKern2.sensitivity)/(8*lfmKern1.mass*lfmKern2.mass*prod(omega));
    end
    K = K0*sK;
end


function [IND INT]=imsphere_IND_INT(r,v_low,res,size_IM,Xp,Yp,Zp)

% function [I J K]=imsphere_IND_INT(r,v_low,res,size_IM,Xp,Yp,Zp)
% ------------------------------------------------------------------------
%
%
% Kevin Mattheus Moerman
% kevinmoerman@hotmail.com
% 07/08/2008
% ------------------------------------------------------------------------

%% Setting up voxel dimensions and voxel radius dimensions
v_high=v_low/res;
r_v_high=r/v_high(1);
r_v_low=r/v_low(1);

%% Calculating image coordinate of point in high resolution image
[Ip_high Jp_high Kp_high]=cart2im(Xp,Yp,Zp,v_high);

%% Finding the index of the voxel in which point is found or is closest to
Ir_high=round(Ip_high); Jr_high=round(Jp_high); Kr_high=round(Kp_high);

%% Calculating shift from the centre of this voxel
Ic_high=Ir_high-Ip_high; Jc_high=Jr_high-Jp_high; Kc_high=Kr_high-Kp_high;

%% Finding indices of points in a sphere at specified location
[I_sph_high J_sph_high K_sph_high]=sphere_index(r_v_high,Ic_high,Jc_high,Kc_high);

%Calculating image coordinate of point in low resolution image
[Ip_low Jp_low Kp_low]=cart2im(Xp,Yp,Zp,v_low);
%Finding the index of the voxel in which point is found
Ir_low=round(Ip_low); Jr_low=round(Jp_low); Kr_low=round(Kp_low);
%Calculating shift from the centre of this voxel
Ic_low=Ir_low-Ip_low; Jc_low=Jr_low-Jp_low; Kc_low=Kr_low-Kp_low;

%Calculating shift
I_low_high_shift=((Ir_low*res)-(res/2)+0.5)-Ir_high;
J_low_high_shift=((Jr_low*res)-(res/2)+0.5)-Jr_high;
K_low_high_shift=((Kr_low*res)-(res/2)+0.5)-Kr_high;

size_M=(res+(2*round(r_v_low+1)*res))*ones(1,3);
M=ones(size_M);
IJK_middle=round(size_M/2);

IJK_voxel_center(1)=IJK_middle(1)-I_low_high_shift;
IJK_voxel_center(2)=IJK_middle(2)-J_low_high_shift;
IJK_voxel_center(3)=IJK_middle(3)-K_low_high_shift;
%round_IJK_voxel_center=round(IJK_voxel_center);

I_sph_high=round(I_sph_high+IJK_voxel_center(1));
J_sph_high=round(J_sph_high+IJK_voxel_center(2));
K_sph_high=round(K_sph_high+IJK_voxel_center(3));
IND = sub2ind(size_M,I_sph_high,J_sph_high,K_sph_high);
M(IND)=0;

%disp(['High resolution sphere simulated using ', num2str(numel(IND)), ' voxels']);

%Calculating low resolution image
[M]=voxelate(M,res);

IND=find(M<1);
INT=M(IND);
[I,J,K]=ind2sub(size(M),IND);
IJK_middle=round(size(M)/2);
I=I-IJK_middle(1); J=J-IJK_middle(2); K=K-IJK_middle(3);

I=I+Ir_low;
J=J+Jr_low;
K=K+Kr_low;

[IND]=sub2ind(size_IM,I,J,K);

%% END

%%
% _*GIBBON footer text*_
%
% License: <https://github.com/gibbonCode/GIBBON/blob/master/LICENSE>
%
% GIBBON: The Geometry and Image-based Bioengineering add-On. A toolbox for
% image segmentation, image-based modeling, meshing, and finite element
% analysis.
%
% Copyright (C) 2006-2022 Kevin Mattheus Moerman and the GIBBON contributors
%
% 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/>.


function [fx,dfdx,dfdp] = f_HRF3(Xt,P,ut,in)
% Balloon (HRF) model evolution function in log-space
% function [fx,dfdx,dfdp] = f_HRF2(Xt,P,ut,in)
% This function evaluates the evolution function derived from the balloon
% model for the hemodynamic response function. It can be called in two
% ways: (i) as a "stand-alone" evolution function, whereby the system's
% states are hemodynamic states of the balloon model, or (ii) as a
% generalized observation function, where the real system's states are the
% neuronal states of a DCM for fMRI model. Note that the hemodynamic states
% are in log-space, for positivity constrints.

deltat = in.deltat;
n = size(Xt,1);

% Get parameters
[E0,V0,tau0,kaf,kas,epsilon,alpha] = BOLD_parameters;
nreg = numel(in.n5); %n./5;
ind1 = in.ind1;
ind2 = in.ind2;
ind3 = in.ind3;
ind4 = in.ind4;
n1 = in.n1;
n2 = in.n2;
n3 = in.n3;
n4 = in.n4;
n5 = in.n5;
epsilon = 1;
alpha = alpha.*exp(P(in.ind5));

E0 = 1./(1+exp(-(P(ind1)-0.6633)));
dsdp = E0.*(1-E0);
tau0 = tau0.*exp(P(ind2));
kaf = kaf.*exp(P(ind3));
kas = kas.*exp(P(ind4));

% hemodynamic states ...
if in.linearized
    x1 = zeros(nreg,1);
    x2 = ones(nreg,1);
    x3 = ones(nreg,1);
    x4 = ones(nreg,1);
else
    % vasodilatory signal s(t)
    x1 = Xt(n1,:);
    % blood inflow f(t)
    if ~in.logx2
        x2 = Xt(n2,:) + 1; % deviation to steady-state!
    else
        x2 = exp(Xt(n2,:)) + in.xshift;
    end
    % blood volume v(t)
    x3 = exp(Xt(n3,:)) + 0.*in.xshift;
    % dHb content q(t)
    x4 = exp(Xt(n4,:)) + 0.*in.xshift;
end
% blood outflow
fv = x3.^(1./alpha);
% d[blood flow]/dXt(3)
dfvdx = (1./alpha).*fv;
% oxygen extraction
ff = (1-(1-E0).^(1./x2))./E0;
% d[O2 extraction]/dXt(2)
if ~in.logx2
    dffdx = log(1-E0).*(1-E0).^(1./x2)./(E0.*x2.^2);
else
    dffdx = log(1-E0).*(1-E0).^(1./x2)./(E0.*x2);
end
% ... and flow field, derivatives, etc...
f = zeros(n,1);
J = zeros(n,n);
dfdp = zeros(size(P,1),n);

% Evaluate flow field
f(n1) = (epsilon.*ut - kas.*x1 - kaf.*(x2 - 1));
if ~in.logx2
    f(n2) = x1;
else
    f(n2) = x1./x2;
end
f(n3) = (x2 - fv)./(tau0.*x3);
f(n4) = (x2.*ff./x4 - fv./x3)./tau0;

% Evaluate jacobian and gradients wrt parameters
for i=1:nreg

    if ~in.logx2
        J(n1(i),n1(i):n1(i)+3) = [ -kas(i) , 1, 0, 0 ];
        J(n2(i),n1(i):n1(i)+3) = [ -kaf(i), 0, ...
            1./(tau0(i).*x3(i)), ...
            (ff(i)+x2(i).*dffdx(i))./(x4(i).*tau0(i))];
    else
        J(n1(i),n1(i):n1(i)+3) = [ -kas(i) , 1./x2(i), 0, 0 ];
        J(n2(i),n1(i):n1(i)+3) = [ -kaf(i).*x2(i), -x1(i)./x2(i), ...
            x2(i)./(tau0(i).*x3(i)), ...
            x2(i).*(ff(i)+dffdx(i))./(tau0(i).*x4(i))];
    end
    J(n3(i),n1(i):n1(i)+3) = [ 0, 0,...
        -f(n3(i)) - dfvdx(i)./(tau0(i).*x3(i)) , ...
        (fv(i)-dfvdx(i))./(tau0(i).*x3(i))];
    J(n4(i),n1(i):n1(i)+3) = [ 0, 0, 0, ...
        -(x2(i).*ff(i))./(tau0(i).*x4(i)) ];
    J(n5(i),n1(i)) = epsilon;

    if in.linearized

        tmp = log(1-E0(i))./E0(i);

        dfdp(ind1(i),n1(i):n1(i)+3) = [0,0,0,...
            -dsdp(i).*(1+tmp).*Xt(n2(i))./(tau0(i).*E0(i))];
        dfdp(ind2(i),n1(i):n1(i)+3) = [0,0,...
            (-Xt(n2(i))+Xt(n3(i))./alpha(i))./tau0(i),...
            -((ff(i)+dffdx(i)).*Xt(n2(i)) + (fv(i)-dfvdx(i)).*Xt(n3(i)) ...
            - ff(i).*Xt(n4(i)))./tau0(i)];
        dfdp(ind3(i),n1(i):n1(i)+3) = kaf(i).*[-Xt(n2(i)),0,0,0];
        dfdp(ind4(i),n1(i):n1(i)+3) = kas(i).*[-Xt(n1(i)),0,0,0];
        dfdp(in.ind5(i),n1(i):n1(i)+3) = [0,0,...
            Xt(n3(i))./(tau0(i).*alpha(i)),...
            Xt(n3(i))./(tau0(i).*alpha(i))];

    else

        % gradient wrt parameters
        dfdp(ind1(i),n1(i):n1(i)+3) = [0,0,0,...
            (((1-E0(i)).^(-1+1./x2(i)))-x2(i).*ff(i))...
            .*dsdp(i)./(tau0(i).*x4(i).*E0(i))];
        dfdp(ind2(i),n1(i):n1(i)+3) = [0,0,...
            -(x2(i) - fv(i))./(tau0(i).*x3(i)),...
            -(x2(i).*ff(i)./x4(i) - fv(i)./x3(i))./tau0(i)];
        dfdp(ind3(i),n1(i):n1(i)+3) = kaf(i).*[-x2(i)+1,0,0,0];
        dfdp(ind4(i),n1(i):n1(i)+3) = kas(i).*[-x1(i),0,0,0];
        dfdp(in.ind5(i),n1(i):n1(i)+3) = [0,0,...
            log(x3(i)).*fv(i)./(tau0(i).*x3(i).*alpha(i)),...
            log(x3(i)).*fv(i)./(tau0(i).*x3(i).*alpha(i))];

    end

end

% Apply Euler discretization
dfdp = deltat.*dfdp;
dfdx = eye(n) + deltat.*J;
if in.linearized
    fx = dfdx'*Xt;
else
    fx = Xt + deltat.*f;
end


classdef SeqObsModel_Gaussian < SeqObsModel
    % SeqObsModel : defines Gaussian emission model for observed data in HMM
    %   Allows for multivariate observations of any dimension.
    %   Assumes conjugate prior on mean and covariance (Normal-Inverse Wishart).
    %   Extends the base object "SeqObsModel", filling in Gaussian details
    methods
        % ====================================================== CONSTRUCTOR
        function obj = SeqObsModel_Gaussian( K, D)
            obj.K = K;
            obj.D = D;
            obj.theta = repmat( struct('mu',zeros(D,1),'invSigma', zeros(D,D) ), 1, K );
        end

        % setPrior() sets prior on emission parameters to specified values.
        % USAGE:
        %   To set to default values given data object and model params
        %     ObsM = ObsM.setPrior( data, modelParamStruct )
        %     where modelParamStruct has attributes:
        %        degFree : # deg of freedom
        %        precMu  : precision parameter for mean
        %        Scoef   : scalar multiplier for Scale Matrix
        %        doEmpCovScalePrior: boolean.
        %               true  = set ScaleMat to Scoef*emp covariance
        %               false = set ScaleMat to Scoef*eye(D)
        %  To initialize directly to given parameters
        %      ObsM = ObsM.setPrior( mu, precMu, degFree, ScaleMat )
        function obj = setPrior( obj, varargin )
            obj.priorDef = struct();
            obj.priorDef.doEmpCovScalePrior =0;
            obj.priorDef.Scoef = 0;
            if strfind( class(varargin{1} ), 'SeqData' )
                data = varargin{1};
                obsM = varargin{2};
                obj.prior.mu = zeros( data.D, 1 );
                obj.prior.precMu = obsM.precMu;
                obj.prior.degFree = max( obsM.degFree, data.D+2 );
                obj.priorDef.doEmpCovScalePrior = obsM.doEmpCovScalePrior;
                obj.priorDef.Scoef = obsM.Scoef;
                if obsM.doEmpCovScalePrior
                    obj.prior.ScaleMat = obsM.Scoef * cov( data.Xdata', 1);
                else
                    obj.prior.ScaleMat = obsM.Scoef * eye( data.D );
                end
            elseif length( varargin ) == 1
                PStruct = varargin{1};
                obj.prior.mu = PStruct.Mu(:);
                obj.prior.precMu = PStruct.precMu;
                obj.prior.degFree = PStruct.degFree;
                obj.prior.ScaleMat = PStruct.ScaleMat;
            else
                obj.prior.mu = varargin{1}(:);
                obj.prior.precMu = varargin{2};
                obj.prior.degFree = varargin{3};
                obj.prior.ScaleMat = varargin{4};
            end
        end

        % ==================================================== GET Xstats
        function Xstats = getXSuffStats( obj, Xkk )
            Xstats.nObs = size(Xkk,2);
            if Xstats.nObs > 0
                Xstats.Xsum = sum( Xkk, 2);
                Xstats.XXsum = Xkk*Xkk';
            else
                Xstats.Xsum = [];
                Xstats.XXsum = [];
            end
        end

        function obj = incXStats( obj, ii, data, stateSeq, featIDs )
            for kk = featIDs
                Xkk = data.seq(ii);
                Xkk = Xkk(:, stateSeq(ii).z == kk );
                nNew = size(Xkk,2);
                if length(obj.Xstats)<kk || obj.Xstats(kk).nObs == 0
                    obj.Xstats(kk).nObs = nNew;
                    obj.Xstats(kk).Xsum = sum(Xkk,2);
                    obj.Xstats(kk).XXsum = Xkk*Xkk';
                else
                    obj.Xstats(kk).nObs = obj.Xstats(kk).nObs + nNew;
                    obj.Xstats(kk).Xsum = obj.Xstats(kk).Xsum + sum(Xkk,2);
                    obj.Xstats(kk).XXsum = obj.Xstats(kk).XXsum + Xkk*Xkk';
                end
            end
        end

        function obj = decXStats( obj, ii, data, stateSeq, featIDs )
            for kk = featIDs
                Xkk = data.seq(ii);
                Xkk = Xkk(:, stateSeq(ii).z == kk );
                nNew = size(Xkk,2);
                if obj.Xstats(kk).nObs-nNew <= 0
                    obj.Xstats(kk).nObs=0;
                    obj.Xstats(kk).Xsum=[];
                    obj.Xstats(kk).XXsum=[];
                else
                    obj.Xstats(kk).nObs = obj.Xstats(kk).nObs - nNew;
                    obj.Xstats(kk).Xsum = obj.Xstats(kk).Xsum - sum(Xkk,2);
                    obj.Xstats(kk).XXsum = obj.Xstats(kk).XXsum - Xkk*Xkk';
                end
            end
        end

        % ==================================================== GET params
        function retStr = getParamDescr( obj, PP )
            if ~exist('PP','var')
                PP = obj.prior;
            end
            retStr = 'Norm-InvWish. Mu0=0, dFree=%d,';
            if obj.priorDef.doEmpCovScalePrior
                retStr = [retStr ' S0=%.2f*EmpCov.'];
            else
                retStr = [retStr ' S0=%.2f*eye.'];
            end
            retStr = sprintf( retStr, obj.prior.degFree, obj.priorDef.Scoef );
        end

        function PP = getPosteriorParams( obj, Xstats )
            if Xstats.nObs > 0
                N = Xstats.nObs;
                degFreeN = N + obj.prior.degFree;
                precMuN  = N + obj.prior.precMu;
                Xdiff = Xstats.Xsum/N - obj.prior.mu;
                XdiffMat = ( obj.prior.precMu * N  )/( obj.prior.precMu + N ) .* ( Xdiff*Xdiff' );
                CovMat = (Xstats.XXsum - Xstats.Xsum*Xstats.Xsum'/N);
                ScaleMatN = obj.prior.ScaleMat + CovMat + XdiffMat;
                MuN = (Xstats.Xsum + obj.prior.mu*obj.prior.precMu)/ ( precMuN );

                PP.mu = MuN;
                PP.precMu = precMuN;
                PP.degFree = degFreeN;
                PP.ScaleMat = ScaleMatN;
            else
                PP = obj.prior;
            end
        end

        % ==================================================== GET theta mean

        % Retrieve the expected value of theta under the given params
        % In Gaussian case,
        %   mu = PP.mu
        %   Sigma = PP.ScaleMat / (degFree-D-1 )
        function [theta] = getTheta_Mean( obj, PP )
            if ~exist('PP','var')
                PP = obj.prior;
            end
            theta.mu  = PP.mu;
            theta.invSigma = (PP.degFree-obj.D-1)*( PP.ScaleMat \ eye(obj.D) );
        end

        % ==================================================== GET theta samps
        function [theta, PP] = sampleTheta_FromParams( obj, PP )
            if ~exist('PP','var')
                PP = obj.prior;
            end
            [~, sqrtInvSigma] = randiwishart( PP.ScaleMat, PP.degFree );
            theta.mu = ( sqrt( PP.precMu ) .* sqrtInvSigma ) \randn(obj.D,1)  + PP.mu;
            theta.invSigma = sqrtInvSigma'*sqrtInvSigma;
        end

        function [theta, PN] = sampleTheta( obj, Xkk )
            if ~exist( 'Xkk', 'var' )
                Xkk = [];
            end
            Xstats = obj.getXSuffStats( Xkk );
            PN = obj.getPosteriorParams( Xstats );
            [theta] = obj.sampleTheta_FromParams( PN );
        end

        % ==================================================== GET log probs

        % Calc probability of a realization of a emission parameter "theta"
        %   given the provided prior parameter struct "PP"
        function logPr = calcLogPrTheta( obj, theta, PP )
            if ~exist('PP','var')
                PP = obj.prior;
            end
            logPr = calcLogPrNormalInvWishart( theta.mu, theta.invSigma, PP );
        end

        % Calculate soft evidence for particular sequence "ii"
        %   given all available emission parameters (indicated by "kIDs")
        %OUTPUT
        %  logSoftEv : K x T matrix, where entry k,t gives log probability
        %                of the t-th observation under theta(kk)
        function logSoftEv = calcLogSoftEv(obj, ii, data, kIDs )
            if ~exist( 'kIDs', 'var' )
                kIDs = 1:obj.K;
            end
            Xseq = data.seq(ii);
            T = size( Xseq,2);
            logSoftEv = -inf( obj.K, T );
            for kk = kIDs
                cholInvSigma = chol( obj.theta(kk).invSigma );
                logDetInvSigma = 2*sum( log( diag( cholInvSigma) ) );
                XdiffMu = bsxfun(@minus, Xseq, obj.theta(kk).mu );
                U = XdiffMu'*cholInvSigma';
                logSoftEv(kk,:) = 0.5*logDetInvSigma - 0.5*sum( U.^2,2);
            end
            logSoftEv = logSoftEv - 0.5*obj.D*log(2*pi);
        end

        % Calculate marginal probability of observed data
        %  given the state sequence assignments of every observation
        % USAGE:
        %  logPr = ThetaM.calcMargPrData()
        %    calculate using *stored* sufficient statistics in this object
        %  logPr = ThetaM.calcMargPrData( data, stateSeq )
        %    calculate using freshly computed suff stats
        function logPr = calcMargPrData(obj, data, stateSeq, ks )
            if ~exist( 'ks', 'var' )
                ks = 1:obj.K;
            end
            if exist( 'data', 'var' ) && ~isempty( data )
                obj = obj.updateAllXSuffStats(  horzcat(stateSeq(:).z), data, ks);
            end
            LOG_PI =  1.144729885849400;
            P0 = obj.prior;
            D = obj.D;
            logPr = zeros(1, obj.K);
            for kk = ks
                Nkk = obj.Xstats(kk).nObs;
                if Nkk > 0
                    PN = obj.getPosteriorParams( obj.Xstats(kk) );

                    logPr(kk) = ...
                        logMvGamma(0.5*PN.degFree,D) - logMvGamma(0.5*P0.degFree,D) ...
                        + 0.5*P0.degFree*log( det( P0.ScaleMat ) )   ...
                        - 0.5*PN.degFree*log( det( PN.ScaleMat ) ) ...
                        + 0.5*D*log( P0.precMu / PN.precMu );
                end
            end
            N = sum( [obj.Xstats(ks).nObs] );
            logPr = sum(logPr) - 0.5*N*D*LOG_PI;
        end

    end
end


function matrix = genMatCurlVIFE3D(pde,mesh,fem,meshI,femI,option)
%% Generate global matrices and load vector of FEM for 3D ellipic eq
%     -div(A grad u)  = f,    x\in \Omega
% INPUTS:
% pde --- given data function from equation, e.g.
%         pde.A --- diffusion coefficient
%         pde.f --- right hand side function
%         pde.gD --- Dirichlet boundary value function
%         pde.one --- constant function 1.
% mesh --- mesh structure.
% fem --- global degree of freedom of FEM
%
% OUTPUTS:
% matrix.S --- stiffness matrix (w/o boundary condition)
% matrix.A --- final FEM matrix (after boundary condition)
% matrix.rhsF --- load vector (w/o boundary condition)
% matrix.f --- final RHS matrix (after boundary condition)

% Last Modified: 08/07/2020 by Xu Zhang

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
g2ldofNint = femI.g2ldofNint;
node = meshI.node;
gdof = femI.gdof;
NEdof = size(gdof,1);

% %% 1. Stiffness Matrix
feEvalBas1 = @EvalNed1Bas3D;
feEvalBas2 = @EvalNed1Bas3D;

dof1 = 6; dof2 = 6; nloc = dof1*dof2;
ntID = find(mesh.tLoc > 0); ntN = length(ntID);

%% 1. Matrix on noninterface elements
MatInd = 1;
AN = fem.area(ntID);
gxN = fem.gx(ntID,:); gyN = fem.gy(ntID,:); gzN = fem.gz(ntID,:); gw = fem.gw;
XN = zeros(nloc*ntN, 1);

coefN = feval(pde.A,gxN,gyN,gzN);
Ibasx = cell(dof1,1); Ibasy = cell(dof1,1); Ibasz = cell(dof1,1);
Jbasx = cell(dof2,1); Jbasy = cell(dof2,1); Jbasz = cell(dof2,1);
t_e_orit = fem.t_e_orit(ntID,:);

for i = 1:dof1
    Ibasx{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 1).*t_e_orit(:,i);
    Ibasy{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 2).*t_e_orit(:,i);
    Ibasz{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 3).*t_e_orit(:,i);
end
for j = 1:dof2
    Jbasx{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 1).*t_e_orit(:,j);
    Jbasy{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 2).*t_e_orit(:,j);
    Jbasz{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 3).*t_e_orit(:,j);
end

IN = reshape(repmat(g2ldofNint(:,1:6),6,1),nloc*ntN,1);
JN = repmat(reshape(g2ldofNint(:,1:6),dof2*ntN,1),6,1);
ind = 0;
for i = 1:dof1
    for j = 1:dof2
        XN(ind+1:ind+ntN) = AN.*(sum(((Ibasx{i}.*(coefN.*Jbasx{j})).*gw'),2) + ...
            sum(((Ibasy{i}.*(coefN.*Jbasy{j})).*gw'),2) + ...
            sum(((Ibasz{i}.*(coefN.*Jbasz{j})).*gw'),2));
        ind = ind + ntN;
    end
end
ID = find(XN~=0);
SN = sparse(IN(ID),JN(ID),XN(ID),NEdof,NEdof);

%%%%%%

MatInd = 0;
AN = fem.area(ntID);
gxN = fem.gx(ntID,:); gyN = fem.gy(ntID,:); gzN = fem.gz(ntID,:); gw = fem.gw;
XN = zeros(nloc*ntN, 1);

coefN = feval(pde.B,gxN,gyN,gzN);
Ibasx = cell(dof1,1); Ibasy = cell(dof1,1); Ibasz = cell(dof1,1);
Jbasx = cell(dof2,1); Jbasy = cell(dof2,1); Jbasz = cell(dof2,1);

for i = 1:dof1
    Ibasx{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 1).*t_e_orit(:,i);
    Ibasy{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 2).*t_e_orit(:,i);
    Ibasz{i} = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, MatInd, 3).*t_e_orit(:,i);
end
for j = 1:dof2
    Jbasx{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 1).*t_e_orit(:,j);
    Jbasy{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 2).*t_e_orit(:,j);
    Jbasz{j} = feEvalBas2(fem.bas, ntID, gxN, gyN, gzN, j, MatInd, 3).*t_e_orit(:,j);
end

IN = reshape(repmat(g2ldofNint(:,1:6),6,1),nloc*ntN,1);
JN = repmat(reshape(g2ldofNint(:,1:6),dof2*ntN,1),6,1);
ind = 0;
for i = 1:dof1
    for j = 1:dof2
        XN(ind+1:ind+ntN) = AN.*(sum(((Ibasx{i}.*(coefN.*Jbasx{j})).*gw'),2) + ...
            sum(((Ibasy{i}.*(coefN.*Jbasy{j})).*gw'),2) + ...
            sum(((Ibasz{i}.*(coefN.*Jbasz{j})).*gw'),2));
        ind = ind + ntN;
    end
end
ID = find(XN~=0);
MN = sparse(IN(ID),JN(ID),XN(ID),NEdof,NEdof);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

dof = 6;  nloc = dof;
X = zeros(nloc*ntN, 1);

fN1 = feval(pde.f1,gxN,gyN,gzN);
fN2 = feval(pde.f2,gxN,gyN,gzN);
fN3 = feval(pde.f3,gxN,gyN,gzN);
ind = 0;
I = reshape(g2ldofNint(:,1:6),nloc*ntN,1);
for i = 1:dof
    ibas1 = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, 0, 1);
    ibas2 = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, 0, 2);
    ibas3 = feEvalBas1(fem.bas, ntID, gxN, gyN, gzN, i, 0, 3);
    X(ind+1:ind+ntN) = AN.*sum((ibas1.*fN1+ibas2.*fN2+ibas3.*fN3).*gw',2).*t_e_orit(:,i);
    ind = ind + ntN;
end
rhsN = sparse(I,1,X,NEdof,1);

if isfield(option,'HM')
    if option.HM == 1
        Atotal = SN - pde.kappa*MN + femI.KI - pde.kappa*femI.SI;
    end
else
    Atotal = SN + MN + femI.KI + femI.SI;
end

rhs = rhsN + femI.b;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% %% 3. Dirichlet Boundary Conditions

% non-interface edges
vSign = meshI.vSign;
eSign = sum(vSign(femI.gdof),2)/2;
eIdm = find(eSign==-1/2); eIdp = find(eSign==1/2);
[gew,gex,gey,gez] = gaussPedge(node(gdof(:,1),:),node(gdof(:,2),:),2);
% tu1 = sum(feval(pde.exactu1,gex,gey,gez).*gew,2);
% tu2 = sum(feval(pde.exactu2,gex,gey,gez).*gew,2);
% tu3 = sum(feval(pde.exactu3,gex,gey,gez).*gew,2);
tu1 = sum(feval(pde.exactu1,gex,gey,gez).*gew,2);
tu2 = sum(feval(pde.exactu2,gex,gey,gez).*gew,2);
tu3 = sum(feval(pde.exactu3,gex,gey,gez).*gew,2);
tu1(eIdm) = sum(feval(pde.um1,gex(eIdm,:),gey(eIdm,:),gez(eIdm,:)).*gew(eIdm,:),2);
tu1(eIdp) = sum(feval(pde.up1,gex(eIdp,:),gey(eIdp,:),gez(eIdp,:)).*gew(eIdp,:),2);
tu2(eIdm) = sum(feval(pde.um2,gex(eIdm,:),gey(eIdm,:),gez(eIdm,:)).*gew(eIdm,:),2);
tu2(eIdp) = sum(feval(pde.up2,gex(eIdp,:),gey(eIdp,:),gez(eIdp,:)).*gew(eIdp,:),2);
tu3(eIdm) = sum(feval(pde.um3,gex(eIdm,:),gey(eIdm,:),gez(eIdm,:)).*gew(eIdm,:),2);
tu3(eIdp) = sum(feval(pde.up3,gex(eIdp,:),gey(eIdp,:),gez(eIdp,:)).*gew(eIdp,:),2);
tgt = node(gdof(:,2),:) - node(gdof(:,1),:);
tgt = tgt./sum(tgt.^2,2).^(1/2);
tu = tu1.*tgt(:,1) + tu2.*tgt(:,2) + tu3.*tgt(:,3);

bcind = fem.bcind;
[bc,mapper] = boundaryEdge3D(node,gdof,bcind);
bdidx = zeros(NEdof,1);
isBdEdge = true(NEdof,1);
isBdEdge(mapper) = false;
bdidx(isBdEdge) = 1;
Tbd = spdiags(bdidx,0,NEdof,NEdof);
T = spdiags(1-bdidx,0,NEdof,NEdof);
A = T*Atotal*T + Tbd;

ub = tu;
ub(mapper) = 0;
rhsB = Atotal*ub;
f = rhs - rhsB;
f(isBdEdge) = tu(isBdEdge);

%% Outputs
matrix = struct('A', A, 'f', f,'tu',tu,'isBdEdge',isBdEdge);