function [tout, info] = transfdiff(opt, tp, tm)
% TRANSFDIFF Transform diffusion algorithm for sequence of images.
%
% This function is part of a larger algorithm that solves the following
% registration problem:
%
% We have a sequence of images I=1,2,...,N. We want to register each image
% to its adjacent neighbours to form a volume. But instead of registering
% I=2 to I=1, and then I=3 to the result, and then I=4 to the result, etc.,
% it can be shown that all images can be aligned in parallel using an
% iterative process that we call "transform diffusion registration".
%
% TRANSFDIFF solves the "transform diffusion" step in the algorithm:
%
% Let's assume that somebody has already registered each image I to its
% neighbours I-1 and I+1. TRANSFDIFF takes that set of registration
% parameters and "diffuses" them to produce a transform for each image.
% Applying these transforms to the images produces the globally aligned
% sequence.
%
%
% TRANSFDIFF applies a diffusion process to those neighbour transforms. The
% output is the N transforms that applied to the N images best align them.
%
% [TOUT, INFO] = TRANSFDIFF(OPT, ...)
%
%   OPT is a string with the name of the transform to apply, or a struct
%   with the transform name and algorithm parameters. Options common to all
%   transformations:
%
%     'MaxIter': (def 50) Stopping criterion. The algorithm stops after
%                MaxIter diffusion iterations.
%
%     'Alpha':   (def 0.45) Diffusion coefficient. The smaller the value of
%                alpha, more iterations are needed for the same result. But
%                alpha >0.45 may not smooth high frequencies very well, and
%                values >=0.5 cause oscillations.
%
%   TOUT is an array with N output transforms for the N slices.
%
%   INFO is a struct with information about the algorithm internals:
%
%     'NumIter': Number of diffusion iterations until stop condition.
%
% -------------------------------------------------------------------------
%
% * OPT='TranslationTransform'
% * OPT.Transform='TranslationTransform'
%
% [TOUT, INFO] = TRANSFDIFF(OPT, TP)
%
%   TP is a matrix with N-1 rows. TP(I, :) is the translation from slice I
%   to slice I+1.
%
%   This case assumes that the transforms are symmetric, i.e. the
%   translation from slice I+1 to I is -TP(I, :).
%
%   Specific OPT options:
%
%     'Epsilon': (def 0.0) Stopping criterion. It stops when all
%                translation components are small:
%                for all t, abs(t)<=Epsilon.
%
% -------------------------------------------------------------------------
%
% * OPT='AffineTransform'
% * OPT.Transform='AffineTransform'
%
% * OPT='EulerTransform'
% * OPT.Transform='EulerTransform'
%
% [TOUT, INFO] = TRANSFDIFF(OPT, TP)
%
%   TP is an array of affine transforms in homogeneous coordinates. TP(I,:)
%   is the affine transform from slice I to slice I+1. Each transform has
%   the format
%
%   TP(:, :, I) = [A 0]
%                 [d 1]
%
%   where A is a matrix and d is a translation vector, and the affine
%   transform of a row vector X is defined as Y=X*TP(:,:,I).
%
%   This case assumes that the transforms are symmetric, i.e. the
%   transformation from slice I+1 to I is inv(TP(I, :)).
%
%   Specific OPT options:
%
%     'Epsilon': (def 0.0) Stopping criterion. It stops when the norm of
%                the difference between the transformation and the identity
%                matrix is small:
%                for all t, norm(t-eye(3))<=Epsilon.
%
% -------------------------------------------------------------------------
%
% * OPT.Transform='BsplineTransform'
%
% [TOUT, INFO] = TRANSFDIFF(OPT, TP, TM)
%
%   TP is a matrix with N-1 rows. Each row contains the B-spline
%   coefficients with the transformation from slice I to slice I+1. The
%   format of the B-spline coefficient vector is
%   [p0x, p1x, ..., pNx, p0y, p1y, ..., pNy].
%
%   TM is similar to TP, but TM(I, :) is the transformation from slice I+1
%   to I.
%
%   Specific OPT options:
%
%     'Epsilon': (def 0.0) Stopping criterion. It stops when all
%                coefficient displacement components are small:
%                for all t, abs(t)<=Epsilon.
%
%     'tbsp': (def []) For Choi and Lee (2000) injectivity criteria. These
%             are sufficient criteria that guarantee the B-spline doesn't
%             fold over. Transform struct in elastix format. The
%             tbsp.TransformParameters are ignored, only the grid
%             information is used. If not provided or empty, the conditions
%             are not used as a stopping criterion, and the resulting
%             B-splines may fold over.
%
%
% -------------------------------------------------------------------------
%
% See also: transdiffreg.

% Author: Ramon Casero <rcasero@gmail.com>
% Copyright © 2015-2016 University of Oxford
% Version: 0.4.4
%
% University of Oxford means the Chancellor, Masters and Scholars of
% the University of Oxford, having an administrative office at
% Wellington Square, Oxford OX1 2JD, UK.
%
% This file is part of Gerardus.
%
% 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. The offer of this
% program under the terms of the License is subject to the License
% being interpreted in accordance with English Law and subject to any
% action against the University of Oxford being under the jurisdiction
% of the English Courts.
%
% You should have received a copy of the GNU General Public License
% along with this program.  If not, see <http://www.gnu.org/licenses/>.

% check arguments
switch (opt.Transform)

    case {'TranslationTransform', 'EulerTransform', 'AffineTransform'}

        narginchk(2, 2);

    case {'BSplineTransform'}

        narginchk(2, 3);

    otherwise

        error(['Transform ' opt.Transform ' not implemented'])

end
nargoutchk(0, 2);

% defaults
if (isempty(opt))
    error('OPT cannot be empty');
elseif (ischar(opt))
    aux = opt;
    clear opt
    opt.Transform = aux;
end
if (~isfield(opt, 'Epsilon'))
    opt.Epsilon = 0.0;
end
if (~isfield(opt, 'MaxIter'))
    opt.MaxIter = 50;
end
if (~isfield(opt, 'Alpha'))
    opt.Alpha = 0.45;
end
if (strcmp(opt.Transform, 'BSplineTransform'))
    if (~isfield(opt, 'tbsp'))
        opt.tbsp = [];
    end
end

if (opt.Alpha < 0 || opt.Alpha > 0.5)
    warning('alpha must be in the interval [0.0, 0.5]. alpha=0.5 can produce oscillations. We recommend alpha<=0.45')
end

% convert to array format if transforms are provided in elastix format
[tp, tp0] = elastix2mat(tp);
if (nargin > 2)
    [tm, tm0] = elastix2mat(tm);
end

% preprocessing of inputs
switch (opt.Transform)

    case {'TranslationTransform', 'BSplineTransform'}

        % compute "tm" from "tp"
        if (nargin < 3 || isempty(tm))
            tm = -tp;
        end

        % add dummy transforms at the beginning of "tm" and at the end of
        % "tp". This makes it easier to operate, because then:
        % * tp(I) and tm(I) are the neighbours of I for intermediate slices
        % * extreme slices can be dealt with in the same way as
        %   intermediate slices
        tp = tp([1:end end], :);
        tm = tm([1 1:end], :);
        tp(end, :) = tm(end, :);
        tm(1, :) = tp(1, :);

    case {'EulerTransform', 'AffineTransform'}

        % compute "tm" from "tp"
        if (nargin < 3 || isempty(tm))
            for I = 1:size(tp, 3)
                tm(:, :, I) = inv(tp(:, :, I));
            end
        end

        % add dummy transforms at the beginning of "tm" and at the end of
        % "tp". This makes it easier to operate (see note above)
        tp = tp(:, :, [1:end end]);
        tm = tm(:, :, [1 1:end]);
        tp(:, :, end) = tm(:, :, end);
        tm(:, :, 1) = tp(:, :, 1);

end

% check inputs
if any(size(tp) ~= size(tm))
    error('TP and TM must have the same size')
end

% if no transforms provided, then we don't need to run the algorithm
if (isempty(tp))
    tout = [];
    return
end

% apply registration diffusion
switch (opt.Transform)

    case {'TranslationTransform', 'BSplineTransform'}

        [tout, info] = translation_diffusion(opt, tp, tm);

    case {'EulerTransform', 'AffineTransform'}

        [tout, info] = affine_diffusion(opt, tp, tm);

    otherwise

        error('Transform not implemented')

end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% LOCAL FUNCTIONS

%% convert transform from elastix format to array, if necessary
function [tout, t0] = elastix2mat(t)

if (isstruct(t))

    % save copy of input so that we can convert the output back to elastix
    % format
    t0 = t;

    switch (t(1).Transform)

        case {'TranslationTransform', 'BSplineTransform'}

            tout = cat(1, t(:).TransformParameters);

        % TODO: AffineTransform

        otherwise

            error('Conversion from this elastix transform to Matlab array not implemented')
    end

else % already in array format

    % no need to save a copy of the input
    t0 = [];

    % output is just the input
    tout = t;

end

end

%% translation_diffusion: Apply diffusion to translation transforms
function [tout, info] = translation_diffusion(opt, tp, tm)

% init accumulated transform to apply to each slice
tout = zeros(size(tp));

I = 0;
while (1)

    % iteration number
    I = I + 1;

    % transform to apply to each slice in this iteration (transform update)
    t = opt.Alpha * (tp + tm);

    % stopping criterion: B-spline injectivity
    if (strcmp(opt.Transform, 'BSplineTransform'))

        % find control points that don't guarantee injectivity and
        % constrain them so that they do
        [info.NConstrained(I), aux] = ...
            check_bspline_injectivity_choi2000(tout + t, opt.tbsp);

        % apply the control point modifications to the B-spline update
        t = aux - tout;

    end

    % compose with the previous transforms to obtain a total transform
    tout = tout + t;

    % update neighbour transforms ...

    % ... of internal images
    tp(1:end-1, :) = t(2:end, :) + tp(1:end-1, :) - t(1:end-1, :);
    tm(2:end, :) = t(1:end-1, :) + tm(2:end, :) - t(2:end, :);

    %... of extreme images
    tp(end, :) = tm(end, :);
    tm(1, :) = tp(1, :);

    % stopping criterion: absolute value of each translation component
    tabs = abs(t(:));
    info.MaxAbs(I) = max(tabs);
    info.MeanAbs(I) = mean(tabs);
    info.MedAbs(I) = median(tabs);
    if (info.MaxAbs(I) <= opt.Epsilon)
        break
    end

    % stopping criterion: maximum number of iterations
    if (I == opt.MaxIter)
        break
    end

end

info.NumIter = I;

% DEBUG:
% isInjective = check_bspline_injectivity_global(tout + t, opt.tbsp);
% nnz(~isInjective)

end

%% affine_diffusion: Apply diffusion to affine transforms
function [tout, info] = affine_diffusion(opt, tp, tm)

% init accumulated transform to apply to each slice
tout = eye(size(tp(:, :, 1)));
tout = repmat(tout, 1, 1, size(tp, 3));

% allocate memory for transform update
t = zeros(size(tp));

I = 0;
while (1)

    % iteration number
    I = I + 1;

    % compute the transform to apply to each slice in this iteration
    for J = 1:size(tp, 3)

        % transform to apply to each slice in this iteration (transform
        % update)
        t(:, :, J) = real(expm(opt.Alpha ...
            * (logm(tp(:, :, J)) + logm(tm(:, :, J)))));

        % compose with the previous transforms to obtain a total transform
        tout(:, :, J) = tout(:, :, J) * t(:, :, J);

    end

    % update neighbour transforms of internal images
    for J = 2:size(tm, 3)
        tm(:, :, J) = t(:, :, J) \ tm(:, :, J) * t(:, :, J-1);
    end
    for J = 1:size(tp, 3)-1
        tp(:, :, J) = t(:, :, J) \ tp(:, :, J) * t(:, :, J+1);
    end

    % update neighbour transforms of extreme images
    tp(:, :, end) = tm(:, :, end);
    tm(:, :, 1) = tp(:, :, 1);

    % stopping criteria...
    % ... norm of the transform update different from identity
    tnorm = zeros(1, size(t, 3));
    for J = 1:size(t, 3)
        tnorm(J) = norm(t(:, :, J) - eye(size(t(:, :, 1))), 2);
    end
    info.MaxTNorm(I) = max(tnorm);
    info.MeanTNorm(I) = mean(tnorm);
    info.MedTNorm(I) = median(tnorm);
    if (all(tnorm <= opt.Epsilon))
        break
    end

    % ... maximum number of iterations
    if (I == opt.MaxIter)
        break
    end

end

info.NumIter = I;

end

%% check_bspline_injectivity_global:
%
% Check B-spline global injectivity triangulating the grid and looking for
% flipped triangles.
%
% tout: (N, K)-matrix, N number of slices, K number of B-spline control
%       points
% tbsp: elastix struct with the B-spline info (we need the grid spacing)
function isInjective = check_bspline_injectivity_global(tout, tbsp)

DEBUG = false;

% number of slices
N = size(tout, 1);

% grid of control points for 0 deformation
tbsp.TransformParameters(:) = 0;
[gx0, gy0] = elastix_bspline_grid(tbsp);
tri = delaunay(gx0(:), gy0(:));

% area of the triangles
atri = trifacet_signed_area(tri, [gx0(:), gy0(:)]);

% flip triangles with negative area, so that all are positive
idx = atri < 0;
tri(idx, :) = tri(idx, end:-1:1);

% DEBUG: plot control points
if (DEBUG)
    trimesh(tri, gx0(:), gy0(:))
end

isInjective = true(1, N);
for J = 1:N

    % transfer coefficients to the B-spline
    tbsp.TransformParameters = tout(J, :);

    % apply transformation to the control points (note:
    aux = transformix_pts(tbsp, [gx0(:) gy0(:)]);
    gx1 = reshape(aux(:, 1), size(gx0, 1), size(gx0, 2));
    gy1 = reshape(aux(:, 2), size(gy0, 1), size(gy0, 2));

    % DEBUG: plot control points
    if (DEBUG)
        trimesh(tri, gx1(:), gy1(:))
    end

    % area of the triangles
    atri = trifacet_signed_area(tri, [gx1(:), gy1(:)]);

    % if any triangle is flipped, the B-spline is not injective
    isInjective(J) = all(atri > 0);

end

end

%% check_bspline_injectivity_choi2000:
%
% Check B-spline injectivity using sufficient conditions in Choi and Lee
% (2000). Constrain control points to guarantee injectivity.
%
% This function checks Choi and Lee (2000) conditions for injectivity of
% the B-spline. The conditions are sufficient, i.e. if any is fulfilled,
% then we know the B-spline is injective. If not, the B-spline may or may
% not be injective, but we'll assume it's not.
%
% tout: (N, K)-matrix, N number of slices, K number of B-spline control
%       points
% tbsp: elastix struct with the B-spline info (we need the grid spacing)
%
% nConstrained: number of control points that had to be constrained to
%               guarantee injectivity
function [nConstrained, tout] = check_bspline_injectivity_choi2000(tout, tbsp)

% constants from Choi and Lee (2000), to check violations
K2 = 2.046392675;
A2 = sqrt((3/2)^2 + (K2 - 3/2)^2);

% number of slices
N = size(tout, 1);

% Note: the grid has row->x, col->y coordinates, instead of the Matlab
% convention, which is the opposite

% number of B-spline control points
L = length(tbsp.TransformParameters)/2;
Nx = tbsp.GridSize(1);
Ny = tbsp.GridSize(2);

% loop slices
isInjective = false(Nx, Ny, N);
for J = 1:N

    %% find control points that are guaranteed to not cause injectivity
    %% problems

    % transfer coefficients to the B-spline
    tbsp.TransformParameters = tout(J, :);

    % easy nomenclature for x, y coordinates of control points
    cx = tbsp.TransformParameters(1:L);
    cy = tbsp.TransformParameters(L+1:end);

    % reshape coefficients into grid
    cx = reshape(cx, Nx, Ny);
    cy = reshape(cy, Nx, Ny);

    % normalize control point displacement to grid spacing = 1
    cx = cx / tbsp.GridSpacing(1);
    cy = cy / tbsp.GridSpacing(2);

    % Theorem 1 from Choi and Lee (2000): first sufficient condition
    % for injectivity
    cond1 = (abs(cx) < 1 / K2) & (abs(cy) < (1 / K2));

    % Theorem 2 from Choi and Lee (2000): second sufficient condition
    % for injectivity
    cond2 = cx.^2 + cy.^2 < (1 / A2)^2;

    % if either condition is fulfilled, then the B-spline is injective
    isInjective(:, :, J) = cond1 | cond2;

    %% correct control points that potential can cause non-injectivity

    % find problematic points (linear index easier than row/column
    % index)
    idx = find(~isInjective(:, :, J));

    % correction factor for the problematic coefficients
    S = 1 ./ sqrt(cx(idx).^2 + cy(idx).^2) / (A2 - 0.01);

    % correct problematic control points
    cx(idx) = cx(idx) .* S;
    cy(idx) = cy(idx) .* S;

    % transfer corrected coefficients to the output
    aux = tout(J, :);
    aux(idx) = tbsp.GridSpacing(1)*cx(idx);
    aux(L + idx) = tbsp.GridSpacing(2)*cy(idx);
    tout(J, :) = aux;

end

% count number of control points that had to be constrained to guarantee
% they are injective
nConstrained = nnz(~isInjective);

end

%% check_bspline_injectivity_chun2009:
%
% Check B-spline injectivity using sufficient conditions in Chun and
% Fessler (2009).
%
% This function checks Chun and Fessler (2009) conditions for injectivity
% of the B-spline. If all the conditions are fulfilled, then we know the
% B-spline is injective. If not, the B-spline may or may not be injective,
% but we'll assume it's not.
%
% Because the conditions check grid edge lengths, it's not trivial to
% associate the conditions to vertices.
%
% tout: (N, K)-matrix, N number of slices, K number of B-spline control
%       points
% tbsp: elastix struct with the B-spline info (we need the grid spacing)
function isInjective = check_bspline_injectivity_chun2009(tout, tbsp)

DEBUG = false;

% constants from Chun and Fessler (2009), to check for violations
kx = 0.5 - 0.01;
ky = 0.5 - 0.01;

% number of slices
N = size(tout, 1);

% number of B-spline coefficients
L = length(tbsp.TransformParameters);
Nx = tbsp.GridSize(1);
Ny = tbsp.GridSize(2);

% DEBUG: plot control points
if (DEBUG)
    % grid of control points for 0 deformation
    tbsp.TransformParameters(:) = 0;
    [gx0, gy0] = elastix_bspline_grid(tbsp);
    tri = delaunay(gx0(:), gy0(:));

    trimesh(tri, gx0(:), gy0(:))
end

% loop slices
isInjective = true;
for J = 1:N

    % transfer coefficients to the B-spline
    tbsp.TransformParameters = tout(J, :);

    % easy nomenclature for x, y coordinates of control points
    cx = tbsp.TransformParameters(1:L/2);
    cy = tbsp.TransformParameters(L/2+1:end);

    % reshape coefficients into grid
    cx = reshape(cx, Nx, Ny);
    cy = reshape(cy, Nx, Ny);

    % transpose grid so that we have x->cols, y->rows
    cx = cx';
    cy = cy';

    % grid spacing
    mx = tbsp.GridSpacing(1);
    my = tbsp.GridSpacing(2);

    % conditions (i->x, j->y):

    % -mx * kx <= c_{i+1,j}^x - c_{i,j}^x
    % we also add a column to the right to account for the lost column
    cond = (diff(cx, 1, 2) >= -mx * kx);
    isInjective = isInjective && all(cond(:));

    % -my * ky <= c_{i,j+1}^y - c_{i,j}^y
    % we also add a row to the bottom to account for the lost row
    cond = (diff(cy, 1, 1) >= -my * ky);
    isInjective = isInjective && all(cond(:));

    % |c_{i,j+1}^x - c_{i,j}^x| <= mx * kx
    % we also add a row to the bottom to account for the lost row
    cond = (abs(diff(cx, 1, 1)) <= mx * kx);
    isInjective = isInjective && all(cond(:));

    % |c_{i+1,j}^y - c_{i,j}^y| <= my * ky
    % we also add a column to the right to account for the lost column
    cond = (abs(diff(cy, 1, 2)) <= my * ky);
    isInjective = isInjective && all(cond(:));

end

end


function y = vl_nnloss_modified(x, c, varargin)
%VL_NNLOSS CNN categorical or attribute loss.
%   Y = VL_NNLOSS(X, C) computes the loss incurred by the prediction
%   scores X given the categorical labels C.
%
%   The prediction scores X are organised as a field of prediction
%   vectors, represented by a H x W x D x N array. The first two
%   dimensions, H and W, are spatial and correspond to the height and
%   width of the field; the third dimension D is the number of
%   categories or classes; finally, the dimension N is the number of
%   data items (images) packed in the array.
%
%   While often one has H = W = 1, the case W, H > 1 is useful in
%   dense labelling problems such as image segmentation. In the latter
%   case, the loss is summed across pixels (contributions can be
%   weighed using the `InstanceWeights` option described below).
%
%   The array C contains the categorical labels. In the simplest case,
%   C is an array of integers in the range [1, D] with N elements
%   specifying one label for each of the N images. If H, W > 1, the
%   same label is implicitly applied to all spatial locations.
%
%   In the second form, C has dimension H x W x 1 x N and specifies a
%   categorical label for each spatial location.
%
%   In the third form, C has dimension H x W x D x N and specifies
%   attributes rather than categories. Here elements in C are either
%   +1 or -1 and C, where +1 denotes that an attribute is present and
%   -1 that it is not. The key difference is that multiple attributes
%   can be active at the same time, while categories are mutually
%   exclusive. By default, the loss is *summed* across attributes
%   (unless otherwise specified using the `InstanceWeights` option
%   described below).
%
%   DZDX = VL_NNLOSS(X, C, DZDY) computes the derivative of the block
%   projected onto the output derivative DZDY. DZDX and DZDY have the
%   same dimensions as X and Y respectively.
%
%   VL_NNLOSS() supports several loss functions, which can be selected
%   by using the option `type` described below. When each scalar c in
%   C is interpreted as a categorical label (first two forms above),
%   the following losses can be used:
%
%   Classification error:: `classerror`
%     L(X,c) = (argmax_q X(q) ~= c). Note that the classification
%     error derivative is flat; therefore this loss is useful for
%     assessment, but not for training a model.
%
%   Top-K classification error:: `topkerror`
%     L(X,c) = (rank X(c) in X <= K). The top rank is the one with
%     highest score. For K=1, this is the same as the
%     classification error. K is controlled by the `topK` option.
%
%   Log loss:: `log`
%     L(X,c) = - log(X(c)). This function assumes that X(c) is the
%     predicted probability of class c (hence the vector X must be non
%     negative and sum to one).
%
%   Softmax log loss (multinomial logistic loss):: `softmaxlog`
%     L(X,c) = - log(P(c)) where P(c) = exp(X(c)) / sum_q exp(X(q)).
%     This is the same as the `log` loss, but renormalizes the
%     predictions using the softmax function.
%
%   Multiclass hinge loss:: `mhinge`
%     L(X,c) = max{0, 1 - X(c)}. This function assumes that X(c) is
%     the score margin for class c against the other classes.  See
%     also the `mmhinge` loss below.
%
%   Multiclass structured hinge loss:: `mshinge`
%     L(X,c) = max{0, 1 - M(c)} where M(c) = X(c) - max_{q ~= c}
%     X(q). This is the same as the `mhinge` loss, but computes the
%     margin between the prediction scores first. This is also known
%     the Crammer-Singer loss, an example of a structured prediction
%     loss.
%
%   When C is a vector of binary attribures c in (+1,-1), each scalar
%   prediction score x is interpreted as voting for the presence or
%   absence of a particular attribute. The following losses can be
%   used:
%
%   Binary classification error:: `binaryerror`
%     L(x,c) = (sign(x - t) ~= c). t is a threshold that can be
%     specified using the `threshold` option and defaults to zero. If
%     x is a probability, it should be set to 0.5.
%
%   Binary log loss:: `binarylog`
%     L(x,c) = - log(c(x-0.5) + 0.5). x is assumed to be the
%     probability that the attribute is active (c=+1). Hence x must be
%     a number in the range [0,1]. This is the binary version of the
%     `log` loss.
%
%   Logistic log loss:: `logistic`
%     L(x,c) = log(1 + exp(- cx)). This is the same as the `binarylog`
%     loss, but implicitly normalizes the score x into a probability
%     using the logistic (sigmoid) function: p = sigmoid(x) = 1 / (1 +
%     exp(-x)). This is also equivalent to `softmaxlog` loss where
%     class c=+1 is assigned score x and class c=-1 is assigned score
%     0.
%
%   Hinge loss:: `hinge`
%     L(x,c) = max{0, 1 - cx}. This is the standard hinge loss for
%     binary classification. This is equivalent to the `mshinge` loss
%     if class c=+1 is assigned score x and class c=-1 is assigned
%     score 0.
%
%   VL_NNLOSS(...,'OPT', VALUE, ...) supports these additionals
%   options:
%
%   InstanceWeights:: []
%     Allows to weight the loss as L'(x,c) = WGT L(x,c), where WGT is
%     a per-instance weight extracted from the array
%     `InstanceWeights`. For categorical losses, this is either a H x
%     W x 1 or a H x W x 1 x N array. For attribute losses, this is
%     either a H x W x D or a H x W x D x N array.
%
%   TopK:: 5
%     Top-K value for the top-K error. Note that K should not
%     exceed the number of labels.
%
%   See also: VL_NNSOFTMAX().

% Copyright (C) 2014-15 Andrea Vedaldi.
% Copyright (C) 2016 Karel Lenc.
% All rights reserved.
%
% This file is part of the VLFeat library and is made available under
% the terms of the BSD license (see the COPYING file).

if ~isempty(varargin) && ~ischar(varargin{1})  % passed in dzdy
    dzdy = varargin{1} ;
    varargin(1) = [] ;
else
    dzdy = [] ;
end

opts.marginMat = [] ;
opts.marginAlpha_ = 1 ;
opts.instanceWeights = [] ;
opts.classWeights = [] ;
opts.threshold = 0 ;
opts.loss = 'softmaxlog' ;
opts.topK = 5 ;
opts = vl_argparse(opts, varargin, 'nonrecursive') ;

inputSize = [size(x,1) size(x,2) size(x,3) size(x,4)] ;

% Form 1: C has one label per image. In this case, get C in form 2 or
% form 3.
c = gather(c) ;
if numel(c) == inputSize(4)
    c = reshape(c, [1 1 1 inputSize(4)]) ;
    c = repmat(c, inputSize(1:2)) ;
end

switch lower(opts.loss)
    case {'cosinesimilaritylogloss', 'cosinesimilaritymmloss', 'cosinesimilarityregloss', 'cosinesimilarityregloss_full', ...
            'cosinesimlaritylogloss', 'cosinesimlaritymmloss', 'cosinesimlarityregloss', 'cosinesimlarityregloss_full', ...
           'cosinesimlairtylogloss', 'cosinesimlairtymmloss', 'cosinesimlairtyregloss', 'cosinesimlairtyregloss_full', ...
           'cosinesimilarityabsregloss', 'cosinesimilarityabsregloss_full'}
        hasIgnoreLabel = 0;
    otherwise
        hasIgnoreLabel = any(c(:) == 0);
end

% --------------------------------------------------------------------
% Spatial weighting
% --------------------------------------------------------------------

% work around a bug in MATLAB, where native cast() would slow
% progressively
if isa(x, 'gpuArray')
    switch classUnderlying(x) ;
        case 'single', cast = @(z) single(z) ;
        case 'double', cast = @(z) double(z) ;
    end
    c = gpuArray(c);
else
    switch class(x)
        case 'single', cast = @(z) single(z) ;
        case 'double', cast = @(z) double(z) ;
    end
end

labelSize = [size(c,1) size(c,2) size(c,3) size(c,4)] ;
assert(isequal(labelSize(1:2), inputSize(1:2))) ;
assert(labelSize(4) == inputSize(4)) ;
instanceWeights = [] ;
switch lower(opts.loss)
    case {'classerror', 'topkerror', 'log', 'softmaxlog', 'mhinge', 'mshinge'}
        % there must be one categorical label per prediction vector
        assert(labelSize(3) == 1) ;

        if hasIgnoreLabel
            % null labels denote instances that should be skipped
            instanceWeights = cast(c(:,:,1,:) ~= 0) ;
        end

    case {'binaryerror', 'binarylog', 'logistic', 'hinge'}

        % there must be one categorical label per prediction scalar
        assert(labelSize(3) == inputSize(3)) ;

        if hasIgnoreLabel
            % null labels denote instances that should be skipped
            instanceWeights = cast(c ~= 0) ;
        end
    case {'cosinesimilaritylogloss', 'cosinesimilaritymmloss', 'cosinesimilarityregloss', 'cosinesimilarityregloss_full', ...
           'cosinesimlaritylogloss', 'cosinesimlaritymmloss', 'cosinesimlarityregloss', 'cosinesimlarityregloss_full', ...
           'cosinesimlairtylogloss', 'cosinesimlairtymmloss', 'cosinesimlairtyregloss', 'cosinesimlairtyregloss_full',...
           'cosinesimilarityabsregloss', 'cosinesimilarityabsregloss_full'}
        assert(labelSize(1) == inputSize(1)) ;
        assert(labelSize(2) == inputSize(2)) ;
        assert(labelSize(3) == inputSize(3)) ;
        assert(labelSize(4) == inputSize(4)) ;
        instanceWeights = opts.instanceWeights;
    otherwise
        error('Unknown loss ''%s''.', opts.loss) ;
end

if ~isempty(opts.instanceWeights)
    % important: this code needs to broadcast opts.instanceWeights to
    % an array of the same size as c
    if isempty(instanceWeights) && isempty(strfind(lower(opts.loss), 'cosinesimilarity'))  && isempty(strfind(lower(opts.loss), 'cosinesimlarity'))
        instanceWeights = bsxfun(@times, onesLike(c), opts.instanceWeights) ;
%         instanceWeights = c;
%         instanceWeights(instanceWeights~=11) = 3;
%         instanceWeights(instanceWeights==11) = 1;
%         instanceWeights = bsxfun(@times, instanceWeights, opts.instanceWeights) ;
    elseif isempty(instanceWeights) && ~isempty(strfind(lower(opts.loss), 'cosinesimilarity'))  && ~isempty(strfind(lower(opts.loss), 'cosinesimlarity'))
        instanceWeights = bsxfun(@times, instanceWeights, opts.instanceWeights);
    else
    end
end

% --------------------------------------------------------------------
% Do the work
% --------------------------------------------------------------------

switch lower(opts.loss)
    case {'log', 'softmaxlog', 'mhinge', 'mshinge'}
        % from category labels to indexes
        numPixelsPerImage = prod(inputSize(1:2)) ;
        numPixels = numPixelsPerImage * inputSize(4) ;
        imageVolume = numPixelsPerImage * inputSize(3) ;

        n = reshape(0:numPixels-1,labelSize) ;
        offset = 1 + mod(n, numPixelsPerImage) + ...
            imageVolume * fix(n / numPixelsPerImage) ;
        ci = offset + numPixelsPerImage * max(c - 1,0) ;
end

if nargin <= 2 || isempty(dzdy)
    switch lower(opts.loss)
        case 'classerror'
            [~,chat] = max(x,[],3) ;
            t = cast(c ~= chat) ;
        case 'topkerror'
            [~,predictions] = sort(x,3,'descend') ;
            t = 1 - sum(bsxfun(@eq, c, predictions(:,:,1:opts.topK,:)), 3) ;
        case 'log'
            t = - log(x(ci)) ;
        case 'softmaxlog'
            Xmax = max(x,[],3) ;
            ex = exp(bsxfun(@minus, x, Xmax)) ;
            t = Xmax + log(sum(ex,3)) - x(ci) ;
        case 'mhinge'
            t = max(0, 1 - x(ci)) ;
        case 'mshinge'
            Q = x ;
            Q(ci) = -inf ;
            t = max(0, 1 - x(ci) + max(Q,[],3)) ;
        case 'binaryerror'
            t = cast(sign(x - opts.threshold) ~= c) ;
        case 'binarylog'
            t = -log(c.*(x-0.5) + 0.5) ;
        case 'logistic'
            %t = log(1 + exp(-c.*X)) ;
            a = -c.*x ;
            b = max(0, a) ;
            t = b + log(exp(-b) + exp(a-b)) ;
        case 'hinge'
            t = max(0, 1 - c.*x) ;
        case 'cosinesimilaritylogloss'
            t = - (log(x).*c + log(1-x).*(1-c));
        case {'cosinesimilaritymmloss', 'cosinesimlaritymmloss', 'cosinesimlairtymmloss'}
            if isempty(opts.marginMat)
                t = max(0, x-opts.marginAlpha_) .* (1-c) ;
            else
                t = max(0, x-opts.marginMat) .* (1-c) ;
            end
        case {'cosinesimilarityregloss', 'cosinesimlarityregloss', 'cosinesimlairtyregloss'}
            t = c.* ((x-c).^2); % including positive pairs only
        case {'cosinesimilarityregloss_full', 'cosinesimlarityregloss_full', 'cosinesimlairtyregloss_full'}
            t = ((x-c).^2); % including positive and negative pairs
        case {'cosinesimilarityabsregloss', 'cosinesimlarityabsregloss', 'cosinesimlairtyabsregloss'}
            t = c.* (c-x); % including positive pairs only
        case {'cosinesimilarityabsregloss_full', 'cosinesimlarityabsregloss_full', 'cosinesimlairtyabsregloss_full'}
            t =  (c-x); % including positive pairs only
    end
    if ~isempty(instanceWeights)
        y = instanceWeights(:)' * t(:) ;
    else
        y = sum(t(:));
    end
else
    if ~isempty(instanceWeights)
        dzdy = dzdy * instanceWeights ;
    end
    switch lower(opts.loss)
        case {'classerror', 'topkerror'}
            y = zerosLike(x) ;
        case 'log'
            y = zerosLike(x) ;
            y(ci) = - dzdy ./ max(x(ci), 1e-8) ;
        case 'softmaxlog'
            Xmax = max(x,[],3) ;
            ex = exp(bsxfun(@minus, x, Xmax)) ;
            y = bsxfun(@rdivide, ex, sum(ex,3)) ;
            y(ci) = y(ci) - 1 ;
            y = bsxfun(@times, dzdy, y) ;
        case 'mhinge'
            % t = max(0, 1 - x(ci)) ;
            y = zerosLike(x) ;
            y(ci) = - dzdy .* (x(ci) < 1) ;
        case 'mshinge'
            Q = x ;
            Q(ci) = -inf ;
            [~, q] = max(Q,[],3) ;
            qi = offset + numPixelsPerImage * (q - 1) ;
            W = dzdy .* (x(ci) - x(qi) < 1) ;
            y = zerosLike(x) ;
            y(ci) = - W ;
            y(qi) = + W ;
        case 'binaryerror'
            y = zerosLike(x) ;
        case 'binarylog'
            y = - dzdy ./ (x + (c-1)*0.5) ;
        case 'logistic'
            % t = exp(-Y.*X) / (1 + exp(-Y.*X)) .* (-Y)
            % t = 1 / (1 + exp(Y.*X)) .* (-Y)
            y = - dzdy .* c ./ (1 + exp(c.*x)) ;
        case 'hinge'
            y = - dzdy .* c .* (c.*x < 1) ;
        case {'cosinesimilaritylogloss', 'cosinesimlaritylogloss', 'cosinesimlairtylogloss'}
            y = -dzdy .* ( c./x - (1-c) ./ (1-x) );
        case {'cosinesimilaritymmloss', 'cosinesimlaritymmloss', 'cosinesimlairtymmloss'}

            if isempty(opts.marginMat)
                % t = max(0, x-opts.marginAlpha_) .* (1-c) ;
                y = x - opts.marginAlpha_;
                y(y<0) = 0;
                y = y.* (1-c);
                y = dzdy .* y;
            else
                y = x - opts.marginMat;
                y(y<0) = 0;
                y = y.* (1-c);
                y = dzdy .* y;
            end
        case {'cosinesimilarityregloss', 'cosinesimlarityregloss', 'cosinesimlairtyregloss'}
            y = c.* dzdy .* (x-c);     % including positive pairs only
        case {'cosinesimilarityregloss_full', 'cosinesimlarityregloss_full', 'cosinesimlairtyregloss_full'}
            y = dzdy .* (x-c); % including positive and negative pairs
        case {'cosinesimilarityabsregloss', 'cosinesimlarityabsregloss', 'cosinesimlairtyabsregloss'}
            %t = c.* (c-x); % including positive pairs only
            y = c.* dzdy .* (-x);     % including positive pairs only
        case {'cosinesimilarityabsregloss_full', 'cosinesimlarityabsregloss_full', 'cosinesimlairtyabsregloss_full'}
            %t =  (c-x); % including positive pairs only
            y = dzdy .* (-x);     % including positive pairs only
    end
end

% --------------------------------------------------------------------
function y = zerosLike(x)
% --------------------------------------------------------------------
if isa(x,'gpuArray')
    y = gpuArray.zeros(size(x),classUnderlying(x)) ;
else
    y = zeros(size(x),'like',x) ;
end

% --------------------------------------------------------------------
function y = onesLike(x)
% --------------------------------------------------------------------
if isa(x,'gpuArray')
    y = gpuArray.ones(size(x),classUnderlying(x)) ;
else
    y = ones(size(x),'like',x) ;
end


classdef chebop2
%CHEBOP2   CHEBOP2 class for representing linear partial differential operators.
%
% Class used to solve linear PDEs defined on rectangular domains that have unique and
% globally smooth solutions.
%
% N = CHEBOP2(@(u) op(u)) constructs an operator N representing the operator
% given in @(u)op(u) acting on functions of two variables on [-1,1] by [-1,1].
%
% N = CHEBOP2(@(u) op(u), [a b c d]) constructs an operator N acting on
% functions of two variables defined on [a,b] by [c,d].
%
% N = CHEBOP2(@(x,y,u) op(x,y,u),...) constructs a variable coefficient PDE
% operator. If a partial differential operator is constant coefficient, we
% recommend the @(u) notation rather than @(x,y,u) as it is more efficient.
%
% Boundary conditions are imposed via the syntax N.lbc, N.rbc, N.ubc, and N.dbc.
%
% Example 1: (Poisson with zero Dirichlet conditions):
%    N = chebop2(@(u) diff(u,2,1) + diff(u,2,2));
%    N.bc = 0;
%    u = N \ 1;
%
% Example 2: (Helmholtz equation with gravity)
%    N = chebop2(@(x,y,u) laplacian(u) - 10*y.^2.*u, [-1 1 -3 0]);
%    N.bc = 1;
%    u = N \ 0;
%
% Example 3: (Klein-Gordon equation)
%    N = chebop2(@(u) diff(u,2,1) - diff(u,2,2) + 5*u,[-1 1 0 3]);
%    N.lbc = 0; N.rbc = 0;
%    N.dbc = @(x,u) [u - exp(-30*x.^2) ; diff(u)];
%    u = N \ 0;
%
% Example 4: (Poisson equation with general Dirichlet conditions and rhs)
%    N = chebop2( @(u) lap(u));
%    N.lbc = @(y) -y.^2;
%    N.rbc = @(y) y.^2;
%    N.ubc = @(x) x;
%    N.dbc = @(x) x;
%    u = N \ chebfun2(@(x,y) cos(x.*y));
%
% For further details about the PDE solver, see:
%
% A. Townsend and S. Olver, The automatic solution of partial differential
% equations using a global spectral method, J. Comput. Phys., 299 (2015),
% pp. 106-123.
%
% Warning: This PDE solver is an experimental new feature. It has not been
% publicly advertised.  Chebop2 cannot do nonlinear problems as more
% algorithmic advances are needed.

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

    %% PROPERTIES.
    properties ( GetAccess = 'public', SetAccess = 'public' )

        domain = [];  % Domain of the operator.
        op = [];      % The operator.
        ubc = [];     % Up boundary condition(s).
        lbc = [];     % Left boundary condition(s).
        rbc = [];     % Right boundary condition(s).
        dbc = [];     % Down boundary condition(s).
        dim = [];     % Size of the system (number of eqns).
        scale = [];   % Relative solution scale.
        coeffs = [];  % Matrix storing constant coefficients.
        rhs = [];     % Righthand side, if given by user.
        xorder = 0;   % Diff order in the x-variable.
        yorder = 0;   % Diff order in the y-variable.
        U             %
        S             % Low rank form of the partial differential operator.
        V             %

    end

    %% CONSTRUCTOR.
    methods

        function N = chebop2(varargin)
            % CHEBOP2 CONSTRUCTOR.

            % Get CHEBFUN2 preferences.
            pref = chebfunpref();
            tol = pref.cheb2Prefs.chebfun2eps;

            % If empty input arguments then return an empty CHEBOP2 object.
            if ( isempty(varargin) )
                return
            end

            % What domain is the operator defined on?
            if ( numel(varargin) > 1 )
                ends = varargin{2}; % Second argument should be a domain.
                if ( length(ends) == 4 )
                    % Valid domain?
                    if ( diff(ends(1:2)) > 0 && diff(ends(3:4)) > 0 )
                        dom = ends;
                    else
                        error('CHEBFUN:CHEBOP2:chebop2:emptyDomain', ...
                            'Empty domain.');
                    end
                else
                    error('CHEBFUN:CHEBOP2:chebop2:badDomain',...
                        'Argument should be a domain given by four doubles.')
                end
            else
                if ( isa(varargin{1}, 'function_handle') )
                    % Pick the default domain.
                    rect1 = [-1, 1];
                    rect2 = [-1, 1];
                    dom = [rect1, rect2];
                elseif ( isa( varargin{1}, 'double') )
                    % Set up identity operator on the domain.
                    N = chebop2(@(u) u, varargin{1});
                    return
                else
                    error('CHEBFUN:CHEBOP2:chebop2:badArg',...
                        'First argument is not an operator or domain.')
                end
            end

            % First argument in the constructor is the operator. If the
            % operator is univariate then it's a constant coefficient PDE,
            % otherwise assume it is a variable coefficient.
            if ( isa(varargin{1}, 'function_handle') )
                fh = varargin{1};

                if ( nargin(fh) == 1 )  % The PDE has constant coefficients.
                    % Trust that the user has formed the CHEBFUN2 objects
                    % outside of CHEBOP2.

                    % Extract out rhs:
                    x = chebfun2(@(x,y) x, dom);
                    RHS = fh(0*x);
                    % In this case, the RHS must be a constant chebfun2.
                    if ( norm( RHS ) > 0 )
                        fh = @(u) fh(u) - RHS;
                        N.rhs = -RHS;             % store for later.
                    else
                        N.rhs = 0;
                    end

                    u = adchebfun2(chebfun2(@(x,y) x.*y, dom));
                    v = fh(u);
                    % If the PDO has constant coefficients then convert to
                    % double:
                    try
                        A = cell2mat(v.jacobian).';
                    catch
                        % PDO has variable coefficients, keep them in a
                        % cell array:
                        A = v.jacobian;
                    end

                elseif ( nargin(fh) == 2 )
                    error('CHEBFUN:CHEBOP2:chebop2:badOp1', ...
                        'Did you intend to have @(x,y,u)?')
                elseif ( nargin(fh) == 3 )
                    % The coefficients of the PDE are now variable
                    % coefficient.

                    % Setup a chebfun2 on the right domain
                    u = adchebfun2(chebfun2(@(x,y) x.*y, dom));
                    x = chebfun2(@(x,y) x, dom);
                    y = chebfun2(@(x,y) y, dom);

                    % Extract out rhs:
                    RHS = fh(x, y, 0*x);
                    % The RHS is a chebfun2. Check if it is the zero. Do
                    % not change function handle unless we have to because
                    % we want the user to be able to see the PDO
                    % untouched if possible.
                    if ( norm( RHS ) > 0 )
                        fh = @(x, y, u) fh(x, y, u) - RHS;
                        N.rhs = -RHS;             % store for later.
                    else
                        N.rhs = 0;
                    end

                    % Apply it to the operator.
                    v = fh(x, y, u);
                    A = v.jacobian;  % Cell array of variable coefficients.

                    % If we have a variable coefficient PDO, then compute the
                    % separable representation immediately. We need it now.
                    [cellU, matS, cellV] = chebop2.separableFormat( A,...
                                                    size(A,2), size(A,1), dom );
                    N.U = cellU;
                    N.S = matS;
                    N.V = cellV;
                else
                    error('CHEBFUN:CHEBOP2:chebop2:badOp2',...
                        'Operator should be @(u) or @(x,y,u).')
                end

            else
                error('CHEBFUN:CHEBOP2:chebop2:badOp3',...
                    'First argument should be an operator')
            end

            % Often the coefficients are obtained with small rounding errors
            % and it is important to remove the very small non-zero ones to
            % have rank(A) correct.
            if ( iscell(A) )
                for jj = size(A, 1)
                    for kk = size(A, 2)
                        if ( isa(A{jj,kk}, 'double') && abs(A{jj,kk}) < 10*tol )
                            A{jj,kk} = 0;
                        end
                    end
                end
            else
                A(abs(A) < 10*tol) = 0;
            end

            % Construct CHEBOP2 object. The boundary conditions will be
            % given later.
            N.domain = dom;
            N.op = fh;
            N.coeffs = A;

            % Calculate xorder and yorder of PDE.
            % Find the differential order of the PDE operator.
            if ( iscell(A) )
                xdifforder = size(A, 2) - 1;
                ydifforder = size(A, 1) - 1;
            elseif ( min(size(A)) > 1 )
                xdifforder = find(sum(abs(A), 2) > 100*tol, 1, 'last') - 1;
                ydifforder = find(sum(abs(A)) > 100*tol, 1, 'last' ) - 1;
            else
                if ( size(A, 1) == 1 )
                    ydifforder = length(A) - 1;
                    xdifforder = 0;
                else
                    xdifforder = length(A) - 1;
                    ydifforder = 0;
                end
            end
            N.xorder = xdifforder;
            N.yorder = ydifforder;

            % Issue a warning to the user for the first CHEBOP2:
            warning('CHEBFUN:CHEBOP2:chebop2:experimental',...
                ['CHEBOP2 is a new experimental feature.\n'...
                'It has not been tested to the same extent as other'...
                'parts of the software.']);
            % Turn it off:
            warning('off', 'CHEBFUN:CHEBOP2:chebop2:experimental');

        end

    end

    %% STATIC HIDDEN METHODS.
    methods ( Static = true, Hidden = true )

        % Matrix equation solver for AX - XB = F (p and q are ADI shifts):
        X = adi( A, B, F, p, q );

        % Matrix equation solver for A*X-X*B = M*N.' (p and q are ADI shifts):
        [UX, DX, VX] = fadi( A, B, M, N, p, q );

        % ADI shift parameters for ADI method:
        [p, q] = adiShifts(a, b, c, d, tol);

        % Matrix equation solver for AXB^T + CXD^T = E. xsplit, ysplit = 1 if
        % the even and odd modes (coefficients) decouple.
        X = bartelsStewart(A, B, C, D, E, xsplit, ysplit);

        % Use automatic differentiation to pull out the coeffs of the
        % operator:
        deriv = chebfun2deriv(op);

        % This is used to discretize the linear constrains:
        [bcrow, bcvalue] = constructBC(bcArg, bcpos,...
            een, bcn, dom, scl, order);

        % Recover coefficient functions of a linear operator:
        p = recoverCoeffs(L);

        % Checks the type of the boundary conditions:
        [bctype, g] = checkBC(N, m, n);

        % This is used to discretize the PDE:
        [CC, rhs, bb, gg, Px, Py, xsplit, ysplit] =...
            discretize(N, f, m, n, flag);

        % Convert all the different user inputs for bc into uniform format:
        bc = createBC(bcArg, ends);

        % Method for deciding how to solve the matrix equation:
        X = denseSolve(N, f, m, n);

        % Compute the separable representation of a PDO:
        [cellU, S, cellV] = separableFormat(A, xorder, yorder, dom);

        % Setup Laplace operator on domain DOM:
        N = setupLaplace( dom );

        % Remove trailing coefficients.
        a = truncate(a, tol);

    end

end


function model = mogOptimise(model, display, iters)

% MOGOPTIMISE Optimise an MOG model.
% FORMAT
% DESC optimises an mixtures of Gaussians model via the
% expectation maximisation algorithm.
% ARG model : the model to be optimised.
% RETURN model : the optimised model.
%
% SEEALSO : mmpcaCreate, modelOptimise
%
% COPYRIGHT : Neil D. Lawrence, 2006

% MLTOOLS

diffll = 1;
iter = 0;
ll = mogLowerBound(model);
while abs(diffll)>1e-6 & iter<iters
  iter = iter + 1;
  model =  mogEstep(model);
  if display > 1
    [ll, oldll] = boundCheck(model, ll, 'E-step');
  end
  model = mogUpdatePrior(model);
  if display > 1
    [ll, oldll] = boundCheck(model, ll, 'Prior Update');
  end
  model = mogUpdateMean(model);
  if display > 1
    [ll, oldll] = boundCheck(model, ll, 'Mean Update');
  end
  model = mogUpdateCovariance(model);
  if display > 1
    [ll, oldll] = boundCheck(model, ll, 'Covariance Update');
  end
  if display > 1
  else
    oldll = ll;
    ll = mogLowerBound(model);
  end
  diffll = ll -oldll;
  if display
    fprintf('Iteration %d log-likelihood: %2.6f\n', iter, ll)
  end
end

function [ll, oldll] = boundCheck(model, oldll, step)

% BOUNDCHECK Helper function for checking bound.

ll = mogLowerBound(model);
diffll = ll - oldll;
if ll -oldll < 0
  warning(['Log likelihood went down by ' num2str(diffll) 'in ' ...
           'step: ' step ' in mogOptimise'])
end


function monotonicity = DeriveMonotonicityFromStationary(properties,xL,xU)

monotonicity = 'none';
if isequal(properties.convexity,'convex')
    if xU <= properties.stationary(1)
        monotonicity = 'decreasing';
    elseif xL >= properties.stationary(1)
        monotonicity = 'increasing';
    end
elseif isequal(properties.convexity,'concave')
    if xU <= properties.stationary(1)
        monotonicity = 'increasing';
    elseif xL >= properties.stationary(1)
        monotonicity = 'decreasing';
    end
end


function ans = calc_Dx(doseBinsV, volsHistV, x, volType)
% Returns the lowest dose in x% of structure'
% given the DVH data and the parameter percent.
%
%  MODIFICATION ALERT:  THIS FUNCTION IS UTILIZED BY THE DREXLER CODEBASE
%
%  Last modified: AJH 11/05
%
%  Usage: calc_Dx(doseBinsV, volsHistV, percent)
%
% Copyright 2010, Joseph O. Deasy, on behalf of the CERR development team.
%
% This file is part of The Computational Environment for Radiotherapy Research (CERR).
%
% CERR development has been led by:  Aditya Apte, Divya Khullar, James Alaly, and Joseph O. Deasy.
%
% CERR has been financially supported by the US National Institutes of Health under multiple grants.
%
% CERR is distributed under the terms of the Lesser GNU Public License.
%
%     This version of CERR 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.
%
% CERR 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 CERR.  If not, see <http://www.gnu.org/licenses/>.

if isstruct(x)  %for use with ROE
    volType = x.volType.val;
    x = x.x.val;
end

if ~exist('volType','var')
    %warning('Input volume assumed to be in percentage. Set volType=0 for absolute values.');
else
    if ~volType
    x = x/sum(volsHistV)*100;
    end
end

cumVolsV = cumsum(volsHistV);
	cumVols2V = cumVolsV(end) - cumVolsV;
	ind = min(find([ cumVols2V/cumVolsV(end) < x/100 ])); %CHANGED
	if isempty(ind)
        ans = 0;
	else
    	ans = doseBinsV(ind);
	end
return;


function [grains,grainId,mis2mean] = calcGrains(ebsd,varargin)
% grains reconstruction from 2d EBSD data
%
% Syntax
%
%   [grains, ebsd.grainId] = calcGrains(ebsd,'angle',10*degree)
%
%   % reconstruction low and high angle grain boundaries
%   lagb = 2*degree;
%   hagb = 10*degree;
%   grains = calcGrains(ebsd,'angle',[hagb lagb])
%
%   % allow grains to grow into not indexed regions
%   grains = calcGrains(ebsd('indexed'),'angle',10*degree)
%
%   % do not allow grains to grow into not indexed regions
%   grains = calcGrains(ebsd,'unitCell')
%
%   % follow non convex outer boundary
%   grains = calcGrains(ebsd,'boundary','tight')
%
%   % specify phase dependent thresholds
%   % thresholds follow the same order as ebsd.CSList and should have the same length
%   grains = calcGrains(ebsd,'angle',{angl_1 angle_2 angle_3})
%
%   % markovian clustering algorithm
%   p = 1.5;    % inflation power (default = 1.4)
%   maxIt = 10; % number of iterations (default = 4)
%   delta = 5*degree % variance of the threshold angle
%   grains = calcGrains(ebsd,'method','mcl',[p maxIt],'soft',[angle delta])
%
% Input
%  ebsd   - @EBSD
%
% Output
%  grains       - @grain2d
%  ebsd.grainId - grainId of each pixel
%
% Options
%  threshold, angle - array of threshold angles per phase of mis/disorientation in radians
%  boundary         - bounds the spatial domain ('convexhull', 'tight')
%  maxDist          - maximum distance to for two pixels to be in one grain (default inf)
%  fmc       - fast multiscale clustering method
%  mcl       - markovian clustering algorithm
%  custom    - use a custom property for grain separation
%
% Flags
%  unitCell - omit voronoi decomposition and treat a unitcell lattice
%  qhull    - use qHull for the voronin decomposition
%
% References
%
% * F.Bachmann, R. Hielscher, H. Schaeben, Grain detection from 2d and 3d
% EBSD data - Specification of the MTEX algorithm: Ultramicroscopy, 111,
% 1720-1733, 2011
%
% * C. McMahon, B. Soe, A. Loeb, A. Vemulkar, M. Ferry, L. Bassman,
%   Boundary identification in EBSD data with a generalization of fast
%   multiscale clustering, <https://doi.org/10.1016/j.ultramic.2013.04.009
%   Ultramicroscopy, 2013, 133:16-25>.
%
% See also
% GrainReconstruction GrainReconstructionAdvanced

% subdivide the domain into cells according to the measurement locations,
% i.e. by Voronoi teselation or unit cell
[V,F,I_FD] = spatialDecomposition([ebsd.prop.x(:), ebsd.prop.y(:)],ebsd.unitCell,varargin{:});
% V - list of vertices
% F - list of faces
% D - cell array of cells
% I_FD - incidence matrix faces to vertices

% determine which cells to connect
[A_Db,I_DG] = doSegmentation(I_FD,ebsd,varargin{:});
% A_db - neigbhouring cells with (inner) grain boundary
% I_DG - incidence matrix cells to grains

% compute grain ids
[grainId,~] = find(I_DG.');

% phaseId of each grain
phaseId = full(max(I_DG' * ...
  spdiags(ebsd.phaseId,0,numel(ebsd.phaseId),numel(ebsd.phaseId)),[],2));
phaseId(phaseId==0) = 1; % why this is needed?

% compute boundary this gives
% I_FDext - faces x cells external grain boundaries
% I_FDint - faces x cells internal grain boundaries
[I_FDext, I_FDint, Fext, Fint] = calcBoundary;

if check_option(varargin,'removeQuadruplePoints')
  qAdded = removeQuadruplePoints;
end

% setup grains
grains = grain2d( makeBoundary(Fext,I_FDext), ...
  calcPolygons(I_FDext * I_DG,Fext,V), ...
  [], ebsd.CSList, phaseId, ebsd.phaseMap, varargin{:});

grains.grainSize = full(sum(I_DG,1)).';
grains.innerBoundary = makeBoundary(Fint,I_FDint);
grains.scanUnit = ebsd.scanUnit;

% merge quadruple grains
if check_option(varargin,'removeQuadruplePoints') && qAdded > 0
  mergeQuadrupleGrains;
end

% calc mean orientations, GOS and mis2mean
% ----------------------------------------

[d,g] = find(I_DG);

grainRange    = [0;cumsum(grains.grainSize)];        %
firstD        = d(grainRange(2:end));
q             = quaternion(ebsd.rotations);
meanRotation  = q(firstD);
GOS = zeros(length(grains),1);

% choose between equivalent orientations in one grain such that all are
% close together
for pId = grains.indexedPhasesId
  ndx = ebsd.phaseId(d) == pId;
  if ~any(ndx), continue; end
  q(d(ndx)) = project2FundamentalRegion(q(d(ndx)),ebsd.CSList{pId},meanRotation(g(ndx)));
end

% TODO: this can be done more efficiently using accumarray

% compute mean orientation and GOS
doMeanCalc = find(grains.grainSize>1 & grains.isIndexed);
for k = 1:numel(doMeanCalc)
  qind = subSet(q,d(grainRange(doMeanCalc(k))+1:grainRange(doMeanCalc(k)+1)));
  mq = mean(qind,'robust');
  meanRotation = setSubSet(meanRotation,doMeanCalc(k),mq);
  GOS(doMeanCalc(k)) = mean(angle(mq,qind));
end

% save
grains.prop.GOS = GOS;
grains.prop.meanRotation = reshape(meanRotation,[],1);
mis2mean = inv(rotation(q(:))) .* grains.prop.meanRotation(grainId(:));

% assign variant and parent Ids for variant-based grain computation
if check_option(varargin,'variants')
    variantId = get_option(varargin,'variants');
    grains.prop.variantId = variantId(firstD,1);
    grains.prop.parentId = variantId(firstD,2);
end

  function [A_Db,I_DG] = doSegmentation(I_FD,ebsd,varargin)
    % segmentation
    %
    %
    % Output
    %  A_Db - adjecency matrix of grain boundaries
    %  A_Do - adjecency matrix inside grain connections

    % extract segmentation method
    grainBoundaryCiterions = dir([mtex_path '/EBSDAnalysis/@EBSD/private/gbc*.m']);
    grainBoundaryCiterions = {grainBoundaryCiterions.name};
    gbcFlags = regexprep(grainBoundaryCiterions,'gbc_(\w*)\.m','$1');

    gbc      = get_flag(varargin,gbcFlags,'angle');
    gbcValue = ensurecell(get_option(varargin,{gbc,'threshold','delta'},15*degree,{'double','cell'}));

    if numel(gbcValue) == 1 && length(ebsd.CSList) > 1
      gbcValue = repmat(gbcValue,size(ebsd.CSList));
    end

    % get pairs of neighbouring cells {D_l,D_r} in A_D
    A_D = I_FD'*I_FD==1;
    [Dl,Dr] = find(triu(A_D,1));

    if check_option(varargin,'maxDist')
      xyDist = sqrt((ebsd.prop.x(Dl)-ebsd.prop.x(Dr)).^2 + ...
        (ebsd.prop.y(Dl)-ebsd.prop.y(Dr)).^2);

      dx = sqrt(sum((max(ebsd.unitCell)-min(ebsd.unitCell)).^2));
      maxDist = get_option(varargin,'maxDist',3*dx);
      % maxDist = get_option(varargin,'maxDist',inf);
    else
      maxDist = 0;
    end

    connect = zeros(size(Dl));

    for p = 1:numel(ebsd.phaseMap)

      % neighboured cells Dl and Dr have the same phase
      if maxDist > 0
        ndx = ebsd.phaseId(Dl) == p & ebsd.phaseId(Dr) == p & xyDist < maxDist;
      else
        ndx = ebsd.phaseId(Dl) == p & ebsd.phaseId(Dr) == p;
      end

      connect(ndx) = true;

      % check, whether they are indexed
      ndx = ndx & ebsd.isIndexed(Dl) & ebsd.isIndexed(Dr);

      % now check for the grain boundary criterion
      if any(ndx)

        connect(ndx) = feval(['gbc_' gbc],...
          ebsd.rotations,ebsd.CSList{p},Dl(ndx),Dr(ndx),gbcValue{p},varargin{:});

      end
    end

    % adjacency of cells that have no common boundary
    ind = connect>0;
    A_Do = sparse(double(Dl(ind)),double(Dr(ind)),connect(ind),length(ebsd),length(ebsd));
    if check_option(varargin,'mcl')

      param = get_option(varargin,'mcl');
      if isempty(param), param = 1.4; end
      if length(param) == 1, param = [param,4]; end

      A_Do = mclComponents(A_Do,param(1),param(2));
      A_Db = sparse(double(Dl),double(Dr),true,length(ebsd),length(ebsd));
      A_Db(A_Do~=0) = false;

    else

      A_Db = sparse(double(Dl(connect<1)),double(Dr(connect<1)),true,...
        length(ebsd),length(ebsd));

    end
    A_Do = A_Do | A_Do.';

    % adjacency of cells that have a common boundary
    A_Db = A_Db | A_Db.';

    % compute I_DG connected components of A_Do
    % I_DG - incidence matrix cells to grains
    I_DG = sparse(1:length(ebsd),double(connectedComponents(A_Do)),1);

  end

  function [I_FDext, I_FDint, Fext, Fint] = calcBoundary
    % distinguish between interior and exterior grain boundaries

    % cells that have a subgrain boundary, i.e. a boundary with a cell
    % belonging to the same grain
    sub = ((A_Db * I_DG) & I_DG)';                 % grains x cell
    [i,j] = find( diag(any(sub,1))*double(A_Db) ); % all adjacence to those
    sub = any(sub(:,i) & sub(:,j),1);              % pairs in a grain

    % split grain boundaries A_Db into interior and exterior
    A_Db_int = sparse(i(sub),j(sub),1,size(I_DG,1),size(I_DG,1));
    A_Db_ext = A_Db - A_Db_int;                    % adjacent over grain boundray

    % create incidence graphs
    I_FDbg = diag( sum(I_FD,2)==1 ) * I_FD;
    D_Fbg  = diag(any(I_FDbg,2));

    [ix,iy] = find(A_Db_ext);
    D_Fext  = diag(sum(abs(I_FD(:,ix)) & abs(I_FD(:,iy)),2)>0);

    I_FDext = (D_Fext| D_Fbg)*I_FD;

    [ix,iy] = find(A_Db_int);
    D_Fsub  = diag(sum(abs(I_FD(:,ix)) & abs(I_FD(:,iy)),2)>0);
    I_FDint = D_Fsub*I_FD;

    % remove empty lines from I_FD, F, and V
    isExt = full(any(I_FDext,2));
    I_FDext = I_FDext.'; I_FDext = I_FDext(:,isExt).';

    isInt = full(any(I_FDint,2));
    I_FDint = I_FDint.'; I_FDint = I_FDint(:,isInt).';

    % remove vertices that are not needed anymore
    [inUse,~,F] = unique(F(isExt | isInt,:));
    V = V(inUse,:);
    F = reshape(F,[],2);
    Fext = F(isExt(isExt | isInt),:); % external boundary segments
    Fint = F(isInt(isExt | isInt),:); % internal boundary segments

  end

  function gB = makeBoundary(F,I_FD)

    % compute ebsdInd
    [eId,fId] = find(I_FD.');
    eId = eId(:); fId = fId(:);

    % replace fId that appears a second time by fId + length(F)+1
    % such that it refers to the second column
    d = diff([0;fId]);
    fId = cumsum(d>0) + (d==0)*size(F,1);

    %  ebsdInd - [Id1,Id2] list of adjecent EBSD pixels for each segment
    ebsdInd = zeros(size(F,1),2);
    ebsdInd(fId) = eId;

    % compute misrotations
    mori = rotation.nan(size(F,1),1);
    isNotBoundary = all(ebsdInd,2);
    mori(isNotBoundary) = ...
      inv(ebsd.rotations(ebsdInd(isNotBoundary,2))) ...
      .* ebsd.rotations(ebsdInd(isNotBoundary,1));

    gB = grainBoundary(V,F,ebsdInd,grainId,ebsd.phaseId,mori,ebsd.CSList,ebsd.phaseMap,ebsd.id);

  end

  function qAdded = removeQuadruplePoints

    quadPoints = find(accumarray(reshape(Fext(full(any(I_FDext,2)),:),[],1),1) == 4);
    qAdded = 0;

    if isempty(quadPoints), return; end

    % find the 4 edges connected to the quadpoints
    I_FV = sparse(repmat((1:size(Fext,1)).',1,2),Fext,ones(size(Fext)));

    quadPoints = find(sum(I_FV) == 4).';
    [iqF,~] = find(I_FV(:,quadPoints));

    % this is a length(quadPoints x 4 list of edges
    iqF = reshape(iqF,4,length(quadPoints)).';

    % find the 4 vertices adfacent to each quadruple point
    qV = [Fext(iqF.',1).';Fext(iqF.',2).'];
    qV = qV(qV ~= reshape(repmat(quadPoints.',8,1),2,[]));
    qV = reshape(qV,4,[]).';

    % compute angle with respect to quadruple point
    qOmega = reshape(atan2(V(qV,1) - V(repmat(quadPoints,1,4),1),...
      V(qV,2) - V(repmat(quadPoints,1,4),2)),[],4);

    % sort the angles
    [~,qOrder] = sort(qOmega,2);

    % find common pixels for pairs of edges - first we try 1/4 and 2/3
    s = size(iqF);
    orderSub = @(i) sub2ind(s,(1:s(1)).',qOrder(:,i));

    iqD = I_FDext(iqF(orderSub(1)),:) .* I_FDext(iqF(orderSub(4)),:) + ...
      I_FDext(iqF(orderSub(2)),:) .* I_FDext(iqF(orderSub(3)),:);

    % if not both have one common pixel
    switchOrder = full(sum(iqD,2))~= 2;

    % switch to 3/4 and 1/2
    qOrder(switchOrder,:) = qOrder(switchOrder,[4 1 2 3]);
    orderSub = @(i) sub2ind(s,(1:s(1)).',qOrder(:,i));

    iqD = I_FDext(iqF(orderSub(1)),:) .* I_FDext(iqF(orderSub(4)),:) + ...
      I_FDext(iqF(orderSub(2)),:) .* I_FDext(iqF(orderSub(3)),:);

    % some we will not be able to remove
    ignore = full(sum(iqD,2)) ~= 2;
    iqD(ignore,:) = [];
    quadPoints(ignore) = [];
    iqF(ignore,:) = [];
    qV(ignore,:) = [];
    qOrder(ignore,:) = [];
    s = size(iqF);
    orderSub = @(i) sub2ind(s,(1:s(1)).',qOrder(:,i));

    % add an additional vertex (with the same coordinates) for each quad point
    newVid = (size(V,1) + (1:length(quadPoints))).';
    V = [V;V(quadPoints,:)];

    % include new vertex into face list, i.e. replace quadpoint -> newVid
    Ftmp = Fext(iqF(orderSub(1)),:).';
    Ftmp(Ftmp == quadPoints.') = newVid;
    Fext(iqF(orderSub(1)),:) = Ftmp.';

    Ftmp = Fext(iqF(orderSub(2)),:).';
    Ftmp(Ftmp == quadPoints.') = newVid;
    Fext(iqF(orderSub(2)),:) = Ftmp.';

    %F(iqF(orderSub(1)),:) = [qV(orderSub(1)),newVid];
    %F(iqF(orderSub(2)),:) = [newVid,qV(orderSub(2))];
    sw = Fext(:,1) > Fext(:,2);
    Fext(sw,:) = fliplr(Fext(sw,:));

    [iqD,~] = find(iqD.'); iqD = reshape(iqD,2,[]).';

    % if we have different grains - we need a new boundary
    newBd = full(sum(I_DG(iqD(:,1),:) .* I_DG(iqD(:,2),:),2)) == 0;

    % add new edges
    Fext = [Fext; [quadPoints(newBd),newVid(newBd)]];
    qAdded = sum(newBd);

    % new rows to I_FDext
    I_FDext = [I_FDext; ...
      sparse(repmat((1:qAdded).',1,2), iqD(newBd,:), 1, ...
      qAdded,size(I_FDext,2))];

    % new empty rows to I_FDint
    %I_FDint = [I_FDint; sparse(qAdded,size(I_FDint,2))];

  end

  function mergeQuadrupleGrains

    gB = grains.boundary; gB = gB(length(gB)+1-(1:qAdded));
    toMerge = false(size(gB));

    for iPhase = ebsd.indexedPhasesId

      % restrict to the same phase
      iBd = all(gB.phaseId == iPhase,2);

      if ~any(iBd), continue; end

      % check for misorientation angle % TODO
      toMerge(iBd) = angle(gB(iBd).misorientation) < 5 * degree;

    end

    [grains, parentId] = merge(grains,gB(toMerge));

    % update I_DG
    I_PC = sparse(1:length(parentId),parentId,1);
    I_DG = I_DG * I_PC;

    % update grain ids
    [grainId,~] = find(I_DG.');
  end

end

function poly = calcPolygons(I_FG,F,V)
%
% Input:
%  I_FG - incidence matrix faces to grains
%  F    - list of faces
%  V    - list of vertices

poly = cell(size(I_FG,2),1);

if isempty(I_FG), return; end

% for all grains
for k=1:size(I_FG,2)

  % inner and outer boundaries are circles in the face graph
  EC = EulerCycles(F(I_FG(:,k)>0,:));

  % first cicle should be positive and all others negatively oriented
  for c = 1:numel(EC)
    if xor( c==1 , polySgnArea(V(EC{c},1),V(EC{c},2))>0 )
      EC{c} = fliplr(EC{c});
    end
  end

  % this is needed
  for c=2:numel(EC), EC{c} = [EC{c} EC{1}(1)]; end

  poly{k} = [EC{:}];

end

end


%% brute_force_tune
% Code to tune parameters by brute force. This is for heading init.
% Obviously theoretically completely silly.
%
% Works sorta like RANSAC I guess?
%
% Adam Werries 2016, see Apache 2.0 license.

k_max = 60;
% Specify ranges
roll = linspace(-45,45,100);
pitch = linspace(-45,45,100);
yaw = linspace(-180,180,100);
% Repeat arrays
roll = repmat(roll, [1 k_max]);
pitch = repmat(pitch, [1 k_max]);
yaw = repmat(yaw, [1 k_max]);
% Generate random selections of each vector
num_items = length(roll);
roll_i = randperm(num_items);
pitch_i = randperm(num_items);
yaw_i = randperm(num_items);
rms_error_filter = zeros(1,num_items);
max_error_filter = zeros(1,num_items);
parfor i = 1:num_items
    fprintf('Iteration: %d, Roll: %08.5f, Pitch: %08.5f, Yaw: %08.5f\n', i, roll(roll_i(i)), pitch(pitch_i(i)), yaw(yaw_i(i)));
    temp_cond = init_cond;
    temp_cond(7:9) = [deg2rad(roll(roll_i(i))) deg2rad(pitch(pitch_i(i))) deg2rad(yaw(yaw_i(i)))];
    [out_profile,out_IMU_bias_est,out_KF_SD] = Loosely_coupled_INS_GNSS(temp_cond, filter_time, epoch, lla, novatel, imu, LC_KF_config, est_IMU_bias);
    xyz = out_profile(:,2:4);
    llh = ecef2lla(xyz);
    [x,y] = deg2utm(llh(:,1),llh(:,2));
    x = x-min_x;
    y = y-min_y;

    distance = ((ground_truth_full(:,1)-x).^2 + (ground_truth_full(:,2)-y).^2).^0.5;
    rms_error_filter(i) = rms(distance);
    max_error_filter(i) = max(distance);
end

[minmax, i] = min(max_error_filter);
fprintf('\nBest max: %08.5f, rms is %08.5f\n', minmax, rms_error_filter(i));
fprintf('Best iteration for max: %d, Roll: %08.5f, Pitch: %08.5f, Yaw: %08.5f\n', i, roll(roll_i(i)), pitch(pitch_i(i)), yaw(yaw_i(i)));
[minrms, i] = min(rms_error_filter);
fprintf('Best rms: %08.5f, max is %08.5f\n', minrms, max_error_filter(i));
fprintf('Best iteration for rms: %d, Roll: %08.5f, Pitch: %08.5f, Yaw: %08.5f\n', i, roll(roll_i(i)), pitch(pitch_i(i)), yaw(yaw_i(i)));
[minrms, i] = min((rms_error_filter+max_error_filter)/2);
fprintf('Best average of RMS and max: %08.4f, rms is  %08.4f, max is %08.4f\n', minrms, rms_error_filter(i), max_error_filter(i));
fprintf('Best iteration for rms: %d, Roll: %08.5f, Pitch: %08.5f, Yaw: %08.5f\n', i, roll(roll_i(i)), pitch(pitch_i(i)), yaw(yaw_i(i)));


function alpha = complexNormalAngle(varargin)
%COMPLEXNORMALANGLE compute normal angle of a vertex of a cellular complex
%
%   ALPHA = complexNormalAngle(NODES, EDGES, FACES, INDEX)
%   ALPHA = complexNormalAngle(NODES, EDGES, FACES, CELLS, INDEX)
%   Compute the nortmal angle of the polyhedral reconstruction defined be
%   nodes NODES, edges EDGES and faces FACES. For 3D reconstructions, it
%   can also contain cells CELLS. INDEX is the index of NODES for which the
%   normal angle ALPHA is computed.
%   Result is normalised between 0 and 2*PI.
%
%   ALPHA = complexNormalAngle(GRAPH, INDEX)
%   Internal data are stored in a structure GRAPH, with fields : 'nodes',
%   'edges', 'faces', and eventually 'cells'.
%
%
%   ALPHA = complexNormalAngle(..., INDICES)
%   If INDICES is an array of indices, the normal angle is computed for
%   each element of NODES(INDICES,:). The result ALPHA has the same size
%   than INDICES.
%
%
% ------
% Author: David Legland
% e-mail: david.legland@jouy.inra.fr
% Created: 2005-12-19
% Copyright 2005 INRA - CEPIA Nantes - MIAJ (Jouy-en-Josas).

%   2006-04-19 fix bug for small number of faces
%   2006-04-26 returns 0 and not [] for null complexes
%   2006-10-25 revert to return value=[] for null complex
%   2008-08-11 code clean up

cells = [];
if length(varargin)==4
    % no cells in cellular complex
    nodes = varargin{1};
    edges = varargin{2};
    faces = varargin{3};
    ind   = varargin{4};

elseif length(varargin)==5
    % cells are given
    nodes = varargin{1};
    edges = varargin{2};
    faces = varargin{3};
    cells = varargin{4};
    ind   = varargin{5};

elseif length(varargin)==2
    % data stored as structure
    graph = varargin{1};
    nodes = graph.nodes;
    edges = graph.edges;
    faces = graph.faces;
    if isfield(graph, 'cells')
        cells = graph.cells;
    end
    ind   = varargin{2};
else
    error('wrong number of arguments');
end

alpha0 = zeros([length(ind) 1]);
alpha1 = zeros([length(ind) 1]);
alpha2 = zeros([length(ind) 1]);
alpha3 = zeros([length(ind) 1]);
alpha = [];

if size(nodes, 2)==2
    % 2 dimensions

    if iscell(faces)

        % process faces as cell array
        for i=1:length(ind)
            % check that vertex is contained in the complex
            if ind(i)>size(nodes, 1)
                continue;
            end

            % normal angle of vertex
            alpha0(i) = 2*pi;

            % normal angle of edges
            alpha1(i) = length(find(sum(edges==ind(i), 2)))*pi;

            % normal angle of faces
            alpha2(i) = 0;
            for j=1:length(faces)
                face = faces{j};
                indf = find(face==ind(i));

                if ~isempty(indf)
                    alpha2(i) = alpha2(i) + polygonNormalAngle(nodes(face,:), indf);
                end
            end
        end

    else
        % process faces as arrays
        for i=1:length(ind)

            % check that vertex is contained in the complex
            if ind(i)>size(nodes, 1)
                continue;
            end

            % normal angle of vertex
            alpha0(i) = 2*pi;

            % normal angle of edges
            alpha1(i) = length(find(sum(edges==ind(i), 2)))*pi;

            % normal angle of faces
            alpha2(i) = 0;
            for j=1:size(faces, 1)
                face = faces(j,:);
                indf = find(face==ind(i));

                if ~isempty(indf)
                    alpha2(i) = alpha2(i) + polygonNormalAngle(nodes(face,:), indf);
                end
            end

        end
    end

    % compute total normal angle of reconstruction
    alpha = alpha0 - alpha1 + alpha2;

elseif size(nodes, 2)==3
    % 3 dimensions
    for i=1:length(ind)

        % check that vertex is contained in the complex
        if ind(i)>size(nodes, 1)
            continue;
        end

        % normal angle of vertex
        alpha0(i) = 4*pi;

        % normal angle of edges
        alpha1(i) = length(find(sum(edges==ind(i), 2)))*2*pi;

        % normal angle of faces
        alpha2(i) = 0;
        if iscell(faces)
            % process faces as cell array
            for j=1:length(faces)
                face = faces{j};
                indf = find(face==ind(i));

                if ~isempty(indf)
                    alpha2(i) = alpha2(i) + polygon3dNormalAngle(nodes(face,:), indf);
                end
            end

        else
            % process faces as array of double
            for j=1:size(faces, 1)
                face = faces(j,:);
                indf = find(face==ind(i));

                if ~isempty(indf)
                    alpha2(i) = alpha2(i) + polygon3dNormalAngle(nodes(face,:), indf);
                end
            end
        end

        % normal angle of cells
        alpha3(i) = 0;
        for j=1:length(cells)
            cell = cells{j};
            if iscell(faces)
                cellFaces = faces(cell);
            else
                cellFaces = faces(cell, :);
            end

            alpha3(i) = alpha3(i) + polyhedronNormalAngle(nodes, cellFaces, ind(i));
        end

        % compute total normal angle of reconstruction
        alpha = alpha0 - alpha1 + alpha2 - alpha3;
   end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% set of internal functions, for angle computations in 2D and 3D

function theta = polygonNormalAngle(points, ind)
%POLYGONNORMALANGLE compute normal angle at a vertex of the polygon

% number of points
np = size(points, 1);

% number of angles to compute
nv = length(ind);

theta = zeros(nv, 1);

for i=1:nv
    p0 = points(ind(i), :);

    % previous vertex
    if ind(i)==1
        p1 = points(np, :);
    else
        p1 = points(ind(i)-1, :);
    end

    % next vertex
    if ind(i)==np
        p2 = points(1, :);
    else
        p2 = points(ind(i)+1, :);
    end

    % compute angles
    theta1  = mod(atan2(p1(2)-p0(2), p1(1)-p0(1)) + 2*pi, 2*pi);
    theta2  = mod(atan2(p2(2)-p0(2), p2(1)-p0(1)) + 2*pi, 2*pi);
    dtheta  = mod(theta2-theta1+2*pi, 2*pi);

    % use simplification due to the fact that cells are convex
    dtheta  = min(dtheta, 2*pi-dtheta);
    theta(i)= pi - dtheta;
end

return;

function theta = polyhedronNormalAngle(nodes, faces, ind)
%POLYHEDRONNORMALANGLE compute normal angle at a vertex of a 3D polyhedron
%
%   THETA = polyhedronNormalAngle(NODES, FACES, IND);
%   where NODES is a set of 3D points, and FACES a set of faces, whose
%   elements are indices to NODES array, compute the normal angle at the
%   vertex whose index is given by IND.

% number of angles to compute
na = length(ind);

theta = zeros(na, 1);
for i=1:na

    % find faces containing given vertex,
    % and compute normal angle at each face containing vertex
    if iscell(faces)
        for j=1:length(faces)
            if ismember(ind(i), faces{j})
                % create 3D polygon
                face = nodes(faces{j}, :);

                % index of point in polygon
                indp = find(faces{j}==ind(i));

                % compute face angle
                thetaf = [thetaf polygon3dInnerAngle(face, indp)];
            end
        end
    else
        indf = find(sum(ismember(faces, ind(i)), 2));

        thetaf = zeros(length(indf), 1);
        for j=1:length(indf)
            ind2 = faces(indf(j), :);
            face = nodes(ind2, :);
            indp = find(ind2==ind(i));
            thetaf(j) = polygon3dInnerAngle(face, indp);
        end
    end

    % compute normal angle of polyhedron, by use of angle defect formula
    if ~isempty(thetaf)
        theta(i) = 2*pi - sum(thetaf);
    end
end

return;

function theta = polygon3dNormalAngle(points, ind)
%POLYGON3DNORMALANGLE compute normal angle at a vertex of the 3D polygon

theta = 2*pi - 2*polygon3dInnerAngle(points, ind);

return;

function theta = polygon3dInnerAngle(points, ind)
%POLYGON3DNORMALANGLE compute normal angle at a vertex of the 3D polygon

% number of points
np = size(points, 1);

% number of angles to compute
nv = length(ind);

theta = zeros(nv, 1);

for i=1:nv
    p0 = points(ind(i), :);

    % previous vertex
    if ind(i)==1
        p1 = points(np, :);
    else
        p1 = points(ind(i)-1, :);
    end

    % next vertex
    if ind(i)==np
        p2 = points(1, :);
    else
        p2 = points(ind(i)+1, :);
    end

    theta(i) = angle3d(p1, p0, p2);
    theta(i) = min(theta(i), 2*pi-theta(i));
    % todo: solve case for CW oriented polygons
end

return;


% a basic 3D ZTE/PETRA sequence
% achieves "TE" about 70 us and possibly below

%% high-level sequence parameters

fov=256e-3; dx=2e-3;            % Define FOV and resolution
alpha=4;                        % flip angle
Nr=300;                         % number of readout points (with some oversampling)
R= 8;                           % acceleration/undersampling (for the outer shell)
R_inner= 1;                     % acceleration/undersampling (angular direction of the inner area)
xSpoil=3;%0.6;  %0.6 for 2mm       % the amount of spoiling after the end of the readout (used to ramp to the next point)

Kmax=1/2/dx;
dK=1/fov;

%% more detailed and derived params

rf_duration=10e-6;              % duration of the excitation pulse
ro_duration=300e-6;             % read-out time: controls RO bandwidth and T2-blurring
minRF_to_ADC_time=50e-6;        % the parameter wich defines TE together with ro_discard
%ro_discard=4;                   % how many ADC samples are contaminated by RF switching artifacts and alike
rfSpoilingInc=117;              % RF spoiling increment

% system limits
sys = mr.opts('MaxGrad', 36, 'GradUnit', 'mT/m', ...
    'MaxSlew', 180, 'SlewUnit', 'T/m/s', 'rfRingdownTime', 10e-6, ...
    'rfDeadTime', 100e-6, 'adcDeadTime', 10e-6, 'gamma',42.576e6); % 1H: 42.576e6  23Na: 11.262e6

seq=mr.Sequence(sys);           % Create a new sequence object
seq_sar=mr.Sequence(sys);       % Create an auxillary sequence object for SAR testing

%% create main sequence elements

% %create alpha-degree block pulse
% rf = mr.makeBlockPulse(alpha*pi/180,'Duration',rf_duration,'system',sys);
%or create alpha-degree gaussian pulse
rf = mr.makeGaussPulse(alpha*pi/180,'Duration',rf_duration,'timeBwProduct',3,'system',sys);
Tenc=rf_duration/2+minRF_to_ADC_time+ro_duration; %encoding time
Tg=sys.rfDeadTime+rf_duration/2+Tenc+sys.adcDeadTime; % constant gradient time
Tt=ceil((Tenc*(1+xSpoil)-Tg)/sys.gradRasterTime)*sys.gradRasterTime; % transition time
Ag=Kmax/Tenc; % read gradient grdient amplitude
% % Gr=mr.makeTrapezoid('z','Amplitude',Ag,'flatTime',Tg,'riseTime',0); % this graient has no ramps, I wonder if the Matlab mr library would like it...
% Gr=mr.makeExtendedTrapezoid('z','times',[0 Tg],'amplitudes',[Ag Ag]); % this is a constant graient with no ramps
TR=Tg+Tt;

adc=mr.makeAdc(Nr,'Duration', ro_duration, 'Delay', sys.rfDeadTime+rf_duration+minRF_to_ADC_time);

SamplesBookkeeping=[];

%rfbw=1/rf_duration;
rfbw=mr.calcRfBandwidth(rf);
fprintf('Read gradient amplitude %g mT/m, effective "slice thinckess" %f mm\n', 1e3*Ag/sys.gamma, 1/Ag*rfbw*1e3);
FO=50e-3*Ag;
NF=0; % TR is broken now for NF~=0...

%% generate the sampling set on a surface of a sphere
[phi, theta, im]=spherical_samples(Kmax,dK,R);
Ns=length(phi);
SamplesBookkeeping=[SamplesBookkeeping Ns];

%% main ZTE loop
fprintf('Populating ZTE loop (%d TRs)\n', Ns);
tic;
populate_subsequence(sys,seq,rf,adc,phi,theta,im,Ns,Ag,TR,Tt,NF,FO);
toc

%% SPI loop
KstartZTE=Ag*(rf_duration/2+minRF_to_ADC_time+adc.dwell);
nKspi=floor(KstartZTE/dK);
dKspi=KstartZTE/(nKspi+1);
Tenc=rf_duration/2+minRF_to_ADC_time+adc.dwell/2;
fprintf('Populating SPI loop (%d spheres)\n', nKspi);
tic;
for s=nKspi:-1:1
    % generate the sampling set on a surface of a sphere
    [phi, theta, im]=spherical_samples(dKspi*s,dKspi,R_inner); % normally no acceleration
    Ns=length(phi);
    SamplesBookkeeping=[SamplesBookkeeping Ns];
    % update gradients
    Ag=dKspi*s/Tenc; % read gradient grdient amplitude
    % actually create the next sampling sphere
    fprintf('Populating sphere %d (%d TRs)\n',s,Ns);
    fprintf('Effective "slice thinckess" %f mm\n',  1/Ag*rfbw*1e3);
    populate_subsequence(sys,seq,rf,adc,phi,theta,im,Ns,Ag,TR,Tt,NF,FO);
end
% sample the centere of k-space
seq.addBlock(mr.makeDelay(Tt));
seq.addBlock(rf, adc, mr.makeDelay(TR-Tt));
SamplesBookkeeping=[SamplesBookkeeping 1];
toc
fprintf('Total number of SPI samples: %d; a Cartesian patch would require %d\n', sum(SamplesBookkeeping(2:end)), ceil(nKspi^3*pi*4/3));

% %% OLD!!! SPI loop
% Kspi=ceil(Ag*(rf_duration/2+minRF_to_ADC_time+adc.dwell)/dK);
% rf.delay=rf.delay+Tt;
% Tspi1=mr.calcDuration(rf)+sys.rfRingdownTime;
% adc.delay=adc.delay+Tt-Tspi1;
% Tspi2=TR-Tspi1;
%
% g=mr.makeTrapezoid('x','system',sys,'area',dK*Kspi,'duration',minRF_to_ADC_time);
% gx=g;gx.channel='x';
% gy=g;gy.channel='y';
% gz=g;gz.channel='z';
%
% gxr=g;gxr.channel='x';
% gyr=g;gyr.channel='y';
% gzSpoil=mr.makeTrapezoid('x','system',sys,'area',Kmax*(1+xSpoil),'duration',minRF_to_ADC_time);
%
% fprintf('Populating SPI loop (~%d TRs)\n', round(pi/6*((2*Kspi+1)^3)));
% tic;
%
% % SPI loop itself
% gxr.amplitude=0;
% gyr.amplitude=0;
% for k=-Kspi:Kspi
%     for j=-Kspi:Kspi
%         for i=-Kspi:Kspi
%             if i^2+j^2+k^2>Kspi^2
%                 continue;
%             end
%             % refocusing and spoiling for the previous TR
%             seq.addBlock( rf, gxr, gyr, gzSpoil, mr.makeDelay(Tspi1));
%             % new TR
%             gx.amplitude=g.amplitude/Kspi*i;
%             gy.amplitude=g.amplitude/Kspi*j;
%             gy.amplitude=g.amplitude/Kspi*k;
%             seq.addBlock( rf, adc, gx, gy, gz, mr.makeDelay(Tspi2));
%             gxr.amplitude=-gx.amplitude;
%             gyr.amplitude=-gy.amplitude;
%         end
%     end
% end
% toc

%% check whether the timing of the sequence is correct
[ok, error_report]=seq.checkTiming;

if (ok)
    fprintf('Timing check passed successfully\n');
else
    fprintf('Timing check failed! Error listing follows:\n');
    fprintf([error_report{:}]);
    fprintf('\n');
end

%%
%seq.plot('TimeRange',[0 1]);

seq.setDefinition('FOV', [fov fov fov]);
seq.setDefinition('Name', 'petra');
seq.setDefinition('SamplesPerShell', SamplesBookkeeping);

seq.write('zte_petra.seq');
return

%% create an RF-only version of the sequence (e.g. for the SAR or signal evolution testing)

tic;
[total_duration, total_numBlocks]=seq.duration();
for iB=1:total_numBlocks
    b=seq.getBlock(iB);
    bd=seq.blockDurations(iB);
    bs={mr.makeDelay(bd)};
    if ~isempty(b.rf)
        bs{end+1}=b.rf;
    end
    if ~isempty(b.adc)
        bs{end+1}=b.adc;
    end
    seq_sar.addBlock(bs);
end
toc
seq_sar.write('zte_petra_sar.seq');
return

% %% test binary storing
%
% seq.writeBinary('zte_petra.bin');
% seq_bin=mr.Sequence();
% seq_bin.readBinary('zte_petra.bin');
% seq_bin.write('zte_petra_bin.seq');
% return

%% visualize the 3D k-space
tic;
[kfa,~,kf]=seq.calculateKspacePP();
toc

%
figure;plot3(kf(1,:),kf(2,:),kf(3,:));
hold on;plot3(kfa(1,:),kfa(2,:),kfa(3,:),'r.');

%% local functions
function [phi, theta, im]=spherical_samples(Kr,dK,R)
% the number of samples equals the ceil of the area of the sphere divided
% by the area around every sample
Ns=ceil(4*pi*((Kr/dK)^2)/R);
np=0:(Ns-1);
alpha_gold=pi*(3-sqrt(5));
phi=np*alpha_gold;
%theta=0.5*pi*sqrt(np/Ns); % from Davide Piccini  https://doi.org/10.1002/mrm.22898
theta=acos(1-2*np/(Ns-1));  % from  Anton Semechko (2020). Suite of functions to perform uniform sampling of a sphere (https://github.com/AntonSemechko/S2-Sampling-Toolbox), GitHub. Retrieved October 3, 2020.

xp=sin(theta).*cos(phi);
yp=sin(theta).*sin(phi);
zp=cos(theta);
%figure; sphere; colormap gray;
%hold on; plot3(xp,yp,zp,'.');

% looking for the optimal interleaving factor
nm=round(Ns/2); % middle of the trajetory
sr=round(sqrt(Ns));
v0=[xp(nm); yp(nm); zp(nm)];
v=[xp(nm+(1:sr)); yp(nm+(1:sr)); zp(nm+(1:sr))];
%figure;plot(vecnorm((v-v0(:,ones(size(v,2),1)))));
[dKm,im]=min(vecnorm((v-v0(:,ones(size(v,2),1)))));

% fprintf('requested dK=%g achieved minimal dK=%g, min acceleration: %g, Ns=%d\n', dK/sqrt(R), dKm*Kr, (dKm*Kr/dK)^2,Ns);

% v=[xp; yp; zp]*Kr/dK/sqrt(R)*100;
% md=zeros(1,Ns);
% d=zeros(1,Ns);
% for i=1:Ns
%     d=vecnorm(v(:,i*ones(1,Ns))-v);
%     d(i)=NaN;
%     md(i)=min(d);
% end
% fprintf('dKmin=%g%%, dKmed=%g%% dKmax=%g%%\n', min(md),median(md),max(md));

end

function populate_subsequence(sys,seq,rf,adc,phi,theta,im,Ns,Ag,TR,Tt,FN, FO)
Azc=Ag*(TR-Tt)/(TR+Tt);  %*0.35;
%Azc=0; % for MoCo we need "no-gradient" event blocks to be able to apply updates

Gr=mr.makeExtendedTrapezoid('z','times',[0 TR-Tt],'amplitudes',[Ag Ag]); % this is a constant graient with no ramps

% pre-ramp the gradient to Azc
if abs(Azc)>eps
    Tpr=max(2,ceil(Azc/sys.maxSlew/sys.gradRasterTime))*sys.gradRasterTime;
    assert(Tpr<=TR);
    % this "dummy TR" doe not have an RF pulse, which is not good when it is called for the inner shells... but otherwise there would be no enough spoiling...
    seq.addBlock(mr.align('right',mr.makeDelay(TR),mr.makeExtendedTrapezoid('z','system',sys,'times',[0,Tpr],'amplitudes',[0,Azc])));
end

% the loop itself
for j=1:im
    Glast=struct('x',0,'y',0,'z',Azc);
    for i=j:im:Ns
        Gcr=mr.rotate('z',phi(i),mr.rotate('y',theta(i),Gr));
        Gcurr=struct('x',0,'y',0,'z',0);
        for g=1:length(Gcr)
            Gcurr.(Gcr{g}.channel)=Gcr{g}.waveform(1);
        end
        seq.addBlock( ...
            mr.makeExtendedTrapezoid('x','system',sys,'times',[0,Tt],'amplitudes',[Glast.x,Gcurr.x]), ...
            mr.makeExtendedTrapezoid('y','system',sys,'times',[0,Tt],'amplitudes',[Glast.y,Gcurr.y]), ...
            mr.makeExtendedTrapezoid('z','system',sys,'times',[0,Tt],'amplitudes',[Glast.z,Gcurr.z]));
        for f=-FN:FN % 'slices' loop (frequency offsets)
            rf.freqOffset=f*FO;
            rf.phaseOffset=-2*pi*f*FO*rf.t(end)/2;
            seq.addBlock( [{rf, adc}, Gcr] );
        end
        Glast=Gcurr;
    end
    rf_aux=rf;
    rf_aux.delay=rf.delay+Tt;
    if (j==im)
        Azc=0; % on the last interleave we ramp down to 0
    end
    seq.addBlock( rf_aux, ...
        mr.makeExtendedTrapezoid('x','system',sys,'times',[0,TR],'amplitudes',[Glast.x,0]), ...
        mr.makeExtendedTrapezoid('y','system',sys,'times',[0,TR],'amplitudes',[Glast.y,0]), ...
        mr.makeExtendedTrapezoid('z','system',sys,'times',[0,TR],'amplitudes',[Glast.z,Azc]));
end
end