function c = u2nonucfmt(cu, p)
%U2NONUCFMT Uniform to non-uniform filterbank coefficient format
%   Usage:  c=u2nonucfmt(cu,pk)
%
%   Input parameters:
%         cu   : Uniform filterbank coefficients.
%
%   Output parameters:
%         c    : Non-uniform filterbank coefficients.
%         p    : Numbers of copies of each filter.
%
%   `c = u2nonucfmt(cu,pk)` changes the coefficient format from
%   uniform filterbank coefficients *cu* (`M=sum(p)` channels) to
%   non-uniform coefficients *c* (`numel(p)` channels)  such that each
%   channel of *c* consinst of `p(m)` interleaved channels of *cu*.
%
%   The output *c* is a cell-array in any case.
%
%   See also: nonu2ufilterbank
%
%   References: akkva2003

complainif_notenoughargs(nargin,2,mfilename);

if isempty(cu)
   error('%s: cu must be non-empty.',upper(mfilename));
end

if iscell(cu)
    if any(cellfun(@isempty,cu));
      error('%s: Elements of cu must be non-empty.',upper(mfilename));
    end

    M = numel(cu);
    W = size(cu{1},2);

    Lc = size(cu{1},1);
    if any(Lc~=cellfun(@(cEl)size(cEl,1),cu))
        error('%s: Coefficient subbands do not have an uniform length',...
              upper(mfilename));
    end
elseif isnumeric(cu)
    M = size(cu,2);
    W = size(cu,3);
    Lc = size(cu,1);
else
    error('%s: cu must be a cell array or numeric.',upper(mfilename));
end

if isempty(p) || ~isvector(p)
   error('%s: p must be a non-empty vector.',upper(mfilename));
end

if sum(p) ~= M
    error(['%s: Total number of filters in p does not comply with ',...
           'number of channels'],upper(mfilename));
end

Mnonu = numel(p);
c = cell(Mnonu,1);
p = p(:);
pkcumsum = cumsum([1;p]);
crange = arrayfun(@(pEl,pcEl)pcEl:pcEl+pEl-1,p,pkcumsum(1:end-1),'UniformOutput',0);

if iscell(cu)
    for m=1:Mnonu
        ctmp = [cu{crange{m}}].';
        c{m} = reshape(ctmp(:),W,[]).';
    end
else
    for m=1:Mnonu
        c{m} = zeros(p(m)*Lc,W,assert_classname(cu));
        for w=1:W
            c{m}(:,w) = reshape(cu(:,crange{m},w).',1,[]);
        end
    end
end

% Post check whether there is the same number of coefficients
if sum(cellfun(@(cEl) size(cEl,1),c)) ~= M*Lc
    error(['%s: Invalid number of coefficients in subbands.'],upper(mfilename));
end


% Compute derivative k'(x^p,x^q) of a stationary covariance k(d2) (ard or iso)
% w.r.t. to squared distance d2 = (x^p - x^q)'*inv(P)*(x^p - x^q) measure. Here
% P is either diagonal with ARD parameters ell_1^2,...,ell_D^2 where D is the
% dimension of the input space or ell^2 times the unit matrix for isotropic
% covariance.
% The derivatives can only be computed for one of the following eight
% covariance functions: cov{Matern|PP|RQ|SE}{iso|ard}.
%
% Copyright (c) by Hannes Nickisch, 2013-10-28.
%
% See also INFFITC.M, COVFITC.M.

function Kp = cov_deriv_sq_dist(cov,hyp,x,z)
  if nargin<4, z = []; end                                 % make sure, z exists
  xeqz = numel(z)==0; dg = strcmp(z,'diag') && numel(z)>0;      % determine mode

  if iscell(cov), covstr = cov{1}; else covstr = cov; end
  if ~ischar(covstr), covstr = func2str(covstr); end
  if numel([strfind(covstr,'iso'),strfind(covstr,'ard')])==0
    error('Only iso|ard covariances allowed for derivatives w.r.t. xu.')
  elseif numel([strfind(covstr,'covLIN');
                strfind(covstr,'covGabor');
                strfind(covstr,'covPER')])>0
    error('Gabor|LIN|PER covariances not allowed for derivatives w.r.t. xu.')
  end

  [n,D] = size(x);
  if numel(strfind(covstr,'iso')), id = 1:D; else id = 1; end  % *iso covariance
  ell1 = exp(hyp(1));                        % first characteristic length scale

  Kp = feval(cov{:},hyp,x,z,1);              % use derivative w.r.t. log(ell(1))
  % precompute squared distances
  if dg                                                             % vector kxx
    d2 = zeros(n,1);
  else
    if xeqz                                               % symmetric matrix Kxx
      d2 = sq_dist(x(:,id)'/ell1);
    else                                                 % cross covariances Kxz
      d2 = sq_dist(x(:,id)'/ell1,z(:,id)'/ell1);
    end
  end
  Kp = -1/2*Kp./d2; Kp(d2==0) = 0;


function [X,info] = IRrrgmres(A,b,varargin)
% IRrrgmres Range Restricted GMRES for square systems
%
% options  = IRrrgmres('defaults')
% [X,info] = IRrrgmres(A,b);
% [X,info] = IRrrgmres(A,b,K);
% [X,info] = IRrrgmres(A,b,options);
% [X,info] = IRrrgmres(A,b,K,options);
%
% This function implements the Range Restricted GMRES method for solving
% a square system.
%
% With 'defaults' as input returns the default options.  Otherwise outputs
% the iterates specified in K, using max(K) as MaxIter, and using all other
% default options.  With options as input: uses the user-specified options
% and all the other default options.
%
% Inputs:
%  A : either (a) a full or sparse matrix
%             (b) a matrix object that performs the matrix*vector operation
%             (c) user-defined function handle
%  b : right-hand side vector
%  K : (optional) integer vector that specifies which iterates are returned
%      in X; the maximum number of iterations is assumed to be max(K)
%      [ positive integer | vector of positive components ]
%  options : structure with the following fields (optional)
%      x0         - initial guess for the iterations; default = zero vector
%                   [ array | {'none'} ]
%      MaxIter    - maximum allowed number of iterations
%                   [ {100} | positive integer ]
%                   NOTE: K overrules MaxIter if both are assigned
%      x_true     - true solution; allows us to returns error norms with
%                   respect to x_true at each iteration
%                   [ array | {'none'} ]
%      NoiseLevel - norm of noise in rhs divided by norm of rhs
%                   [ {'none'} | nonnegative scalar ]
%      eta        - safety factor for the discrepancy principle
%                   [ {1.01} | scalar greater than (and close to) 1 ]
%      DecompOut  - return the decomposition to the user
%                   [ 'on' | {'off'} ]
%      IterBar    - shows the progress of the outer iterations
%                   [ {'on'} | 'off' ]
%      NoStop     - specifies weather the iterations should proceed
%                   after a stopping criterion has been satisfied
%                   [ 'on' | {'off'} ]
% Note: the options structure can be created using the function IRset.
%
% Outputs:
%   X : computed solutions, stored column-wise (at the iterations listed in K)
%   info: structure with the following fields:
%      its     - number of the last computed iteration
%      saved_iterations - iteration numbers of iterates stored in X
%      StopFlag - stringg that describes the output/stopping condition:
%                   * Performed max number of iterations
%                   * Residual tolerance satisfied (discrepancy principle)
%                   * Normal equationa residual tolerance satisfied
%      Rnrm     - relative residual norms at each iteration
%      Xnrm     - solution norms at each iteration
%      Enrm     - error norms (requires x_true) at each iteration
%      V        - matrix generated by the Arnoldi algorithm: the orthonormal
%                 columns of V are a basis for the RRGMRES solution
%      Q        - orthogonal matrix in QR factorization of Hessenberg matrix
%      T        - triangular matrix in QR factorization of Hessenberg matrix
%      StopReg  - struct containing information about the solution that
%                 satisfies the stopping criterion.  Fields:
%                   It   : iteration where stopping criterion is satisfied
%                   X    : solution satisfying the stopping criterion
%                   Enrm : the corresponding relative error (requires x_true)
%      BestReg  - struct containing information about the solution that
%                 minimizes Enrm (requires x_true). Fields:
%                   It   : iteration where the minimum is attained
%                   X    : best solution
%                   Enrm : best relative error
%
% See also: IRcgls, IRhybrid_gmres, IRget, IRset

% Silvia Gazzola, University of Bath
% Per Christian Hansen, Technical University of Denmark
% James G. Nagy, Emory University
% April, 2018.

% This file is part of the IR Tools package and is distributed under the
% 3-Clause BSD License. A separate license file should be provided as part
% of the package.

% Set default values for options.
defaultopt = struct('x0','none', 'MaxIter',100, 'x_true','none', ...
    'NoiseLevel','none', 'eta',1.01, 'IterBar','on', ...
    'NoStop','off', 'DecompOut','off');

% If input is 'defaults,' return the default options in X.
if nargin==1 && nargout <= 1 && isequal(A,'defaults')
    X = defaultopt;
    return;
end

defaultopt.restart = 'off';
defaultopt.verbosity = 'on';

% Check for acceptable number of optional input arguments.
switch length(varargin)
    case 0
        K = []; options = [];
    case 1
        if isa(varargin{1}, 'double')
            K = varargin{1}; options = [];
        else
            K = []; options = varargin{1};
        end
    case 2
        if isa(varargin{1}, 'double')
            K = varargin{1}; options = varargin{2};
        else
            K = varargin{2}; options = varargin{1};
        end
        if isfield(options, 'MaxIter') && ~isempty(options.MaxIter) && (~isempty(K) && options.MaxIter ~= max(K))
            warning('The value of MaxIter is discarded; the maximum value in K is taken as MaxIter')
        end
    otherwise
        error('Too many input parameters')
end

if isempty(options)
    options = defaultopt;
end

options = IRset(defaultopt, options);

MaxIter    = IRget(options, 'MaxIter',    [], 'fast');
x_true     = IRget(options, 'x_true',     [], 'fast');
NoiseLevel = IRget(options, 'NoiseLevel', [], 'fast');
eta        = IRget(options, 'eta',        [], 'fast');
IterBar    = IRget(options, 'IterBar',    [], 'fast');
NoStop     = IRget(options, 'NoStop',     [], 'fast');
restart    = IRget(options, 'restart',    [], 'fast');
verbose    = IRget(options, 'verbosity',  [], 'fast');
DecompOut  = IRget(options, 'DecompOut',  [], 'fast');

restart = strcmp(restart, 'on');
verbose = strcmp(verbose, 'on');

% Setting K.
if isempty(K)
    K = MaxIter;
end
% Sorting the iterations (in case they are shuffled in input).
K = K(:); K = sort(K,'ascend'); K = unique(K);
if ~((isreal(K) && (all(K > 0)) && all(K == floor(K))))
    error('K must be a vector of positive real integers')
end
if K(end) ~= MaxIter
    MaxIter = K(end);
end

StopIt = MaxIter;

if isempty(NoiseLevel) || strcmp(NoiseLevel,'none')
    Rtol = 0;
else
    Rtol = eta*NoiseLevel;
end

n = length(b);
test_sq = ones(n,1);
try
    test_sq = A_times_vec(A, test_sq);
    if (length(test_sq)~=n)
        error('The matrix A shuold be square; check the length of b.')
    end
catch
    error('The matrix A must be square; check the length of b.')
end

% See if an initial guess is given; if not then use 0 as initial guess.
x0 = IRget(options, 'x0', [], 'fast');
Ab = A_times_vec(A, b);
nrmb = norm(b(:));
if strcmp(x0,'none')
    x0 = zeros(n,1);
    r = b;
    rr = Ab;
else
    try
        Ax = A_times_vec(A, x0);
        r = b - Ax;
        rr = Ab - Ax;
    catch
        error('Check the length of x')
    end
end
if restart
    ktotcount  = IRget(options, 'ktotcount', [],'fast');
    TotIterMax = IRget(options, 'TotIterMax',[],'fast');
    if strcmp(TotIterMax, 'none') || TotIterMax < MaxIter
        TotIterMax = MaxIter;
    end
    if strcmp(ktotcount, 'none')
        error('The total iteration counter must be assigned')
    end
    Ktot = IRget(options, 'Ktot', [], 'fast');
    % No checks on Ktot, it should be given from IRrestart.
end

% Declare matrices and prepare for the iterations.
X       = zeros(n,length(K));
Xnrm    = zeros(MaxIter,1);
Rnrm    = zeros(MaxIter,1);
V       = zeros(n,MaxIter); % Orthonormal vectors spanning Krylov subspace
h       = zeros(MaxIter,1); % New column of Hessenberg matrix H
Q       = zeros(MaxIter+1); % H = Q*T, Q orthogonal
T       = zeros(MaxIter);   % H = Q*T, T upper triangular
invT    = zeros(MaxIter);   % Inverse of the matrix T
coeff   = zeros(MaxIter,1); % Solution of the projected problem, i.e.,
                            % coefficient w.r.t. the Krylov basis.
rhs = zeros(MaxIter+1,1);   % Projected right-hand side

if strcmp(x_true,'none')
    errornorms = false;
else
    errornorms = true;
    Enrm = zeros(max(K),1);
    nrmtrue = norm(x_true(:));
    BestReg.It = [];
    BestReg.X = [];
    BestReg.Enrm = [];
    BestEnrm = 1e10;
    BestReg.Xnrm = [];
    BestReg.Rnrm = [];
end

NoStop = strcmp(NoStop,'on');

if restart
    saved_iterations = zeros(1, length(Ktot));
else
    saved_iterations = zeros(1, length(K));
end

% Prepare for the iterations.
alpharr  = norm(rr);
V(:,1) = rr/alpharr;
Q(1,1) = 1;
rhs(1) = V(:,1)'*r;
y = 0;

% Iterate.
noIterBar = strcmp(IterBar,{'off'});
if ~noIterBar
  h_wait = waitbar(0, 'Running iterations, please wait ...');
end
j = 0;
for k=1:MaxIter
    if restart, ktotcount = ktotcount + 1; end
    if ~noIterBar
        waitbar(k/MaxIter, h_wait)
    end
    w = A_times_vec(A, V(:,k));
    % Modified Gram-Schmidt on the new vector.
    for ll=1:k
        h(ll) = V(:,ll)'*w;
        w = w - V(:,ll)*h(ll);
    end
    alpha = norm(w);
    % Store new Arnoldi vector and update projected rhs.
    V(:,k+1) = w/alpha;
    rhs(k+1) = V(:,k+1)'*r;
    beta = rhs(1:k+1);
    % Apply previous rotations to h.
    T(1:k,k) = Q(1:k,1:k)'*h(1:k);
    % Compute Givens rotation parameters.
    rc = T(k,k);
    if alpha == 0
        c = 1; s = 0;
    elseif abs(alpha) > abs(rc)
        tau = -rc/alpha;
        s = 1 / sqrt(1 + abs(tau)^2);
        c = s*tau;
    else
        tau = -alpha/rc;
        c = 1 / sqrt(1 + abs(tau)^2);
        s = c*tau;
    end
    % Apply givens rotations.
    T(k,k) = c'*rc - s'*alpha;
    Q(1:k,[k,k+1]) = Q(1:k,k)*[c s];
    Q(k+1,[k,k+1]) = [-s c];
    if abs(T(k,k)) <= eps
        disp('Hessenberg matrix is (numerically) singular')
        X(:,j+1) = x;
        X = X(:,1:j+1);
        V = V(:,1:k+1);
        if restart
            saved_iterations(j+1) = ktotcount-1;
        else
            saved_iterations(j+1) = k-1;
        end
        saved_iterations = saved_iterations(1:j+1);
        % if k-1 < StopIt, StopIt = k-1; end
        if k>1
            Xnrm = Xnrm(1:k-1);
            Rnrm = Rnrm(1:k-1);
            if errornorms, Enrm = Enrm(1:k-1); end
        end
        % Stop because the Hessenberg matrix is (numerically) rank def.
        if StopIt == MaxIter
            StopFlag = 'Breakdown of the Arnoldi algorithm';
            StopReg.It = k-1;
            StopReg.X = x;
            if errornorms, StopReg.Enrm = Enrm(k-1); end
        end
        k = k-1;
        break
    end
    % Update the inverse of T.
    invT(1:k-1,k) = -(invT(1:k-1,1:k-1)*T(1:k-1,k))/T(k,k);
    invT(k,k) = 1/T(k,k);
    % Update the previous projected solution.
    coeff(1:k-1) = y;
    y = coeff(1:k) + (Q(1:k+1,k)'*beta)*invT(1:k,k);
    % Update solution.
    x = x0 + V(:,1:k)*y;
    % Compute norms.
    Xnrm(k)    = norm(x(:));
    if k==1
        Rnrm(k) = sqrt( abs( r'*r - abs(Q(1:k+1,k)'*beta)^2 ) )/nrmb;
    else
        Rnrm(k) = sqrt( abs( (Rnrm(k-1)*nrmb)^2 - abs(Q(1:k+1,k)'*beta)^2 ) )/nrmb;
    end
    if errornorms
        Enrm(k) = norm(x_true(:)-x(:))/nrmtrue;
        if Enrm(k)<BestEnrm
            BestReg.It = k;
            BestReg.X = x;
            BestEnrm = Enrm(k);
            BestReg.Enrm = BestEnrm;
            BestReg.Xnrm = Xnrm(k);
            BestReg.Rnrm = Rnrm(k);
        end
    end
    AlreadySaved = 0;
    if any(k==K)
        j = j+1;
        X(:,j) = x;
        saved_iterations(j) = k;
        if restart, saved_iterations(j) = ktotcount; end
        AlreadySaved = 1;
    end
    if restart
        if any(ktotcount == Ktot) && ~ AlreadySaved
            j = j+1;
            X(:,j) = x;
            saved_iterations(j) = ktotcount;
            AlreadySaved = 1;
        end
        if ktotcount == TotIterMax
            if ~ AlreadySaved
                j = j+1;
                saved_iterations(j) = ktotcount;
                X(:,j) = x;
                AlreadySaved = 1;
            end
            StopIt = k;
            StopReg.It = k;
            StopReg.X = x;
            if errornorms
                Enrm = Enrm(1:k);
                StopReg.Enrm = Enrm(k);
            end
            Xnrm     = Xnrm(1:k);
            Rnrm     = Rnrm(1:k);
            X = X(:,1:j);
            saved_iterations = saved_iterations(1:j);
            if verbose
                disp('Reached maximum number of iterations')
            end
            StopFlag = 'Reached maximum number of iterations';
            break
        end
    end
    if (Rnrm(k) <= Rtol) && (StopIt == MaxIter)
        if verbose
            disp('Residual tolerance satisfied')
        end
        StopFlag = 'Residual tolerance satisfied';
        if ~AlreadySaved && ~NoStop
            j = j+1;
            X(:,j) = x;
            if restart
                saved_iterations(j) = ktotcount;
            else
                saved_iterations(j) = k;
            end
            AlreadySaved = 1;
        end
        StopIt = k;
        StopReg.It = k;
        StopReg.X = x;
        if errornorms, StopReg.Enrm = Enrm(k); end
        if ~ NoStop
            Xnrm    = Xnrm(1:k);
            Rnrm    = Rnrm(1:k);
            if errornorms, Enrm = Enrm(1:k); end
            X = X(:,1:j);
            saved_iterations = saved_iterations(1:j);
            break
        end
    end
end
if k == MaxIter
  if StopIt == MaxIter
        % Stop because max number of iterations reached.
        if verbose
            disp('Reached maximum number of iterations')
        end
        StopFlag = 'Reached maximum number of iterations';
        if ~AlreadySaved
            j = j+1;
            X(:,j) = x;
            if restart
                saved_iterations(j) = ktotcount;
            else
                saved_iterations(j) = k;
            end
        end
        StopReg.It = k;
        StopReg.X = x;
        if errornorms, StopReg.Enrm = Enrm(k); end
        Xnrm    = Xnrm(1:k);
        Rnrm    = Rnrm(1:k);
        if errornorms, Enrm = Enrm(1:k); end
        X = X(:,1:j);
        saved_iterations = saved_iterations(1:j);
  end
end
if ~noIterBar, close(h_wait), end
if nargout==2
  info.its = k;
  if restart
      info.ktotcount = ktotcount;
  end
  info.saved_iterations = saved_iterations(1:j);
  info.StopFlag = StopFlag;
  info.StopReg = StopReg;
  info.Rnrm = Rnrm(1:k);
  info.Xnrm = Xnrm(1:k);
  if errornorms
    info.Enrm = Enrm(1:k);
    info.BestReg = BestReg;
  end
  if strcmp(DecompOut,'on')
    info.V = V;
    info.Q = Q;
    info.T = T;
  end
end


classdef KernelPCA < handle
    %{
        Kernel Principal component analysis (KPCA)

        Version 2.2, 14-MAY-2021
        Email: iqiukp@outlook.com
    -------------------------------------------------------------------
    %}

    properties
        data
        label
        numComponents
        explainedLevel
        kernelFunc = Kernel('type', 'gaussian', 'gamma', 0.5)
        lambda
        coefficient %　principal component coefficients
        score %　principal component scores.
        cumContribution % cumulative contribution rate
        newData % Transform the mapping data back to original space
        T2
        SPE
        T2Limit
        SPELimit
        numSPEAlarm
        numT2Alarm
        accuracySPE
        accuracyT2
        eigenvalueTolerance = 1e-8 % tolerance for eigenvalues
        alpha = 1 % hyperparameter of the ridge regression that learns the reconstruction
        theta = 0.7 % experience parameter of fault diagnosis
        significanceLevel = 0.95
        display = 'on'
        temporary
        diagnosis = []
        runningTime
    end

    properties (Dependent)
        numSamples
        numFeatures
    end

    methods
        function obj = KernelPCA(parameter)
            name_ = fieldnames(parameter);
            for i = 1:size(name_, 1)
                obj.(name_{i, 1}) = parameter.(name_{i, 1});
            end
            KernelPCAOption.checkInputForDiagnosis(obj);
        end

        function obj = train(obj, data)
            tStart = tic;
            obj.data = data;
            obj.label = ones(obj.numSamples, 1);

            % compute the kernel matrix
            K = obj.kernelFunc.computeMatrix(obj.data, obj.data);
            % centralize the kernel matrix
            unit = ones(obj.numSamples, obj.numSamples)/obj.numSamples;
            K_c = K-unit*K-K*unit+unit*K*unit;
            % compute the eigenvalues and eigenvectors
            [obj.coefficient, obj.lambda] = obj.computeEigenvalue(K_c);
            %
            obj.score = K_c* obj.coefficient(:, 1:obj.numComponents);
            obj.newData = obj.reconstruct;
            obj.temporary.K = K;
            obj.temporary.K_c = K_c;
            obj.temporary.unit = unit;
            obj.computeControlLimit;

            % compute accuracy
            T2AlarmIndex = find(obj.T2 > obj.T2Limit);
            SPEAlarmIndex = find(obj.SPE > obj.SPELimit);
            obj.numSPEAlarm = length(SPEAlarmIndex);
            obj.numT2Alarm = length(T2AlarmIndex);

            label_ = obj.label;
            label_(SPEAlarmIndex) = -1;
            obj.accuracySPE = sum(label_ == obj.label)/obj.numSamples;
            label_ = obj.label;
            label_(T2AlarmIndex) = -1;
            obj.accuracyT2 = sum(label_ == obj.label)/obj.numSamples;
            %
            obj.runningTime = toc(tStart);
            if strcmp(obj.display, 'on')
                KernelPCAOption.displayTrain(obj)
            end
        end

        function results = test(obj, data, varargin)
            % test the model from the given data
            tStart = tic;
            results.evaluation = 'off';
            if nargin == 3
                results.evaluation = 'on';
                testLabel = varargin{1};
            end
            Kt = obj.kernelFunc.computeMatrix(data, obj.data);
            % centralize the kernel matrix
            unit = ones(size(data, 1), obj.numSamples)/obj.numSamples;
            Kt_c = Kt-unit*obj.temporary.K-Kt*obj.temporary.unit+unit*obj.temporary.K*obj.temporary.unit;
            %
            results.numSamples = size(data, 1);
            results.data = data;
            results.score = Kt_c*obj.coefficient(:, 1:obj.numComponents);

            % compute Hotelling's T2 statistic
            results.T2 = diag(results.score/diag(obj.lambda(1:obj.numComponents))*results.score');
            % compute the squared prediction error (SPE)
            results.SPE = sum((Kt_c*obj.coefficient).^2, 2)-sum(results.score.^2 , 2);

            % compute accuracy
            results.T2AlarmIndex = find(results.T2 > obj.T2Limit);
            results.SPEAlarmIndex = find(results.SPE > obj.SPELimit);
            if strcmp(results.evaluation, 'on')
                label_ = ones(size(results.data, 1), 1);
                label_(results.SPEAlarmIndex) = -1;
                results.accuracySPE = sum(label_ == testLabel)/results.numSamples;
                label_ = ones(size(results.data, 1), 1);
                label_(results.T2AlarmIndex) = -1;
                results.accuracyT2 = sum(label_ == testLabel)/results.numSamples;
            end
            results.numSPEAlarm = length(results.SPEAlarmIndex);
            results.numT2Alarm = length(results.T2AlarmIndex);
            results.temporary.Kt = Kt;
            results.runningTime = toc(tStart);

            if strcmp(obj.display, 'on')
                KernelPCAOption.displayTest(results)
            end

            % fault diagnosis
            if strcmp(obj.diagnosis.switch, 'on')
                results = obj.diagnose(results);
            end
        end

        function newData = reconstruct(obj)
            % Transform the mapping data back to original space.
            % References
            % ----------
            % Bakır G H, Weston J, Schölkopf B. Learning to find pre-images[J].
            % Advances in neural information processing systems, 2004, 16: 449-456.

            K_1 =  obj.kernelFunc.computeMatrix(obj.score, obj.score);
            K_1_ = K_1;
            for i = 1:obj.numSamples
                K_1(i, i) = K_1(i, i)+obj.alpha;
            end
            dual_coef = mldivide(K_1, obj.data);
            K_2 =  K_1_;
            newData = K_2*dual_coef;
        end

        function [coefficient, lambda] = computeEigenvalue(obj, K_c)
            % compute the eigenvalues and eigenvectors
            rng('default')
            [V, D, ~] = svd(K_c/obj.numSamples, 'econ');
            % ill-conditioned matrix
            if ~(isreal(V)) || ~(isreal(D))
                V = real(V);
                D = real(D);
            end
            lambda_ = diag(D);
            obj.cumContribution = cumsum(lambda_/sum(lambda_));

            if isempty(obj.numComponents)
                obj.numComponents = obj.numFeatures;
            else
                if obj.numComponents >= 1
                    obj.numComponents = obj.numComponents;
                else
                    obj.explainedLevel = obj.numComponents;
                    obj.numComponents = find(obj.cumContribution >= obj.numComponents, 1, 'first');
                end
            end
            lambda = lambda_;
            try
                coefficient = V./sqrt(obj.numSamples*lambda_)';
            catch
                coefficient = zeros(obj.numSamples, obj.numSamples);
                for i = 1:obj.numSamples
                    coefficient(:, i) = V(:, i)/sqrt(obj.numSamples*lambda_(i, 1));
                end
            end
        end

        function computeControlLimit(obj)
            % compute the squared prediction error (SPE)
            temp1 = obj.temporary.K_c*obj.coefficient;
            temp2 = obj.temporary.K_c*obj.coefficient(:, 1:obj.numComponents);
            obj.SPE = sum(temp1.^2, 2)-sum(temp2.^2, 2);
            obj.T2 = diag(obj.score/diag(obj.lambda(1:obj.numComponents))*obj.score');

            % compute the T2 limit (the F-Distribution)
            k = obj.numComponents*(obj.numSamples-1)/(obj.numSamples-obj.numComponents);
            obj.T2Limit = k*finv(obj.significanceLevel, obj.numComponents, obj.numSamples-obj.numComponents);

            % compute the SPE limit (the Chi-square Distribution)
            a = mean(obj.SPE);
            b = var(obj.SPE);
            g = b/2/a;
            h = 2*a^2/b;
            obj.SPELimit = g*chi2inv(obj.significanceLevel, h);
        end

        function results = diagnose(obj, results, varargin)
            % falut diagnosis
            tStart = tic;
            fprintf('\n')
            fprintf('*** Fault diagnosis ***\n')
            fprintf('Fault diagnosis start...\n')
            results.diagnosis = obj.diagnosis;
            data_ = results.data;
            results.diagnosis.data = data_(results.diagnosis.start:results.diagnosis.end, :);
            % contribution plots of train data
            if ~exist('.\data', 'dir')
                mkdir data;
            end
            file_ = dir('.\data\*.mat');
            name_ = {file_(1:length(file_)).name}';
            if ismember('diagnosis.mat', name_)
                load('.\data\diagnosis.mat', 'tmp_')
                tmp__ = KernelPCAOption.saveCheckObj(obj);
                if isequal(tmp__, tmp_)
                    load('.\data\diagnosis.mat', 'T2CpsTrain', 'SPECpsTrain')
                else
                    [T2CpsTrain, SPECpsTrain] = obj.computeContribution(results, 'train');
                    tmp_ = KernelPCAOption.saveCheckObj(obj);
                    save('.\data\diagnosis.mat', 'T2CpsTrain', 'SPECpsTrain', 'tmp_')
                end
            else
                [T2CpsTrain, SPECpsTrain] = obj.computeContribution(results, 'train');
                tmp_ = KernelPCAOption.saveCheckObj(obj);
                save('.\data\diagnosis.mat', 'T2CpsTrain', 'SPECpsTrain', 'tmp_')
            end

            % contribution plots of test data
            [T2CpsTest, SPECpsTest] = obj.computeContribution(results, 'test');
            % normalize the contribution plots
            T2CpsTrainMu = mean(T2CpsTrain);
            T2CpsTrainStd = std(T2CpsTrain);

            try
                T2Cps = bsxfun(@rdivide, bsxfun(@minus, T2CpsTest, T2CpsTrainMu), T2CpsTrainStd);
            catch
                mu_array = repmat(T2CpsTrainMu, size(T2CpsTest,1), 1);
                st_array = repmat(T2CpsTrainStd, size(T2CpsTest,1), 1);
                T2Cps = (T2CpsTest-mu_array)./st_array;
            end
            SPECpsTrainMu = mean(SPECpsTrain);
            SPECpsTrainStd = std(SPECpsTrain);
            try
                SPECps = bsxfun(@rdivide, bsxfun(@minus, SPECpsTest, SPECpsTrainMu), SPECpsTrainStd);
            catch
                mu_array = repmat(SPECpsTrainMu, size(SPECpsTest,1), 1);
                st_array = repmat(SPECpsTrainStd, size(SPECpsTest,1), 1);
                SPECps = (SPECpsTest-mu_array)./st_array;
            end

            % store the results
            results.diagnosis.T2Cps = T2Cps;
            results.diagnosis.SPECps = SPECps;

            %
            T2Cps_ = mean(abs(T2Cps), 1);
            results.diagnosis.meanT2Cps = T2Cps_/sum(T2Cps_, 2);

            SPECps_ = mean(abs(SPECps), 1);
            results.diagnosis.meanSPECps = SPECps_/sum(SPECps_, 2);

            %
            [value, index] = sort(results.diagnosis.meanT2Cps, 'descend');
            results.diagnosis.faultVariabeT2.value = value;
            results.diagnosis.faultVariabeT2.index = index;

            [value, index] = sort(results.diagnosis.meanSPECps, 'descend');
            results.diagnosis.faultVariabeSPE.value = value;
            results.diagnosis.faultVariabeSPE.index = index;

            results.diagnosis.runningTime = toc(tStart);
            if strcmp(obj.display, 'on')
                KernelPCAOption.displayDiagnose(results)
            end
        end

        function [T2Cps, SPECps] = computeContribution(obj, result, type)

            %  Compute the Contribution Plots (CPs)
            %
            %  Reference
            %  [1]  Deng X, Tian X. A new fault isolation method based on unified
            %       contribution plots[C]//Proceedings of the 30th Chinese Control
            %       Conference. IEEE, 2011: 4280-4285.
            % -------------------------------------------------------------------
            %  Thanks for the code provided by Rui.
            % --------------------------------------------------------------------

            data_ = obj.data;
            switch type
                case 'train'
                    Kt = obj.temporary.K;
                    Y = data_;
                case 'test'
                    Kt = result.temporary.Kt;
                    Y = result.diagnosis.data;
            end

            K = obj.temporary.K;
            M = size(data_, 1);
            [Mt, d] = size(Y);

            A_T2 = obj.coefficient(:, 1:obj.numComponents)*...
                diag(obj.lambda(1:obj.numComponents))^(-1)*...
                obj.coefficient(:, 1:obj.numComponents)';

            A_SPE = obj.coefficient(:, 1:obj.numComponents)*...
                obj.coefficient(:, 1:obj.numComponents)';
            newY = Y*obj.theta;

            % initialization
            Knew = zeros(Mt, M);
            Knew_d1 = zeros(1, M);
            Knew_d2 = zeros(Mt, M);
            T2Cps = zeros(Mt, d);
            SPECps = zeros(Mt, d);
            Knew_s = zeros(Mt, M);
            sigma = sqrt(1/2/obj.kernelFunc.gamma);

            % compute contribution of statistic
            for i = 1:Mt
                for j = 1:d
                    for k = 1:M
                        Knew(i, k) = Kt(i, k);
                        Knew_d1(k) = Knew(i, k)*2*obj.theta*(newY(i, j)-data_(k, j))/(-sigma^2); % derivative
                        Knew_d2(i, k) = -2*Knew_d1(k);
                    end
                    Knew_d1_s = Knew_d1-ones(1, M)*mean(Knew_d1);
                    Knew_s(i, :) = Knew(i, :)-ones(1, M)*K/M-Knew(i, :)*ones(M) ...
                        /M+ones(1, M)/M*K*ones(M)/M;
                    % contribution of T2
                    T2Cps(i, j) = Y(i, j)*(Knew_d1_s*A_T2*Knew_s(i, :)' ...
                        +Knew_s(i, :)*A_T2*Knew_d1_s');
                    % contribution of SPE
                    SPECps(i, j)= Y(i, j)*mean(Knew_d2(i, :))-Y(i, j) ...
                        *(Knew_d1_s*A_SPE*Knew_s(i, :)'+Knew_s(i, :)*A_SPE*Knew_d1_s');
                end
            end
        end

        function numSamples = get.numSamples(obj)
            numSamples= size(obj.data, 1);
        end

        function numFeatures = get.numFeatures(obj)
            numFeatures= size(obj.data, 2);
        end
    end
end


function [V,Q,QT,QF,QL] = voxel_surface(W,varargin)
  % VOXEL_SURFACE Compute the surface quad mesh of a 3d logical image.
  %
  % [V,Q,QT,QF,QL] = voxel_surface(W,varargin)
  %
  % Inputs:
  %   W  n by m by k binary matix
  %   Optional:
  %     'Centers' followed by n*m*k by 3 list of cell centers
  % Outputs:
  %   V  (n+1)*(m+1)*(k+1) by 3 list of cell corner positions
  %   Q  #Q by 4 list of quads indexing V
  %   QT  n*m*k by 4 list of "top" quad faces indexing V
  %   QF  n*m*k by 4 list of "front" quad faces indexing V
  %   QL  n*m*k by 4 list of "left" quad faces indexing V
  %
  % Examples:
  %   [V,Q] = voxel_surface(W);
  %   [V,IM] = remove_unreferenced(V,Q);
  %   Q = IM(Q);
  %   trisurf(Q,V(:,1),V(:,2),V(:,3));
  %   axis equal;
  %
  % See also: voxelize, voxel_grid

  BC = [];
  % default values
  % Map of parameter names to variable names
  params_to_variables = containers.Map( ...
    {'Centers'}, ...
    {'BC'});
  v = 1;
  while v <= numel(varargin)
    param_name = varargin{v};
    if isKey(params_to_variables,param_name)
      assert(v+1<=numel(varargin));
      v = v+1;
      % Trick: use feval on anonymous function to use assignin to this workspace
      feval(@()assignin('caller',params_to_variables(param_name),varargin{v}));
    else
      error('Unsupported parameter: %s',varargin{v});
    end
    v=v+1;
  end

  dim = ndims(W);

  switch dim
  case 2
    side = [size(W,2) size(W,1)];
    if isempty(BC)
      [X,Y] = meshgrid( ...
        linspace(0,1,side(1)), ...
        linspace(0,1,side(2)));
      % barycenters of cells
      BC = [X(:) Y(:)];
    end
    NV = min(BC);
    XV = max(BC);
    r = (XV-NV)./(side-1);
    [X,Y] = meshgrid( ...
      (NV(1)-0.5*r(1))+linspace(0,1,side(1)+1)*(XV(1)-NV(1)+r(1)), ...
      (NV(2)-0.5*r(2))+linspace(0,1,side(2)+1)*(XV(2)-NV(2)+r(2)));

    [Q,V] = surf2patch(X,Y,0*X);
    V = V(:,1:2);
    QT = [Q(:,[3 4]);Q(:,[1 2]);];
    QF = [Q(:,[2 3]); Q(:,[4 1])];
    if ~any(W(:))
      Q = [];
    else
      assert(false,'non trivial W not supported');
    end

  case 3
    side = [size(W,2) size(W,1) size(W,3)];
    if isempty(BC)
      [X,Y,Z] = meshgrid( ...
        linspace(0,1,side(1)), ...
        linspace(0,1,side(2)), ...
        linspace(0,1,side(3)));
      % barycenters of cells
      BC = [X(:) Y(:) Z(:)];
    end

    NV = min(BC);
    XV = max(BC);
    r = (XV-NV)./(side-1);
    [X,Y,Z] = meshgrid( ...
      (NV(1)-0.5*r(1))+linspace(0,1,side(1)+1)*(XV(1)-NV(1)+r(1)), ...
      (NV(2)-0.5*r(2))+linspace(0,1,side(2)+1)*(XV(2)-NV(2)+r(2)), ...
      (NV(3)-0.5*r(3))+linspace(0,1,side(3)+1)*(XV(3)-NV(3)+r(3)));
    % corners of cells
    V = [X(:) Y(:) Z(:)];

    [II,JJ,KK] = ind2sub([side(2) side(1) side(3)]+1,reshape(1:size(V,1),[side(2) side(1) side(3)]+1));

    QF = [ ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end-1,1:end-1,1:end),JJ(1:end-1,1:end-1,1:end),KK(1:end-1,1:end-1,1:end)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(  2:end,1:end-1,1:end),JJ(  2:end,1:end-1,1:end),KK(  2:end,1:end-1,1:end)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(  2:end,  2:end,1:end),JJ(  2:end,  2:end,1:end),KK(  2:end,  2:end,1:end)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end-1,  2:end,1:end),JJ(1:end-1,  2:end,1:end),KK(1:end-1,  2:end,1:end)),[],1) ...
        ];

    QL = fliplr([ ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end-1,1:end,1:end-1),JJ(1:end-1,1:end,1:end-1),KK(1:end-1,1:end,1:end-1)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(  2:end,1:end,1:end-1),JJ(  2:end,1:end,1:end-1),KK(  2:end,1:end,1:end-1)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(  2:end,1:end,  2:end),JJ(  2:end,1:end,  2:end),KK(  2:end,1:end,  2:end)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end-1,1:end,  2:end),JJ(1:end-1,1:end,  2:end),KK(1:end-1,1:end,  2:end)),[],1) ...
        ]);

    QT = [ ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end,1:end-1,1:end-1),JJ(1:end,1:end-1,1:end-1),KK(1:end,1:end-1,1:end-1)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end,  2:end,1:end-1),JJ(1:end,  2:end,1:end-1),KK(1:end,  2:end,1:end-1)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end,  2:end,  2:end),JJ(1:end,  2:end,  2:end),KK(1:end,  2:end,  2:end)),[],1) ...
        reshape(sub2ind([side(2) side(1) side(3)]+1,II(1:end,1:end-1,  2:end),JJ(1:end,1:end-1,  2:end),KK(1:end,1:end-1,  2:end)),[],1) ...
        ];

    if ~any(W(:))
      Q = [];
    else
      Wp = padarray(W,[1 1 1],0);
      Dy = diff(Wp,1,1);
      Dx = diff(Wp,1,2);
      Dz = diff(Wp,1,3);
      Q = [ ...
        QF(        Dz(2:end-1,2:end-1,1:end)>0.5,:); ...
        QL(        Dx(2:end-1,1:end,2:end-1)>0.5,:); ...
        QT(        Dy(1:end,2:end-1,2:end-1)>0.5,:); ...
        fliplr(QF( Dz(2:end-1,2:end-1,1:end)<-0.5,:)); ...
        fliplr(QL( Dx(2:end-1,1:end,2:end-1)<-0.5,:)); ...
        fliplr(QT( Dy(1:end,2:end-1,2:end-1)<-0.5,:)); ...
        ];
    end
  end
end


% [ wave, f, t, coh, phases, raw, coi, scale, period, scalef ] = getWavelet( x, Fs, fMin, fMax, nbins, graphics )
%
% x:
% vector - wavelet analysis is applied
% multi-column matrix - wavelet analysis is applied to each column
% separately
%
% two-column matrix - treated as two time series and the csd, coherence is
% computed
% 3D array, 3rd dimension is 2 - treated as multiple samples from two time
% series, coherence etc is computed
%
% Fs - sampling frequency
% fMin, fMax, nbins - parameters for spectral analysis - min/max freqnecy
% (Hz) and number of frequency bins (Morlet is always used so this is
% frequency per se and not just scale)
% graphcis - flag 0/1
%
% does: compute the CWT of each segment, then if two signals also compute
% the cross-spectrum/coherence/phase lag. then average over all segments
% and plot
%
% output:
% wave - the PSD/CSD
% f, t - vectors
% coh - only for 2-channel input
% phases - either for the 1-channel or the phase-difference (from the CSD,
%       not the smoothed estimate)
%
%
% call: wavelet, smoothwavelet, phaseplot, myjet, colormaps

% 08-oct-12 ES

% revisions
% 15-nov-12 (1) added single-channel phase; cross-spectrum, coherence, and
%               phase difference estimates
%           (2) local plotting of spectrogram + phasogram (single-channel)
%               or coherogram + phase differences (two-channel)

% to do: organize input / output handling better, plotting etc
%           also compute mean spectra/coherence/phase
% also - external coherence computation for multiple segments (i.e. average
%       across trials)

%function [ wave, f, t, phases, coi, scale, raw ] = getWavelet( x, Fs, fMin, fMax, nbins )
function [ wave, f, t, coh, phases, raw, coi, scale, period, scalef ] = getWavelet( x, Fs, fMin, fMax, nbins, scaling, graphics )

if isa( x, 'int16' ), x = single( x ); end

if nargin < 6 || isempty( scaling )
    scaling = 'var';
end
if nargin < 7 || isempty( graphics )
    if nargout == 0
        graphics = 1;
    else
        graphics = 0;
    end
end

dt = 1 / Fs;
s0 = 1 / fMax;
tMax = 1 / fMin;
dj = log2( tMax/s0 ) / nbins;

mother = 'MORLET';
%k0 = 6;
%fourier_factor = (4*pi)/(k0 + sqrt(2 + k0^2)); % scale->frequency

wave = [];
f = [];
t = [];
coh = [];
phases = [];
raw = [];
coi = [];
scale = [];
period = [];
scalef = [];

er = 0;
sx = size( x );
switch ndims( x )
    case 2
        switch min( sx )
            case 1
                nsignals = 1;
                nsegments = 1;
                nsamples = max( sx );
                x = x( : );
            case 2
                nsignals = 2;
                nsegments = 1;
                if size( x, 1 ) == 2
                    x = x';
                end
                y = x( :, 2 );
                x = x( :, 1 );
                nsamples = max( sx );
            otherwise
                nsignals = 1;
                nsegments = sx( 2 );
                nsamples = sx( 1 );
        end
    case 3
        if size( x, 3 ) ~= 2
            er = 1;
        else
            nsignals = 2;
            nsegments = sx( 2 );
            nsamples = sx( 1 );
            y = x( :, :, 2 );
            x = x( :, :, 1 );
        end
    otherwise
        er = 1;
end
if er
    error( 'not supported' )
end

% parameters
if isa( scaling, 'char' ) && strcmp( scaling, 'var' )
    switch nsignals
        case 1
            scalef = var( x ); % can be a different number for each segment
        case 2
            scalef = std( x ) .* std( y );
    end
elseif isa( scaling, 'double' ) && size( scaling, 1 ) == ( nbins + 1 )
    scalef = scaling; % can be a different number for each frequency
    scaling = 'z';
else
    scalef = ones( 1, nsegments );
    scaling = 'none';
end
flipidx = ( nbins + 1 ) : -1 : 1;
t = ( 1 : nsamples )' / Fs;

% actually comptue
%[wave,period,scale,coi] = wavelet(Y,dt,pad,dj,s0,J1,mother,param);
switch nsignals
    case 1
        xw = zeros( nbins + 1, nsamples, nsegments );
        for i = 1 : nsegments
            [ xw( :, :, i ), period, scale, coi ] = wavelet( x( :, i ), dt, 1, dj, s0, nbins, mother );
        end
        xw = xw( flipidx, :, : ); % freq, time, segments
        wave = abs( xw ) .^ 2;
        phases = angle( xw );
        raw = xw;
    case 2
        xw = zeros( nbins + 1, nsamples, nsegments );
        yw = xw;
        coh = xw;
        for i = 1 : nsegments
            [ xw( :, :, i ), period, scale, coi ] = wavelet( x( :, i ), dt, 1, dj, s0, nbins, mother );
            [ yw( :, :, i ) ] = wavelet( y( :, i ), dt, 1, dj, s0, nbins, mother );
            % coherence (copied as is from wtc.m):
            sinv=1./(scale');
            X = xw( :, :, i );
            Y = yw( :, :, i );
            wxy = X .* conj( Y ); % complex, single trial
            sX=smoothwavelet(sinv(:,ones(1,nsamples)).*(abs(X).^2),dt,period,dj,scale);
            sY=smoothwavelet(sinv(:,ones(1,nsamples)).*(abs(Y).^2),dt,period,dj,scale);
            sWxy=smoothwavelet(sinv(:,ones(1,nsamples)).*wxy,dt,period,dj,scale);
            Rsq=abs(sWxy).^2./(sX.*sY);
            %phases( :, :, i ) = angle( sWxy );
            coh( :, :, i ) = Rsq( flipidx, : );
            %coh = abs( yo( :, 1, 2 ) .^ 2 ) ./ ( yo( :, 1, 1 ) .* yo( :, 2, 2 ) );

%             subplot( 4, 2, 1 ), [ c h ] = contourf( t, f, log2( abs( flipud( X ) ).^2 ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 3 ), [ c h ] = contourf( t, f, log2( abs( flipud( Y ) ).^2 ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 5 ), [ c h ] = contourf( t, f, log2( abs( flipud( wxy ) ).^2 ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 2 ), [ c h ] = contourf( t, f, log2( flipud( sX ) ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 4 ), [ c h ] = contourf( t, f, log2( flipud( sY ) ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 6 ), [ c h ] = contourf( t, f, log2( flipud( sWxy ) ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 7 ), [ c h ] = contourf( t, f, flipud( abs(wxy).^2./(X.*Y) ), 100 ); set( h, 'linestyle', 'none' );
%             subplot( 4, 2, 8 ), [ c h ] = contourf( t, f, flipud( abs(sWxy).^2./(sX.*sY) ), 100 ); set( h, 'linestyle', 'none' );
%
%             x0 = (abs(X).^2);
%             fmat = sinv(:,ones(1,nsamples));
%             sX=smoothwavelet(fmat.*x0,dt,period,dj,scale);

        end

        % individual channels
        xw = xw( flipidx, :, : );
        yw = yw( flipidx, :, : );
        raw( :, :, :, 1 ) = xw;
        raw( :, :, :, 2 ) = yw;
        %xwave = abs( xw ) .^ 2;
        %ywave = abs( yw ) .^ 2;
        %xphases = angle( xw );
        %yphases = angle( yw );
        % cross spectrum
        xyw = xw .* conj( yw );
        phases = angle( xyw ); % from the CSD
        wave = abs( xyw );
        % for multiple trials, one can also compute the trial-averaged coherence and phase lag by:
        %cohTA = mean( abs( xyw ) .^ 2, 3 ) ./ ( mean( abs( xw ) .^ 2, 3 ) .* mean( abs( yw ) .^ 2, 3 ) );
        %phasesTA = mod( atan2( mean( sin( phases ), 3 ), mean( cos( phases ), 3 ) ), 2 * pi );

end

%xw = flipud( xw );
f = 1 ./ period( flipidx );
coi = 1 ./ coi; % minimum freq to consider at each time point
%f = fliplr( 1 ./ period );
%f = fliplr( 1 ./ scale );
%t = ( 1 : length( x ) )' / Fs;

if graphics

    % here the scaling is by the signal variance

    figure

    if nsignals == 1
        nplots = 1; % PSD
    else
        nplots = 4; % [ PSD1 CSD; COH PSD2 ]
    end

    for np = 1 : nplots
    end

    % scale
    %scalef = mean( scalef );
    %scaleres = 0.25;
    scaleres = 10;
    scalename = '{\sigma}^2';
    pow = zeros( size( wave ) );
    switch scaling
        case { 'var', 'none' }
            for i = 1 : nsegments
                pow( :, :, i ) = log2( abs( wave( :, :, i ) / scalef( i ) ) );
            end
        case 'z'

            for i = 1 : length( f )
                pow( i, :, : ) = ( wave( i, :, : ) - scalef( i, 1 ) ) / scalef( i, 2 );
            end
    end
    pow = mean( pow, 3 ); % average over segments (1 signal)
    levels = min( pow(:) ) : scaleres : max( pow( : ) );
    mphases = mod( atan2( mean( sin( phases ), 3 ), mean( cos( phases ), 3 ) ), 2 * pi );
    %mp = []; for i = 1 : size( phases, 1 ), mp( :, i ) = circ_mean( squeeze( phases( i, :, : ) )' ); end

    % plot
    h1 = subplot( 1, 1, 1 ); %subplot( 2, 1, 1 );
    %[ c h ] = contourf( t, f, pow, levels );
    [ c h ] = contourf( t, f, pow, 100 );
    set( h, 'linestyle','none')
    %xlabel( 'Time (sec)' )
    ylabel( 'Frequency (Hz)' )
    title( sprintf( '%d segments, %d signals', nsegments, nsignals ) )

    % center color limits around log2(1)=0
    if strcmp( scaling, 'var' )
        clim=get(gca,'clim');
        clim=[-1 1]*max(clim(2),3);
        set(gca,'clim',clim)
    end

    % add the cone of influence
    line( t, coi, 'color', [ 0 0 0 ] );
    hold on
    tt=[t([1 1])-dt*.5;t;t([end end])+dt*.5];
    hcoi=fill(tt,1./[period([end 1]) 1./coi period([1 end])],'w');
    %hcoi=fill(tt,[f([end 1]) coi f([1 end])],'w');
    set(hcoi,'alphadatamapping','direct','facealpha',.5)
    hold off

    set( h1, 'box', 'off', 'tickdir', 'out' )

    % add phase arrows (copied as is from xwt.m)
    if nsignals == 2
        Args.ArrowDensity = [30 30];
        Args.ArrowSize = 1;
        Args.ArrowHeadSize = 1;
        ad=mean(Args.ArrowDensity);
        Args.ArrowSize=Args.ArrowSize*30*.03/ad;
        Args.ArrowHeadSize=Args.ArrowHeadSize*Args.ArrowSize*220;
        phs_dt=round(length(t)/Args.ArrowDensity(1));
        tidx=max(floor(phs_dt/2),1):phs_dt:length(t);
        phs_dp=round(length(period)/Args.ArrowDensity(2));
        pidx=fliplr( max(floor(phs_dp/2),1):phs_dp:length(period) );
        phaseplot(t(tidx),f(pidx),2*pi-mphases(pidx,tidx),Args.ArrowSize,Args.ArrowHeadSize);
    end

    % add colorbar
    h = colorbar;
    subplot( h )
    barylbls=rats(2.^(get(h,'ytick')'));
    %barylbls([1 end],:)=' ';
    barylbls(:,all(barylbls==' ',1))=[];
    set(h,'yticklabel',barylbls);
    title( scalename )
    set( h, 'box', 'off', 'tickdir', 'out' )

    colormap( h1, myjet )

    if 1
        figure, %h2 = subplot( 2, 1, 2 );
        switch nsignals
            case 1
                % plot phases separately
                [ c h ] = contourf( t, f, mphases, 10 );
                set( h, 'linestyle','none')
                xlabel( 'Time (sec)' )
                ylabel( 'Frequency (Hz)' )
                %colormap( h2, colormaps( myjet ) )
                set( gca, 'clim', [ 0 2*pi ] )

                % add the cone of influence
                lh = line( t, coi, 'color', [ 0 0 0 ] );
                hold on
                tt=[t([1 1])-dt*.5;t;t([end end])+dt*.5];
                hcoi=fill(tt,1./[period([end 1]) 1./coi period([1 end])],'w');
                set(hcoi,'alphadatamapping','direct','facealpha',.5)
                hold off

                h = colorbar;
                subplot( h )
                title( 'Phase (rad)' )
                set( h, 'box', 'off', 'tickdir', 'out' )
                colormap( colormaps( myjet ) )

            case 2
                % also plot coherence (without phase arrows)
                % plot
                %h2 = subplot( 2, 1, 2 );
                [ c h ] = contourf( t, f, mean( coh, 3 ), 100 );
                %[ c h ] = contourf( t, f, cohTA, 100 );
                set( h, 'linestyle','none')
                xlabel( 'Time (sec)' )
                ylabel( 'Frequency (Hz)' )
                title( sprintf( '%d segments, %d signals', nsegments, nsignals ) )
                set( gca, 'clim', [ 0 1 ] )
                set( h1, 'box', 'off', 'tickdir', 'out' )

                % add the cone of influence
                lh = line( t, coi, 'color', [ 0 0 0 ] );
                hold on
                tt=[t([1 1])-dt*.5;t;t([end end])+dt*.5];
                hcoi=fill(tt,1./[period([end 1]) 1./coi period([1 end])],'w');
                set(hcoi,'alphadatamapping','direct','facealpha',.5)
                hold off

                % add phase plots (only for high coherence values inside the coi)
                aaa=2*pi-mphases;
                aaa(mean( coh, 3 )<.5)=NaN;
                aaa( bsxfun( @lt, f' * ones( 1, nsamples ), coi ) ) = NaN;
                phaseplot(t(tidx),f(pidx),aaa(pidx,tidx),Args.ArrowSize,Args.ArrowHeadSize);

                % add colorbar
                h = colorbar;
                subplot( h )
                title( 'Coherence' )
                set( h, 'box', 'off', 'tickdir', 'out' )
                colormap( myjet )
        end

    end

    %figure, xwt( [t x( :, 1 )],[t y( :, 1 ) ], 'Pad', 1, 'Dj', dj, 'S0', s0, 'J1', nbins );
    %figure, [ Rsq,aWxy ] = waveletCoherence( [t x( :, 1 )],[t y( :, 1 ) ], 'Pad', 1, 'Dj', dj, 'S0', s0, 'J1', nbins, 'mcc', 0, 'MakeFigure', 1 ); title( 'COH' )
    %figure, [ Wxy ] = xwt( [t x( :, i )],[t y( :, i ) ], 'Pad', 1, 'Dj', log2( fMax/fMin ) / nbins, 'S0', 1/fMax, 'MaxScale', 1/fMin, 'MakeFigure', 1 ); title( 'CSD' )

end

%t = [ 0 : 1 : length( x ) - 1 ]' / Fs;
%[ rawwave, period, scale, coi sig95 ] = wt( [ t x ], 'Pad', 1, 'dj', dj, 's0', s0, 'j1', nbins, 'mother', mother, 'MakeFigure', 1 );

return

% notes:
% (1) it is clear what the frequencies are (nbins, log-spaced between fMin and
% fMax), but the amp is unclear to me           DONE
% (2) should use a similar approach to filter the signal at various
% frequency ranges and calculate the spiking rate frequency/phase maps

% to do:
% (1) adjust the scale properly                 DONE
% (2) get the multi-segment version working; basically use the same call to
% wavelet.m, but concatenate the segments (with intervening portions), then
% call once with the first segment to get the proper coi.       DONE
% (3) get the plotting of phases working on this diagram        DONE
% (4) 2-signal version: compute the coherence as in Torrence and Compo
% (smoothing in time- and frequency-domains each of the spectra and the
% cross-spectrum) and add the phase                             DONE

% OK now (15nov12), but not super elegant:
% should partition into two function - getWavelet and plotWavelet,
% the first should only compute, the second should plot with options:
% single signal: just power,power + phase plots.
%       e.g. [ 1 1 ] would make two plots, whereas [ 1 0 ] just on
% two signals: power/phase for each signal; csd/coh/phase for the joint
%       i.e. if only power, plots 2x2 - [ s1 s12_csd; s12_coh; s2 ]
%       if also phase, cuts each plot into two and adds the phase

getWavelet( [ x0( :, 1 ) y0( :, 1 ) ], Fs, Fmin, Fmax, nBins );
getWavelet( [ x0( :, 1 ) ], Fs, Fmin, Fmax, nBins );
getWavelet( [ x0( :, 1 : 10 ) ], Fs, Fmin, Fmax, nBins );
xy = [];
xy( :, :, 1 ) = x0( :, 1 : 10 );
xy( :, :, 2 ) = y0( :, 1 : 10 );
getWavelet( xy, Fs, Fmin, Fmax, nBins );


function soln = rungeKutta(problem)
% soln = rungeKutta(problem)
%
% This function transcribes a trajectory optimization problem using the
% multiple shooting, with 4th-order Runge Kutta integration
%
% See Bett's book for details on the method
%
% For details on the input and output, see the help file for optimTraj.m
%
% Method specific parameters:
%
%   problem.options.method = 'rungeKutta'
%   problem.options.rungeKutta = struct with method parameters:
%       .nSegment = number of trajectory segments
%       .nSubStep = number of sub-steps to use in each segment
%       .adaptiveDerivativeCheck = 'off' by default. Set to 'on' to enable
%           numerical checks on the analytic gradients, computed using the
%           derivest package, rather than fmincon's internal checks.
%           Derivest is slower, but more accurate than fmincon. Derivest
%           can be downloaded from the Mathworks File Exchange, file id of
%           13490 - Adaptive Robust Numerical Differentation, John D-Errico
%
%
% NOTES:
%
%   Code for computing analyic gradients of the Runge Kutta method was
%   contributed by Will Wehner.
%
%   If analytic gradients are used, then the sparsity pattern is returned
%   in the struct: soln.info.sparsityPattern. View it using spy().
%
%

%To make code more readable
G = problem.guess;
B = problem.bounds;
F = problem.func;
Opt = problem.options;

% Figure out grid size:
nSegment = Opt.rungeKutta.nSegment;
nSubStep = Opt.rungeKutta.nSubStep;
nGridControl = 2*nSegment*nSubStep + 1;
nGridState = nSegment + 1;

% Print out some solver info if desired:
if Opt.verbose > 0
    fprintf('  -> Transcription via 4th-order Runge-Kutta method \n');
    fprintf('        nSegments = %d \n', nSegment);
    fprintf('        nSubSteps = %d \n', nSubStep);
end

% Interpolate the guess at the transcription grid points for initial guess:
guess.tSpan = G.time([1,end]);
guess.tState = linspace(guess.tSpan(1), guess.tSpan(2), nGridState);
guess.tControl = linspace(guess.tSpan(1), guess.tSpan(2), nGridControl);
guess.state = interp1(G.time', G.state', guess.tState')';
guess.control = interp1(G.time', G.control', guess.tControl')';
[zGuess, pack] = packDecVar(guess.tSpan, guess.state, guess.control);

% Unpack all bounds:
tLow = [B.initialTime.low, B.finalTime.low];
xLow = [B.initialState.low, B.state.low*ones(1,nGridState-2), B.finalState.low];
uLow = B.control.low*ones(1,nGridControl);
zLow = packDecVar(tLow,xLow,uLow);

tUpp = [B.initialTime.upp, B.finalTime.upp];
xUpp = [B.initialState.upp, B.state.upp*ones(1,nGridState-2), B.finalState.upp];
uUpp = B.control.upp*ones(1,nGridControl);
zUpp = packDecVar(tUpp,xUpp,uUpp);

%%%% Set up problem for fmincon:
flagGradObj = strcmp(Opt.nlpOpt.GradObj,'on');
flagGradCst = strcmp(Opt.nlpOpt.GradConstr,'on');
if flagGradObj || flagGradCst
    gradInfo = grad_computeInfo(pack);
end
if flagGradObj
    P.objective = @(z)( ...
        myObjGrad(z, pack, F.dynamics, F.pathObj, F.bndObj, gradInfo) );   %Analytic gradients
    [~, objGradInit] = P.objective(zGuess);
    sparsityPattern.objective = (objGradInit~=0)';
else
    P.objective = @(z)( ...
        myObjective(z, pack, F.dynamics, F.pathObj, F.bndObj) );   %Numerical gradients
end

if flagGradCst
    P.nonlcon = @(z)( ...
        myCstGrad(z, pack, F.dynamics, F.pathObj, F.pathCst, F.bndCst, gradInfo) ); %Analytic gradients
    [~,~,cstIneqInit,cstEqInit] = P.nonlcon(zGuess);
    sparsityPattern.equalityConstraint = (cstEqInit~=0)';
    sparsityPattern.inequalityConstraint = (cstIneqInit~=0)';
else
    P.nonlcon = @(z)( ...
        myConstraint(z, pack, F.dynamics, F.pathObj, F.pathCst, F.bndCst) ); %Numerical gradients
end

% Check analytic gradients with DERIVEST package
if strcmp(Opt.rungeKutta.adaptiveDerivativeCheck,'on')
    if exist('jacobianest','file')
        runGradientCheck(zGuess, pack,F.dynamics, F.pathObj, F.bndObj, F.pathCst, F.bndCst, gradInfo);
        Opt.nlpOpt.DerivativeCheck = [];  %Disable built-in derivative check
    else
        Opt.rungeKutta.adaptiveDerivativeCheck = 'cannot find jacobianest.m';
        disp('Warning: the derivest package is not on search path.');
        disp(' --> Using fmincon''s built-in derivative checks.');
    end
end

% Build the standard fmincon problem struct
P.x0 = zGuess;
P.lb = zLow;
P.ub = zUpp;
P.Aineq = []; P.bineq = [];
P.Aeq = []; P.beq = [];
P.solver = 'fmincon';
P.options = Opt.nlpOpt;

%%%% Call fmincon to solve the non-linear program (NLP)
tic;
[zSoln, objVal,exitFlag,output] = fmincon(P);
[tSpan,~,uSoln] = unPackDecVar(zSoln,pack);
nlpTime = toc;

%%%% Store the results:
[tGrid,xGrid,uGrid] = simulateSystem(zSoln, pack, F.dynamics, F.pathObj);
soln.grid.time = tGrid;
soln.grid.state = xGrid;
soln.grid.control = uGrid;

% Quadratic interpolation over each sub-step for the control:
tSoln = linspace(tSpan(1),tSpan(2),nGridControl);
soln.interp.control = @(t)( interp1(tSoln', uSoln', t','pchip')' );

% Cubic spline representation of the state over each substep:
dxGrid = F.dynamics(tGrid,xGrid,uGrid);
xSpline = pwch(tGrid, xGrid, dxGrid);
soln.interp.state = @(t)( ppval(xSpline,t) );

% General information about the optimization run
soln.info = output;
soln.info.nlpTime = nlpTime;
soln.info.exitFlag = exitFlag;
soln.info.objVal = objVal;
if flagGradCst || flagGradObj
    soln.info.sparsityPattern = sparsityPattern;
end

soln.problem = problem;  % Return the fully detailed problem struct

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%
%%%%                   SUB FUNCTIONS                                   %%%%
%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [decVars,pack] = packDecVar(tSpan,state,control)
%
% This function collapses the time (t), state (x)
% and control (u) matricies into a single vector
%
% INPUTS:
%   tSpan = [1, 2] = time bounds
%   state = [nState, nGridState] = state vector at each grid point
%   control = [nControl, nGridControl] = control vector at each grid point
%
% OUTPUTS:
%   decVars = column vector of 2 + nState*nGridState + nControl*nGridControl) decision variables
%   pack = details about how to convert z back into t,x, and u
%       .nState
%       .nGridState
%       .nControl
%       .nGridControl
%
% NOTES:
% nGridControl = 2*nSegment*nSubStep + 1;
% nGridState = nSegment + 1;
%

[nState, nGridState] = size(state);
[nControl, nGridControl] = size(control);

nSegment = nGridState - 1;
nSubStep = (nGridControl - 1)/(2*nSegment);

xCol = reshape(state, nState*nGridState, 1);
uCol = reshape(control, nControl*nGridControl, 1);

indz = 1:numel(control)+numel(state)+numel(tSpan);

% index of time in decVar
indt = 1:2;

% the z index of the first element of each state over time
indtemp = 2 + (1 : (nState + (2*nSubStep)*nControl ) : numel(control)+numel(state));

% remaining state elements at each time
indx = repmat(indtemp,nState,1) + cumsum(ones(nState,nGridState),1) - 1;

% index of control in decVar
indu = indz;
indu([indt(:);indx(:)])=[];
indu = reshape(indu,nControl,nGridControl);

% pack up decVars
decVars = zeros(numel(indz),1);
decVars(indt(:),1) = tSpan;
decVars(indx(:),1) = xCol;
decVars(indu(:),1) = uCol;

% pack structure
pack.nState = nState;
pack.nGridState = nGridState;
pack.nControl = nControl;
pack.nGridControl = nGridControl;
pack.nSegment = nGridState - 1;
pack.nSubStep = (nGridControl-1)/(2*pack.nSegment);
pack.indt = indt;
pack.indx = indx;
pack.indu = indu;

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [tSpan, state, control] = unPackDecVar(decVars,pack)
%
% This function unpacks the decision variables for
% trajectory optimization into the time (t),
% state (x), and control (u) matricies
%
% INPUTS:
%   decVars = column vector of 2 + nState*nGridState + nControl*nGridControl) decision variables
%   pack = details about how to convert z back into t,x, and u
%       .nState
%       .nGridState
%       .nControl
%       .nGridControl
%
% OUTPUTS:
%   tSpan = [1, 2] = time bounds
%   state = [nState, nGridState] = state vector at each grid point
%   control = [nControl, nGridControl] = control vector at each grid point
%

tSpan = [decVars(1),decVars(2)];

% state = reshape(decVars((2+1):(2+nx)), pack.nState, pack.nGridState);
% control = reshape(decVars((2+nx+1):(2+nx+nu)), pack.nControl, pack.nGridControl);

state = decVars(pack.indx);
control = decVars(pack.indu);

% make sure x and u are returned as vectors, [nState,nTime] and
% [nControl,nTime]
state = reshape(state,pack.nState,pack.nGridState);
control = reshape(control,pack.nControl,pack.nGridControl);
end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function cost = myObjective(decVars, pack,dynamics, pathObj, bndObj)
%
% This function unpacks the decision variables, sends them to the
% user-defined objective functions, and then returns the final cost
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%   bndObj = user-defined boundary objective function
%
% OUTPUTS:
%   cost = scalar cost for this set of decision variables
%
%

% All of the real work happens inside this function:
[t,x,~,~,pathCost] = simulateSystem(decVars, pack, dynamics, pathObj);

% Compute the cost at the boundaries of the trajectory
if isempty(bndObj)
    bndCost = 0;
else
    t0 = t(1);
    tF = t(end);
    x0 = x(:,1);
    xF = x(:,end);
    bndCost = bndObj(t0,x0,tF,xF);
end

cost = bndCost + pathCost;

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [c, ceq] = myConstraint(decVars, pack, dynamics, pathObj, pathCst, bndCst)
%
% This function unpacks the decision variables, computes the defects along
% the trajectory, and then evaluates the user-defined constraint functions.
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%   pathCst = user-defined path-constraint function
%   bndCst = user-defined boundary constraint function
%
% OUTPUTS:
%   c = non-linear inequality constraint
%   ceq = non-linear equatlity cosntraint
%
% NOTE:
%   - path constraints are  satisfied at the start and end of each sub-step
%

[t,x,u,defects] = simulateSystem(decVars, pack, dynamics, pathObj);

%%%% Call user-defined constraints and pack up:
[c, ceq] = collectConstraints(t,x,u,...
    defects,...
    pathCst, bndCst);

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [t,x,u,defects,pathCost] = simulateSystem(decVars, pack, dynFun, pathObj)
%
% This function does the real work of the transcription method. It
% simulates the system forward in time across each segment of the
% trajectory, computes the integral of the cost function, and then matches
% up the defects between the end of each segment and the start of the next.
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%
% OUTPUTS:
%   t = [1 x nGrid] = time vector for the edges of the sub-step grid
%   x = [nState x nGrid] = state vector
%   u = [nControl x nGrid] = control vector
%   defects = [nState x nSegment] = defect matrix
%   pathCost = scalar cost for the path integral
%
% NOTES:
%   - nGrid = nSegment*nSubStep+1
%   - This function is usually called twice for each combination of
%   decision variables: once by the objective function and once by the
%   constraint function. To keep the code fast I cache the old values and
%   only recompute when the inputs change.
%

%%%% CODE OPTIMIZATION %%%%
%
% Prevents the same exact code from being called twice by caching the
% solution and reusing it when appropriate.
%
global RUNGE_KUTTA_t RUNGE_KUTTA_x RUNGE_KUTTA_u
global RUNGE_KUTTA_defects RUNGE_KUTTA_pathCost
global RUNGE_KUTTA_decVars
%
usePreviousValues = false;
if ~isempty(RUNGE_KUTTA_decVars)
    if length(RUNGE_KUTTA_decVars) == length(decVars)
        if ~any(RUNGE_KUTTA_decVars ~= decVars)
            usePreviousValues = true;
        end
    end
end
%
if usePreviousValues
    t = RUNGE_KUTTA_t;
    x = RUNGE_KUTTA_x;
    u = RUNGE_KUTTA_u;
    defects = RUNGE_KUTTA_defects;
    pathCost = RUNGE_KUTTA_pathCost;
else
    %
    %
    %%%% END CODE OPTIMIZATION %%%%

    [tSpan, state, control] = unPackDecVar(decVars,pack);

    nState = pack.nState;
    nSegment = pack.nSegment;
    nSubStep = pack.nSubStep;

    % NOTES:
    %   The following bit of code is a bit confusing, mostly due to the
    %   need for vectorization to make things run at a reasonable speed in
    %   Matlab. Part of the confusion comes because the decision variables
    %   include the state at the beginning of each segment, but the control
    %   at the beginning and middle of each substep - thus there are more
    %   control grid-points than state grid points. The calculations are
    %   vectorized over segments, but not sub-steps, since the result of
    %   one sub-step is required for the next.

    % time, state, and control at the ends of each substep
    nTime = 1+nSegment*nSubStep;
    t = linspace(tSpan(1), tSpan(2), nTime);
    x = zeros(nState, nTime);
    u = control(:,1:2:end); % Control a the endpoints of each segment
    uMid = control(:,2:2:end);  %Control at the mid-points of each segment
    c = zeros(1, nTime-1);  %Integral cost for each segment
    dt = (t(end)-t(1))/(nTime-1);

    idx = 1:nSubStep:(nTime-1);   %Indicies for the start of each segment
    x(:,[idx,end]) = state;   %Fill in the states that we already know

    for iSubStep = 1:nSubStep
        % March forward Runge-Kutta step

        t0 = t(idx);
        x0 = x(:,idx);

        k0 = combinedDynamics(t0,        x0,                         u(:,idx), dynFun,pathObj);
        k1 = combinedDynamics(t0+0.5*dt, x0 + 0.5*dt*k0(1:nState,:), uMid(:,idx), dynFun,pathObj);
        k2 = combinedDynamics(t0+0.5*dt, x0 + 0.5*dt*k1(1:nState,:), uMid(:,idx), dynFun,pathObj);
        k3 = combinedDynamics(t0+dt,     x0 +     dt*k2(1:nState,:), u(:,idx+1), dynFun,pathObj);
        z = (dt/6)*(k0 + 2*k1 + 2*k2 + k3);  %Change over the sub-step

        xNext = x0 + z(1:nState,:);  %Next state
        c(idx) = z(end,:);  %Integral of the cost function over this step

        if iSubStep == nSubStep %We've reached the end of the interval
            % Compute the defect vector:
            defects = xNext - x(:,idx+1);
        else
            % Store the state for next step in time
            idx = idx+1;   %  <-- This is important!!
            x(:,idx) = xNext;
        end

    end

    pathCost = sum(c);  %Sum up the integral cost over each segment

    %%%% Cache results to use on the next call to this function.
    RUNGE_KUTTA_t = t;
    RUNGE_KUTTA_x = x;
    RUNGE_KUTTA_u = u;
    RUNGE_KUTTA_defects = defects;
    RUNGE_KUTTA_pathCost = pathCost;

end

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function dz = combinedDynamics(t,x,u,dynFun,pathObj)
% dz = combinedDynamics(t,x,u,dynFun,pathObj)
%
% This function packages the dynamics and the cost function together so
% that they can be integrated at the same time.
%
% INPUTS:
%   t = [1, nTime] = time vector (grid points)
%   x = [nState, nTime] = state vector at each grid point
%   u = [nControl, nTime] = control vector at each grid point
%   dynamics(t,x,u) = dynamics function handle
%               dx = [nState, nTime] = dx/dt = derivative of state wrt time
%   pathObj(t,x,u) = integral cost function handle
%                 dObj = [1, nTime] = integrand from the cost function
%
% OUTPUTS:
%   dz = [dx; dObj] = combined dynamics of state and cost

dx = dynFun(t,x,u);
if isempty(pathObj)
    dc = zeros(size(t));
else
    dc = pathObj(t,x,u);
end

dz = [dx;dc];  %Combine and return

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%
%%%%                 Analytic Gradient Stuff                           %%%%
%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function gradInfo = grad_computeInfo(pack)
%
% This function computes the matrix dimensions and indicies that are used
% to map the gradients from the user functions to the gradients needed by
% fmincon. The key difference is that the gradients in the user functions
% are with respect to their input (t,x,u) or (t0,x0,tF,xF), while the
% gradients for fmincon are with respect to all decision variables.
%
% INPUTS:
%   nDeVar = number of decision variables
%   pack = details about packing and unpacking the decision variables
%       .nTime
%       .nState
%       .nControl
%
% OUTPUTS:
%   gradInfo = details about how to transform gradients
%

%nTime = pack.nTime;
nState = pack.nState;
nGridState = pack.nGridState;
nControl = pack.nControl;
nGridControl = pack.nGridControl;

nDecVar = 2 + nState*nGridState + nControl*nGridControl;

zIdx = 1:nDecVar;
gradInfo.nDecVar = nDecVar;
[tIdx, xIdx, uIdx] = unPackDecVar(zIdx,pack);

gradInfo.tIdx = tIdx([1,end]);
gradInfo.xIdx = xIdx;
gradInfo.uIdx = uIdx;

nSegment = pack.nSegment;
nSubStep = pack.nSubStep;

% indices of decVars associated with u
indu = 1:2:(1+2*nSegment*nSubStep);
gradInfo.indu = uIdx(:,indu);
% indices of decVars associated with uMid
indumid = 2:2:(1+2*nSegment*nSubStep);
gradInfo.indumid = uIdx(:,indumid);

%%%% For unpacking the boundary constraints and objective:
gradInfo.bndIdxMap = [tIdx(1); xIdx(:,1); tIdx(end); xIdx(:,end)];

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [fail] = runGradientCheck(z_test, pack,dynamics, pathObj, bndObj, pathCst, bndCst, gradInfo)
%
% This function tests the analytic gradients of the objective and
% nonlinear constraints with the DERIVEST package. The finite difference
% calculations in matlab's optimization package were not sufficiently
% accurate.
%
GradientCheckTol = 1e-6;  %Analytic gradients must match numerical within this bound

fail = 0;

fprintf('\n%s\n','____________________________________________________________')
fprintf('%s\n','  DerivativeCheck Information with DERIVEST Package ')

% analytic gradient
[~, dcost] = myObjGrad(z_test, pack, dynamics, pathObj, bndObj, gradInfo);

% check gradient with derivest package
deriv = gradest(@(z) myObjGrad(z, pack, dynamics, pathObj, bndObj, gradInfo),z_test);

% print largest difference in numerical and analytic gradients
fprintf('\n%s\n','Objective function derivatives:')
fprintf('%s\n','Maximum relative difference between user-supplied')
fprintf('%s %1.5e \n','and finite-difference derivatives = ',max(abs(dcost-deriv')))
if any(abs(dcost-deriv') > GradientCheckTol)
    error('Objective gradient did not pass')
end

% analytic nonlinear constraints
[c, ceq,dc, dceq] = myCstGrad(z_test, pack, dynamics, pathObj, pathCst, bndCst, gradInfo);

% check nonlinear inequality constraints with 'jacobianest'
if ~isempty(c)
    jac = jacobianest(@(z) myConstraint(z, pack, dynamics, pathObj, pathCst, bndCst),z_test);

    % print largest difference in numerical and analytic gradients
    fprintf('\n%s\n','Nonlinear inequality constraint function derivatives:')
    fprintf('%s\n','Maximum relative difference between user-supplied')
    fprintf('%s %1.5e \n','and finite-difference derivatives = ',max(max(abs(dc-jac'))))
    if any(any(abs(dc - jac') > GradientCheckTol))
        error('Nonlinear inequality constraint did not pass')
    end
end

% check nonlinear equality constraints with 'jacobianest'
if ~isempty(ceq)
    jac = jacobianest(@(z) myCstGradCheckEq(z, pack, dynamics, pathObj, pathCst, bndCst),z_test);

    % print largest difference in numerical and analytic gradients
    fprintf('\n%s\n','Nonlinear equality constraint function derivatives:')
    fprintf('%s\n','Maximum relative difference between user-supplied')
    fprintf('%s %1.5e \n','and finite-difference derivatives = ',max(max(abs(dceq-jac'))))
    if any(any(abs(dceq - jac') > GradientCheckTol))
        error('Nonlinear equality constraint did not pass')
    end
end

fprintf('\n%s\n','DerivativeCheck successfully passed.')
fprintf('%s\n','____________________________________________________________')
end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function ceq = myCstGradCheckEq(decVars, pack, dynamics, pathObj, pathCst, bndCst)
% This function is necessary for runGradientCheck function
% return only equality constraint (ceq) for use with jacobest.m

[t,x,u,defects] = simulateSystem(decVars, pack, dynamics, pathObj);

%%%% Call user-defined constraints and pack up:
[~, ceq] = collectConstraints(t,x,u,...
    defects,...
    pathCst, bndCst);

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [cost, dcost] = myObjGrad(decVars, pack,dynamics, pathObj, bndObj, gradInfo)
%
% This function unpacks the decision variables, sends them to the
% user-defined objective functions, and then returns the final cost
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%   bndObj = user-defined boundary objective function
%   gradInfo =
%
% OUTPUTS:
%   cost = scalar cost for this set of decision variables
%   dcost = gradient of cost
%     NOTE: gradients are only available for pathCost that depends only on
%     input parameters not states.
%
%

% All of the real work happens inside this function:
[t,x,~,~,pathCost,dxdalpha,dJdalpha] = simSysGrad(decVars, pack, dynamics, pathObj, gradInfo); %#ok<ASGLU>
% dxdalpha is included in outputs to make sure subsequent calls to
% simulateSystem without change a to decVars have access to the correct value
% of dxdalpha - see simulateSystem in which dxdalpha is not calculated unless
% nargout > 5

% Compute the cost at the boundaries of the trajectory
if isempty(bndObj)
    bndCost = 0;
else
    t0 = t(1);
    tF = t(end);
    x0 = x(:,1);
    xF = x(:,end);
    bndCost = bndObj(t0,x0,tF,xF);
end

cost = pathCost + bndCost;

% calculate gradient of cost function
if nargout > 1

    nState = pack.nState;
    nControl = pack.nControl;
    nSegment = pack.nSegment;
    nSubStep = pack.nSubStep;
    nDecVar = 2+nState*(1+nSegment)+nControl*(1+nSegment*nSubStep*2);

    % allocate gradient of cost
    dcost_pth = zeros(nDecVar,1);
    dcost_bnd = zeros(nDecVar,1);

    % gradient assocated with bound objective
    if ~isempty(bndObj)

        % bound costs and gradients w.r.t. t0, x0, tF, xF
        [~, d_bnd] = bndObj(t0,x0,tF,xF);

        % gradients of t0, x0, tF, xF w.r.t. decision parameters (labeled alpha)
        dt0_dalpha = zeros(1,nDecVar);
        dt0_dalpha(1) = 1; % t0 is always the first decVar
        %
        dx0_dalpha = zeros(nState,nDecVar);
        dx0_dalpha(1:nState,gradInfo.xIdx(:,end)) = eye(nState);
        %
        dtF_dalpha = zeros(1,nDecVar);
        dtF_dalpha(2) = 1; % tF is always the second decVar
        %
        dxF_dalpha = zeros(nState,nDecVar);
        dxF_dalpha(1:nState,gradInfo.xIdx(:,end)) = eye(nState);

        % gradient of bound cost
        dcost_bnd(:) = [dt0_dalpha; dx0_dalpha; dtF_dalpha; dxF_dalpha]' * d_bnd';
    end

    % gradient assocated with path objective
    if ~isempty(pathObj)

        dcost_pth = dJdalpha';

    end

    dcost = dcost_pth + dcost_bnd;

end

end

function [c, ceq, dc, dceq] = myCstGrad(decVars, pack, dynamics, pathObj, pathCst, bndCst, gradInfo)
%
% This function unpacks the decision variables, computes the defects along
% the trajectory, and then evaluates the user-defined constraint functions.
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%   pathCst = user-defined path-constraint function
%   bndCst = user-defined boundary constraint function
%   gradInfo =
%
% OUTPUTS:
%   c = non-linear inequality constraint
%   ceq = non-linear equatlity cosntraint
%   dc = gradient of c w.r.t. decVars
%   dceq = gradient of ceq w.r.t. decVars
%
% NOTE:
%   - path constraints are  satisfied at the start and end of each sub-step
%

[t,x,u,defects,pathcost,dxdalpha] = simSysGrad(decVars, pack, dynamics, pathObj, gradInfo); %#ok<ASGLU>

%%%% Call user-defined constraints and pack up:
if nargout <= 2
    [c, ceq] = collectConstraints(t,x,u,...
        defects,...
        pathCst, bndCst);

else

    [c, ceq, dc, dceq] = collectConstraintsGrad(t,x,u,...
        defects,...
        pathCst, bndCst, pack, gradInfo, dxdalpha);

end

end

function [c, ceq, dc, dceq] = collectConstraintsGrad(t,x,u,defects, pathCst, bndCst, pack, gradInfo, dxdalpha)
% [c, ceq, dc, dceq] = collectConstraints(t,x,u,defects, pathCst, bndCst, pack, gradInfo, dxdalpha)
%
% OptimTraj utility function.
%
% Collects the defects, calls user-defined constraints, and then packs
% everything up into a form that is good for fmincon.
%
% INPUTS:
%   t = time vector (time at each substep) nTime = 1+nSegment*nSubStep
%   x = state matrix (states at each time in t)
%   u = control matrix (control at each time in t)
%   defects = defects matrix
%   pathCst = user-defined path constraint function
%   bndCst = user-defined boundary constraint function
%   pack =
%   gradInfo =
%   dxdalpha = partial derivative of state at each substep w.r.t. decVars
%
% OUTPUTS:
%   c = inequality constraint for fmincon
%   ceq = equality constraint for fmincon
%   dc = gradient of c w.r.t. decVars
%   dceq = gradient of ceq w.r.t. decVars
%

% problem dimensions
nState = pack.nState;
nControl = pack.nControl;
nSegment = pack.nSegment;
nSubStep = pack.nSubStep;
nDecVar = 2+nState*(1+nSegment)+nControl*(1+nSegment*nSubStep*2);

%%%% defect constraints
ceq_dyn = reshape(defects,numel(defects),1);

dceq_dyn = zeros(nDecVar,length(ceq_dyn));
Inx = eye(nState);
for j = 1:nSegment
    rows = gradInfo.xIdx(:,j+1);
    cols = (j-1)*nState+(1:nState);
    dceq_dyn(:,cols) = dxdalpha{j}(:,:,end)';  % gradient w.r.t. to x_i(+)
    dceq_dyn(rows,cols) = -Inx; % gradient w.r.t. to x_i
end

%%%% Compute the user-defined constraints:

%%%% path constraints
if isempty(pathCst)
    c_path = [];
    ceq_path = [];
    dc_path = [];
    dceq_path = [];
else
    [c_pathRaw, ceq_pathRaw, c_pathGradRaw, ceq_pathGradRaw] = pathCst(t,x,u);
    c_path = reshape(c_pathRaw,numel(c_pathRaw),1);
    ceq_path = reshape(ceq_pathRaw,numel(ceq_pathRaw),1);

    dc_path = zeros(nDecVar,length(c_path));
    dceq_path = zeros(nDecVar,length(ceq_path));

    % dt/dalpha : gradient of time w.r.t. decVars
    dt_dalpha = zeros(1,nDecVar);
    nTime = 1+nSegment*nSubStep;
    n_time = 0:nTime-1;

    % gradients of path constraints
    nc = size(c_pathRaw,1); % number path constraints at each time
    nceq = size(ceq_pathRaw,1);
    for j = 1:(nSegment+1)
        for i = 1:nSubStep

            % d(t[n])/dalpha
            n_time0 = n_time((j-1)*nSubStep+i);
            dt_dalpha(1) = (1 - n_time0/(nTime-1));
            dt_dalpha(2) = (n_time0/(nTime-1));

            %
            if j < nSegment+1
                dxi_dalpha = dxdalpha{j}(:,:,i);
            else
                dxi_dalpha = zeros(nState,nDecVar);
                cols = gradInfo.xIdx(:,j);
                dxi_dalpha(:,cols) = eye(nState);
            end

            %
            dui_dalpha = zeros(nControl,nDecVar);
            cols = gradInfo.indu(:,(j-1)*nSubStep+i);
            dui_dalpha(:,cols) = eye(nControl);

            % inequality path constraints
            if nc > 0
                cols = (1:nc) + nc*((j-1)*nSubStep+i-1);
                dc_path(:,cols) = [dt_dalpha; dxi_dalpha; dui_dalpha]' * c_pathGradRaw(:,:,nSubStep*(j-1)+i)';
            end

            % equality path constraints
            if nceq > 0
                cols = (1:nceq) + nceq*((j-1)*nSubStep+i-1);
                dceq_path(:,cols) = [dt_dalpha; dxi_dalpha; dui_dalpha]' * ceq_pathGradRaw(:,:,nSubStep*(j-1)+i)';
            end

            % no need to continue with inner loop.
            if j == nSegment+1
                break;
            end
        end
    end

end

%%%% bound constraints
if isempty(bndCst)
    c_bnd = [];
    ceq_bnd = [];
    dc_bnd = [];
    dceq_bnd = [];

else
    t0 = t(1);
    tF = t(end);
    x0 = x(:,1);
    xF = x(:,end);

    % bound constraints and gradients w.r.t. t0, x0, tF, xF
    [c_bnd, ceq_bnd, d_bnd, deq_bnd] = bndCst(t0,x0,tF,xF);

    % gradients of t0, x0, tF, xF w.r.t. decision parameters (labeled alpha)
    dt0_dalpha = zeros(1,nDecVar);
    dt0_dalpha(1) = 1; % t0 is always the first decVar
    %
    dx0_dalpha = zeros(nState,nDecVar);
    cols = gradInfo.xIdx(:,1);
    dx0_dalpha(1:nState,cols) = eye(nState);
    %
    dtF_dalpha = zeros(1,nDecVar);
    dtF_dalpha(2) = 1; % tF is always the second decVar
    %
    dxF_dalpha = zeros(nState,nDecVar);
    cols = gradInfo.xIdx(:,end);
    dxF_dalpha(1:nState,cols) = eye(nState);

    % inequality bound constraints
    dc_bnd = [dt0_dalpha; dx0_dalpha; dtF_dalpha; dxF_dalpha]' * d_bnd';

    % equality bound constraints
    dceq_bnd = [dt0_dalpha; dx0_dalpha; dtF_dalpha; dxF_dalpha]' * deq_bnd';

end

%%%% Pack everything up:
c = [c_path;c_bnd];
ceq = [ceq_dyn; ceq_path; ceq_bnd];

dc = [dc_path, dc_bnd];
dceq = [dceq_dyn, dceq_path, dceq_bnd];

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [t,x,u,defects,pathCost,dxdalpha,dJdalpha] = simSysGrad(decVars, pack, dynFun, pathObj, gradInfo)
%
% This function does the real work of the transcription method. It
% simulates the system forward in time across each segment of the
% trajectory, computes the integral of the cost function, and then matches
% up the defects between the end of each segment and the start of the next.
%
% INPUTS:
%   decVars = column vector of decision variables
%   pack = details about how to convert decision variables into t,x, and u
%   dynamics = user-defined dynamics function handle
%   pathObj = user-defined path-objective function
%
% OUTPUTS:
%   t = [1 x nGrid] = time vector for the edges of the sub-step grid
%   x = [nState x nGrid] = state vector
%   u = [nControl x nGrid] = control vector
%   defects = [nState x nSegment] = defect matrix
%   pathCost = scalar cost for the path integral
%
% NOTES:
%   - nGrid = nSegment*nSubStep+1
%   - This function is usually called twice for each combination of
%   decision variables: once by the objective function and once by the
%   constraint function. To keep the code fast I cache the old values and
%   only recompute when the inputs change.
%

%%%% CODE OPTIMIZATION %%%%
%
% Prevents the same exact code from being called twice by caching the
% solution and reusing it when appropriate.
%
global RUNGE_KUTTA_t RUNGE_KUTTA_x RUNGE_KUTTA_u
global RUNGE_KUTTA_defects RUNGE_KUTTA_pathCost
global RUNGE_KUTTA_decVars RUNGE_KUTTA_dxdalpha RUNGE_KUTTA_dJdalpha
%
usePreviousValues = false;
if ~isempty(RUNGE_KUTTA_decVars)
    if length(RUNGE_KUTTA_decVars) == length(decVars)
        if ~any(RUNGE_KUTTA_decVars ~= decVars)
            usePreviousValues = true;
        end
    end
end
%
if usePreviousValues
    t = RUNGE_KUTTA_t;
    x = RUNGE_KUTTA_x;
    u = RUNGE_KUTTA_u;
    defects = RUNGE_KUTTA_defects;
    pathCost = RUNGE_KUTTA_pathCost;
    dxdalpha = RUNGE_KUTTA_dxdalpha;
    dJdalpha = RUNGE_KUTTA_dJdalpha;
else
    %
    %
    %%%% END CODE OPTIMIZATION %%%%

    [tSpan, state, control] = unPackDecVar(decVars,pack);

    nState = pack.nState;
    nControl = pack.nControl;
    nSegment = pack.nSegment;
    nSubStep = pack.nSubStep;

    % NOTES:
    %   The following bit of code is a bit confusing, mostly due to the
    %   need for vectorization to make things run at a reasonable speed in
    %   Matlab. Part of the confusion comes because the decision variables
    %   include the state at the beginning of each segment, but the control
    %   at the beginning and middle of each substep - thus there are more
    %   control grid-points than state grid points. The calculations are
    %   vectorized over segments, but not sub-steps, since the result of
    %   one sub-step is required for the next.

    % time, state, and control at the ends of each substep
    nTime = 1+nSegment*nSubStep;
    t = linspace(tSpan(1), tSpan(2), nTime);
    x = zeros(nState, nTime);
    u = control(:,1:2:end); % Control a the endpoints of each segment
    uMid = control(:,2:2:end);  %Control at the mid-points of each segment
    c = zeros(1, nTime-1);  %Integral cost for each segment
    dt = (t(end)-t(1))/(nTime-1);

    idx = 1:nSubStep:(nTime-1);   %Indicies for the start of each segment
    x(:,[idx,end]) = state;   %Fill in the states that we already know

    % VARIABLES for analytic gradient evaluations.
    % size of decicion parameters (2 for time), nstate*(nSegment+1), ...
    % dxdalpha = partial derivative of state w.r.t. decVars (alpha)
    nalpha = 2 + nState*(1+nSegment) + nControl*(1+2*nSubStep*nSegment);
    dxdalpha = cell(1,nSegment);
    for i = 1:nSegment
        dxdalpha{i} = zeros(nState,nalpha,nSubStep+1);
        cols = gradInfo.xIdx(:,i);
        dxdalpha{i}(:,cols,1) = eye(nState);
    end
    dTdalpha = zeros(1,nalpha); dTdalpha(1:2) = [-1,1];
    dt_dalpha = zeros(1,nalpha);
    n_time = 0:nTime-1;

    % gradient of path cost
    dJdalpha = zeros(1,nalpha);

    for iSubStep = 1:nSubStep
        % March forward Runge-Kutta step

        t0 = t(idx);
        x0 = x(:,idx);

        %------------------------------------------
        % Code for calculating dxdalpha (partial derivative of state w.r.t.
        % the descision parameters): dxdalpha = nstate x nalpha
        % assume nargout <=5 when using finite difference calculation for
        % gradients in which case dxdalpha is unnecessary.

        % Gradient of time w.r.t. decVars
        % ------------------------------------------------------------
        % dt = (tF-t0)/(nTime-1)
        % t = t0 + n*dt
        % t = t0 + n*(tF-t0)/(nTime-1)
        % t = t0*(1-n/(nTime-1)) + tF*(n/(nTime-1))
        %
        % alpha = [t0, tF, x0, x1, ..., xN, u0, uM0, u1, ..., uN]
        % dt/dalpha = [1 - n/(nTime-1), n/(nTime-1), 0, 0, ... 0]
        % ------------------------------------------------------------

        n_time0 = n_time(idx);

        [k0, dk0] = combinedDynGrad(t0,        x0,                         u(:,idx), dynFun,pathObj);
        [k1, dk1] = combinedDynGrad(t0+0.5*dt, x0 + 0.5*dt*k0(1:nState,:), uMid(:,idx), dynFun,pathObj);
        [k2, dk2] = combinedDynGrad(t0+0.5*dt, x0 + 0.5*dt*k1(1:nState,:), uMid(:,idx), dynFun,pathObj);
        [k3, dk3] = combinedDynGrad(t0+dt,     x0 +     dt*k2(1:nState,:), u(:,idx+1), dynFun,pathObj);
        z = (dt/6)*(k0 + 2*k1 + 2*k2 + k3);  %Change over the sub-step

        for j = 1:nSegment

            % d(t[n])/dalpha
            dt_dalpha(1) = (1 - n_time0(j)/(nTime-1));
            dt_dalpha(2) = (n_time0(j)/(nTime-1));

            % du[n]/dalpha
            du_dalpha = zeros(nControl,nalpha);
            du_dalpha(:,gradInfo.indu(:,idx(j))) = eye(nControl);

            % duMid[n]/dalpha
            duMid_dalpha = zeros(nControl,nalpha);
            duMid_dalpha(:,gradInfo.indumid(:,idx(j))) = eye(nControl);

            % du[n+1]/dalpha
            du1_dalpha = zeros(nControl,nalpha);
            du1_dalpha(:,gradInfo.indu(:,idx(j)+1)) = eye(nControl);

            % dk0/dalpha
            dk0da = dk0(:,:,j) * [dt_dalpha; dxdalpha{j}(:,:,iSubStep); du_dalpha];

            % dk1/dalpha
            dk1da = dk1(:,:,j) * [dt_dalpha + 0.5/(nTime-1)*dTdalpha; dxdalpha{j}(:,:,iSubStep) + 0.5*dt*dk0da(1:nState,:) + 0.5/(nTime-1)*k0(1:nState,j)*dTdalpha; duMid_dalpha];

            % dk2/dalpha
            dk2da = dk2(:,:,j) * [dt_dalpha + 0.5/(nTime-1)*dTdalpha; dxdalpha{j}(:,:,iSubStep) + 0.5*dt*dk1da(1:nState,:) + 0.5/(nTime-1)*k1(1:nState,j)*dTdalpha; duMid_dalpha];

            % dk3/dalpha
            dk3da = dk3(:,:,j) * [dt_dalpha + 1/(nTime-1)*dTdalpha; dxdalpha{j}(:,:,iSubStep) + dt*dk2da(1:nState,:) + 1/(nTime-1)*k2(1:nState,j)*dTdalpha; du1_dalpha];

            dz = (dt/6)*(dk0da + 2*dk1da + 2*dk2da + dk3da)...
                + 1/(6*(nTime-1))*(k0(:,j)+2*k1(:,j)+2*k2(:,j)+k3(:,j))*dTdalpha;

            % update dxdalpha
            dxdalpha{j}(:,:,iSubStep+1) = dxdalpha{j}(:,:,iSubStep) + dz(1:nState,:);

            % update dJdalpha
            dJdalpha  = dJdalpha + dz(nState+1,:);
        end

        xNext = x0 + z(1:nState,:);  %Next state
        c(idx) = z(end,:);  %Integral of the cost function over this step

        if iSubStep == nSubStep %We've reached the end of the interval
            % Compute the defect vector:
            defects = xNext - x(:,idx+1);
        else
            % Store the state for next step in time
            idx = idx+1;   %  <-- This is important!!
            x(:,idx) = xNext;
        end

    end

    pathCost = sum(c);  %Sum up the integral cost over each segment

    %%%% Cache results to use on the next call to this function.
    RUNGE_KUTTA_t = t;
    RUNGE_KUTTA_x = x;
    RUNGE_KUTTA_u = u;
    RUNGE_KUTTA_defects = defects;
    RUNGE_KUTTA_pathCost = pathCost;
    RUNGE_KUTTA_dxdalpha = dxdalpha;
    RUNGE_KUTTA_dJdalpha = dJdalpha;

end

end

%%%%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~%%%%

function [dz, J] = combinedDynGrad(t,x,u,dynFun,pathObj)
% [dz, dJ] = combinedDynGrad(t,x,u,dynFun,pathObj)
%
% This function packages the dynamics and the cost function together so
% that they can be integrated at the same time.
%
% INPUTS:
%   t = [1, nTime] = time vector (grid points)
%   x = [nState, nTime] = state vector at each grid point
%   u = [nControl, nTime] = control vector at each grid point
%   dynamics(t,x,u) = dynamics function handle
%               dx = [nState, nTime] = dx/dt = derivative of state wrt time
%   pathObj(t,x,u) = integral cost function handle
%                 dObj = [1, nTime] = integrand from the cost function
%
% OUTPUTS:
%   dz = [dx; dObj] = combined dynamics of state and cost
%   dJ = [JAC(dynamics), JAC(objective)] = combined jacobian of dynamics
%   and objective w.r.t. (t,x,u)

nState = size(x,1);
nControl = size(u,1);

[dx,Jx] = dynFun(t,x,u);
if isempty(pathObj)
    dc = zeros(size(t));
    Jc = zeros(1,1+nState+nControl,length(t));
else
    [dc,Jc] = pathObj(t,x,u);
    Jc = reshape(Jc,1,1+nState+nControl,length(t));
end

dz = [dx;dc];

J = cat(1,Jx,Jc);

end


function [index_set, abundances, reconstruct, err] = AAM_no_subset( X, L, K )
% AAM_NO_SUBSET Alternating angle minimization implementation of the MESMA
% algorithm, called by AAM
%
% This function executes the AAM algorithm on a set of libraries. For every
% pixel in x, an endmember set is constructed containing a single endmember
% from each library. The indices and abundances are returned. The algorithm
% employed is iterative alternating angle minimization.
%
% Input:  X: the N data points to unmix in d dimensions, (d,N)
%         L: spectral library, cell array of p elements of size (d,~)
%         K: Optional number of iterations. Default 3
% Output: index_set identifying the endmembers from each library, (p,N)
%         abundances with respect to these endmembers, (p,N)
%         reconstruct contains the reconstructed spectral (d,N)
%         err contains the reconstruction error (Euclidean distance) (1,N)
%
%
% Rob Heylen, 2016, University of Antwerp.

% Turn this on to activate sanity checks on the libraries and input values.
% Turn this off if you are sure there are no doubles in the libraries, and
% you do not want to check if pixels are contained in the libraries.
library_check=1;

% Initializations
if nargin==2
    K=3;
end
[d,numpx]=size(X);
p=numel(L);
for i=1:p
    N(i)=size(L{i},2);
end
flag=0;
index_set=zeros(p,numpx);
abundances=zeros(p,numpx);
F=zeros(d,p);
I=ones(p,1);
reconstruct=zeros(d,numpx);
err=zeros(1,numpx);

% Main loop over all pixels
for px=1:numpx
    x=X(:,px);

    % Check if x is a library member. If so, we can finish immediately
    if library_check
        for i=1:p
            %if sum(sum(abs(L{i}-x*ones(1,N(i))))==0)>0
            if numel(find(~sum(abs(L{i}-x*ones(1,N(i))))))>0
                I=ones(p,1);
                I(i)=find(sum(abs(L{i}-x*ones(1,N(i))))==0);
                index_set(:,px)=I;
                flag=1;
                break;
            end
        end
    end
    if flag
        flag=0;
        continue;
    end

    % Create random initial endmember set
    for i=1:p
        I(i)=ceil(rand*N(i));
        F(:,i)=L{i}(:,I(i));
    end

    % Iterate K times
    for it=1:K
        % Alternating angle optimization
        for i=1:p
            % Calculate angles
            Fi=F(:,[1:i-1 i+1:p]);       % Pivot
            Gi=[Fi x];                   % Plane through pivot and pixel
            E1=plane_project2(L{i},Fi);  % Project library onto pivot
            E2=plane_project2(L{i},Gi);  % Project library onto plane
            p1=sqrt(sum((E1-L{i}).^2));  % Distances from library to pivot
            p2=sqrt(sum((E2-L{i}).^2));  % Distances from library to plane
            ang=asin(p2./p1);            % Resulting angles

            % Find angles that should be inverted
            mask=(x-plane_project2(x,Fi))'*(E2-E1)<0;
            ang(mask)=pi-ang(mask);

            % Identify minimal angle, update index set
            [~,I(i)]=min(ang);

            % Update endmember set
            F(:,i)=L{i}(:,I(i));
        end
    end

    % Update index_set with obtained indices
    index_set(:,px)=I;
end

% Unmixing phase
E=zeros(d,p);
go=0;
for px=1:numpx
    for i=1:p
        E(:,i)=L{i}(:,index_set(i,px));
    end
    % Plug in your favorite unmixing program here

    if go==0
        [at,opt]=FCLSU_fast2(X(:,px)',E);
        abundances(:,px)=at';
        go=1;
    else
        at=FCLSU_fast2(X(:,px)',E,opt);
        abundances(:,px)=at';
    end

    % Reconstruction
    reconstruct(:,px)=E*abundances(:,px);

    % Error
    err(px)=norm(reconstruct(:,px)-X(:,px));
end

end