% zica() - Z-transform of ICA activations; useful for studying component SNR
%
% Usage: >> [zact,basesd,maz,mazc,mazf] = zica(activations,frames,baseframes)
%
% Inputs:
%   activations - activations matrix produced by runica()
%   frames      - frames per epoch {0|default ->  length(activations)}
%   baseframes  - vector of frames in z-defining baseline period {default frames}
%
% Outputs:
%   zact        - activations z-scaled and reorded in reverse order of max abs
%   basesd      - standard deviations in each activation row (reverse ordered)
%   maz         - maximum absolute z-value for each activation row (rev ordered)
%   mazc        - component indices of the reverse-sorted max abs z-values
%                 (this is the act -> zact reordering)
%   mazf        - frame indices of the max abs z-values (reverse ordered)
%
% Author: Scott Makeig, SCCN/INC/UCSD, La Jolla, 2-25-98
%
% See also: runica()

function [zact,basesd,maxabsz,maxc,maxabszf] = zica(activations,frames,baseframes)

% Copyright (C) 2-25-98 Scott Makeig, SCCN/INC/UCSD, scott@sccn.ucsd.edu
%
% This program is free software; you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation; either version 2 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

% 3-2-98 added frames variable -sm
% 1-25-01 put revsort subfunction in the core of the zica program -ad
% 01-25-02 reformated help & license, added link -ad

if nargin<1
   help zica
   return
end

[chans,framestot] = size(activations);

if nargin < 3
   baseframes = 0;
end
if nargin < 2
   frames = 0;
end
if frames == 0
   frames = framestot;
end
if baseframes == 0
   baseframes = 1:frames
end
epochs = floor(framestot/frames);
if frames*epochs ~= framestot
   fprintf('zica(): indicated frames does not divide data length.\n');
   return
end

if length(baseframes) < 3
  fprintf('\n  zica() - too few baseframes (%d).\n',length(baseframes));
  help zica
  return
end

if min(baseframes) < 1 | max(baseframes) > frames
  fprintf('\n  zica() - baseframes out of range.\n');
  help zica
  return
end

baselength = length(baseframes);
baseact = zeros(epochs*baselength,chans);
for e=1:epochs
 baseact((e-1)*baselength+1:e*baselength,:) = ...
                        matsel(activations,frames,baseframes,0,e)';
end
basesd = sqrt(covary(baseact));
zact = activations./(basesd'*ones(1,framestot));
[maxabsz,maxabszf] = sort(abs(zact'));
maxabsz  = maxabsz(frames,:);
maxabszf = maxabszf(frames,:);

%
% reorder outputs in reverse order of max abs z
[maxabsz,maxc] = revsort(maxabsz);
zact = zact(maxc,:);
basesd = basesd(maxc);
maxabszf = maxabszf(maxc);

% revsort  - reverse sort columns (biggest 1st, ...)
function [out,i] = revsort(in)

if size(in,1) == 1
   in = in'; % make column vector
end

[out,i] = sort(in);
out = out(size(in,1):-1:1,:);
  i = i(size(in,1):-1:1,:);
return;


function classifier = lapsvmp(options,data)
% {lapsvmp} trains a Laplacian SVM classifier in the primal.
%
%      classifier = lapsvmp(options,data)
%
%      options: a structure with the following fields
%               options.gamma_A: regularization parameter (ambient norm)
%               options.gamma_I: regularization parameter (intrinsic norm)
%
%               [optional fields]
%               options.Cg: {0,1} i.e. train with Newton's method or PCG
%                           (default=0)
%               options.MaxIter: maximum number of iterations (default=200)
%               options.Hinge: {0,1} i.e. train a LapSVM (1) or LapRLSC (0)
%                              (default=1)
%               options.UseBias: {0,1} i.e. use or not a bias (default=0)
%               options.InitAlpha: if it's 0, the initial weights are null;
%                                  if it's <0, they are randomly taken;
%                                  otherwise it is the initial vector
%                                  (default=0)
%               options.InitBias: the initial bias (default=0)
%               options.NewtonLineSearch: {0,1} i.e. use or not exact line
%                                         search with Newton's method
%                                         (default=0).
%               options.NewtonCholesky: {0,1} i.e. use or not Cholesky
%                                       factorization and rank 1 updates of
%                                       the Heassian (default=0 for LapSVM.
%                                       If set to 1, the  Hessian must be
%                                       positive definite).
%               options.CgStopType: {0,1,2,3,4,5,6,7,-1}
%                                   the data-based stopping criterion of cg
%                                   only (default=0), where:
%                                   0: do not stop
%                                   1: stability stop
%                                   2: validation stop
%                                   3: stability & validation stop
%                                   4: gradient norm (normalized)
%                                   5: prec. gradient norm (normalized)
%                                   6: mixed gradient norm (normalized)
%                                   7: relative objective function change
%                                  -1: debug (it saves many stats)
%               options.CgStopIter: number of cg iters after which check
%                                   the stopping condition.
%               options.CgStopParam: the parameter for the selected
%                                    CgStopType. If CgStopType is:
%                                    0: ignored
%                                    1: percentage [0,1] of tolerated
%                                       different decisions between two
%                                       consecutive checks
%                                    2: percentage [0,1] of error rate
%                                       decrement required  between two
%                                       consecutive checks
%                                    3: the two params above (i.e. it is an
%                                       array of two elements)
%                                    4: the minimum gradient norm
%                                    5: the minimum prec. gradient norm
%                                    6: the minimum mixed gradient norm
%                                    7: relative objective function change
%                                       between two consecutive checks
%                                    -1: ignored
%                                    see the code for default values.
%               options.Verbose: {0,1} (default=1)
%
%      data: a structure with the following fields
%            data.X: a N-by-D matrix of N D-dimensional training examples
%            data.K: a N-by-N kernel Gram matrix of N training examples
%            data.Y: a N-by-1 label vector in {-1,0,+1}, where 0=unlabeled
%                    (it is in {0,+1} in the case of One-Class SVM/LapSVM.
%            data.L: a N-by-N matrix of the Laplacian
%
%            [other fields]
%            data.Kv: a V-by-N kernel Gram matrix of V validation examples,
%                     required if CgStopType is 2 or 3 (validation check)
%            data.Yv: a V-by-1 vector of {-1,+1} labels for the validation
%                     examples, required if CgStopType is 2 or 3
%                     (validation check)
%
%      classifier: structure of the trained classifier (see the
%                  'saveclassfier' function).
%
% Author: Stefano Melacci (2012)
%         mela@dii.unisi.it
%         * the One-Class extension is joint work with Salvatore Frandina,
%           salvatore.frandina@gmail.com
%         * the original code structure was based on the primal SVM code of
%           Olivier Chapelle, olivier.chapelle@tuebingen.mpg.de

n=length(data.Y);
nn=nnz(data.Y==-1);

% initializing option structure with default values
if ~isfield(options,'Verbose'),           options.Verbose=1; end
if ~isfield(options,'Hinge'),             options.Hinge=1; end
if ~isfield(options,'Cg'),                options.Cg=0; end
if ~isfield(options,'MaxIter'),           options.MaxIter=200; end
if ~isfield(options,'UseBias'),           options.UseBias=0; end
if ~isfield(options,'InitAlpha'),         options.InitAlpha=false; end
if ~isfield(options,'InitBias'),          options.InitBias=0; end
if ~isfield(options,'NewtonLineSearch'),  options.NewtonLineSearch=0; end
if ~isfield(options,'NewtonCholesky'),    options.NewtonCholesky=0; end
if ~isfield(options,'CgStopType'),        options.CgStopType=0; end
switch options.CgStopType
    case 0 % none
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=0; end
        if ~isfield(options,'CgStopIter'),
            options.CgStopIter=options.MaxIter+1; end
    case 1 % stability
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=0.015; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=round(sqrt(n)/2); end
    case 2 % validation
        v=length(data.Yv);
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=1/v; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=round(sqrt(n)/2); end
    case 3 % stability & validation
        v=length(data.Yv);
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=...
                                                           [1/v,0.015];
        end
        if ~isfield(options,'CgStopIter'),
            options.CgStopIter=round(sqrt(n)/2); end
    case 4 % gradient norm
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=1e-8; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=1; end
    case 5 % preconditioned gradient norm
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=1e-8; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=1; end
    case 6 % mixed gradient norm
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=1e-8; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=1; end
    case 7 % relative objective function decrease
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=1e-6; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=1; end
    case -1 % debug
        if ~isfield(options,'CgStopParam') || ...
            isempty(options.CgStopParam), options.CgStopParam=0; end
        if ~isfield(options,'CgStopIter') || isempty(options.CgStopIter),
            options.CgStopIter=options.MaxIter+1; end
    otherwise
        error('Invalid CgStopType.');
end
if nn==0, oc=1; else oc=0; end

% initial alpha vector
if length(options.InitAlpha)>1
    alpha=options.InitAlpha;
else
    if options.InitAlpha==0
        alpha=[];
    else
        alpha=randn(n,1);
        alpha=alpha./norm(alpha);
    end
end

% checking common error conditions
if oc==1 && ~options.UseBias
    error('One-Class SVM requires the UseBias option to be turned on.');
end
if options.UseBias && options.LaplacianNormalize
    error(['The current implementation does not support a normalized ' ...
           'Laplacian when the UseBias options is turned on.']);
end

% initial bias
b=options.InitBias;

switch options.Cg
    case 0
        % newton
        [alpha,b,t,sec,lsiters]=newton(options,data,alpha,b,oc);
        stats=[];
    case 1
        % pcg
        [alpha,b,t,sec,lsiters,stats]=pcg(options,data,alpha,b,oc);
    otherwise
        error('Invalid solver specified in the field .Cg');
end

svs=find(alpha~=0);
classifier=saveclassifier('lapsvmp',svs,alpha(svs), ...
                          data.X(svs,:),b,options,sec,t,lsiters,stats);

function [alpha,b,t,sec,lsiters] = newton(options,data,alpha,b,oc)
% {newton} trains the classifier using the Newton's method.

tic
n=length(data.Y);
labeled=data.Y~=0;
l=nnz(labeled);
gamma_A=options.gamma_A;
gamma_I=options.gamma_I;

% initial seeding
if isempty(alpha)
    alpha=zeros(n,1);
    Kalpha=zeros(n,1);
else
    Kalpha=data.K*alpha;
end

t=0;
lr=0;
sv=false(n,1);
if nargout>4, lsiters=zeros(options.MaxIter,1); end

if gamma_I~=0, LK=data.L*data.K; end

while 1

    if options.Hinge
        sv_prev=sv;
        hloss=sparse([],[],[],n,1,l);
        hloss(labeled)=1-data.Y(labeled).*(Kalpha(labeled,:)+b);
        sv=hloss>0;
        nsv=nnz(sv);
    else
        sv_prev=sv;
        sv=labeled;
        nsv=l;
    end

    if options.Verbose
        if ~options.Hinge
            hloss=sparse([],[],[],n,1,l);
            hloss(labeled)=1-data.Y(labeled).*(Kalpha(labeled,:)+b);
        end
        if gamma_I~=0
            obj=(gamma_A*alpha'*Kalpha+sum(hloss(sv).^2)+...
                gamma_I*Kalpha'*data.L*Kalpha+oc*b)/2;
        else
            obj=(gamma_A*alpha'*Kalpha+sum(hloss(sv).^2)+oc*b)/2;
        end

        fprintf('[t=%d] obj=%f nev=%d lr=%.4f\n', [t full(obj) nsv lr]);
    end

    % goal conditions
    if t>=options.MaxIter, break, end
    if isequal(sv_prev,sv), break, end

    t=t+1;

    IsvK=sparse([],[],[],n,n,nsv*n);
    IsvK(sv,:)=data.K(sv,:);

    % computing new alphas
    onev=ones(1,n);

    if gamma_I==0 % SVM (sparse solution)
        if options.UseBias
            alpha_new=zeros(n,1);
            alpha_b_new=[0,onev(1:nsv);onev(1:nsv)',...
                         gamma_A*speye(nsv)+IsvK(sv,sv)]\ ...
                         [oc/(2*gamma_A);data.Y(sv)];
            alpha_new(sv)=alpha_b_new(2:end);
            b_new=alpha_b_new(1);
        else
            alpha_new=zeros(n,1);
            alpha_new(sv)=(gamma_A*speye(nsv)+IsvK(sv,sv))\data.Y(sv);
            b_new=0;
        end

    else % LapSVM

        % inversion by factorization
        if options.NewtonCholesky
            if t==1
                % compute the Cholesky factorization of the Hessian
                if options.UseBias
                    sumKsv=sum(data.K(sv,:));
                    hess=chol([nsv,sumKsv;sumKsv',...
                               data.K*(gamma_A*speye(n)+IsvK+gamma_I*LK)]);
                    alpha_b_new=hess\(hess'\([sum(data.Y(sv))-oc/2; ...
                                              data.K(:,sv)*data.Y(sv)]));
                    alpha_new=alpha_b_new(2:end);
                    b_new=alpha_b_new(1);
                else
                    hess=chol(data.K*(gamma_A*speye(n)+IsvK+gamma_I*LK));
                    alpha_new=hess\(hess'\(data.K(:,sv)*data.Y(sv)));
                    b_new=0;
                end
                LK=[];
            else
                % update the Cholesky factorization of the Hessian
                sv_diff=~(sv&sv_prev);
                sv_add=find(sv&sv_diff)';
                sv_rem=find(sv_prev&sv_diff)';
                if options.UseBias
                    if ~isempty(sv_add)
                        for i=sv_add
                            hess=cholupdate(hess,[1;data.K(:,i)],'+');
                        end
                    end
                    if ~isempty(sv_rem)
                        for i=sv_rem
                            hess=cholupdate(hess,[1;data.K(:,i)],'-');
                        end
                    end
                    alpha_b_new=hess\(hess'\([sum(data.Y(sv))-oc/2; ...
                                              data.K(:,sv)*data.Y(sv)]));
                    alpha_new=alpha_b_new(2:end);
                    b_new=alpha_b_new(1);
                else
                    if ~isempty(sv_add)
                        for i=sv_add
                            hess=cholupdate(hess,data.K(:,i),'+');
                        end
                    end
                    if ~isempty(sv_rem)
                        for i=sv_rem
                            hess=cholupdate(hess,data.K(:,i),'-');
                        end
                    end
                    alpha_new=hess\(hess'\(data.K(:,sv)*data.Y(sv)));
                    b_new=0;
                end
            end
        else
            % inversion without factorization
            IsvY=sparse([],[],[],n,1,nsv);
            IsvY(sv)=data.Y(sv);
            if options.UseBias
                alpha_b_new=([0,onev;sv,...
                             gamma_A*speye(n)+IsvK+gamma_I*LK])\ ...
                             [oc/(2*gamma_A);IsvY];
                alpha_new=alpha_b_new(2:end);
                b_new=alpha_b_new(1);
            else
                alpha_new=(gamma_A*speye(n)+IsvK+gamma_I*LK)\IsvY;
                b_new=0;
            end

        end
    end

    % step
    if options.NewtonLineSearch && (options.Hinge || nnz(alpha)>0)
        step=alpha_new-alpha;
        step_b=b_new-b;

        [lr,Kalpha,lsi]=linesearch(data,labeled,step,step_b,Kalpha,b,...
                                   gamma_A,gamma_I,[],[],options.Hinge,...
                                   oc);
        alpha=alpha+lr*step;
        b=b+lr*step_b;
        lsiters(t)=lsi;
    else
        alpha=alpha_new;
        b=b_new;
        lr=1;
        lsiters(t)=0;
        Kalpha=data.K*alpha;
    end
end

if nargout>4, lsiters=lsiters(1:t); end
sec=toc;

function [alpha,b,t,sec,lsiters,stats] = pcg(options,data,alpha,b,oc)
% {pcg} trains the classifier using preconditioned conjugate gradient.

tic
n=length(data.Y);
labeled=data.Y~=0;
unlabeled=data.Y==0;
l=nnz(labeled);
u=n-l;

if isfield(data,'Yv')
    v=length(data.Yv);
end
gamma_A=options.gamma_A;
gamma_I=options.gamma_I;

if isempty(alpha)
    alpha=zeros(n,1);
    Kalpha=zeros(n,1);
    if gamma_I~=0, LKalpha=zeros(n,1); else LKalpha=[]; end
    go=data.Y-b*labeled;
    obj0=(sum((1-data.Y(labeled)*b*options.UseBias).^2)+oc*b)/2;
    if options.UseBias, go_b=sum(data.Y(labeled)-b)-oc/2;
    else go_b=0; b=0; end
else
    Kalpha=data.K*alpha;
    if gamma_I~=0, LKalpha=data.L*Kalpha; else LKalpha=[]; end
    out=sparse([],[],[],n,1,l);
    out(labeled)=Kalpha(labeled)+b;
    if options.Hinge
        sv=false(n,1);
        sv(labeled)=(data.Y(labeled).*out(labeled)<1);
    else
        sv=labeled;
    end
    go=-gamma_A*alpha;
    if gamma_I~=0, go=go-gamma_I*LKalpha; end
    obj0=(Kalpha'*(-g0)+sum((out(sv)-data.Y(sv)).^2)+oc*b)/2;
    go(sv)=go(sv)-(out(sv)-data.Y(sv));

    if options.UseBias, go_b=sum(data.Y(sv)-Kalpha(sv)-b)-oc/2;
    else go_b=0; b=0; end

end
d=go; % initial search direction
d_b=go_b;
Kgo=data.K*go;
Kstep=Kgo;

t=0;
stats=[];

switch options.CgStopType
    case 4, ng0=sqrt(sum(Kgo.^2)+go_b^2);
    case 5, ngp0=sqrt(sum(go.^2)+go_b^2);
    case 6, ngm0=sqrt(sum(Kgo'*go)+go_b^2);
    case -1
        stats=zeros(options.MaxIter+1,5+n+1);
        ng0=sqrt(sum(Kgo.^2)+go_b^2);
        ngp0=sqrt(sum(go.^2)+go_b^2);
        ngm0=sqrt(sum(Kgo'*go)+go_b^2);
        stats(t+1,:)=[t,obj0/obj0,ng0/ng0,ngp0/ngp0,ngm0/ngm0,alpha',b];
end

if nargout>3, lsiters=zeros(options.MaxIter,1); end
valerr_prev=1;
obj_prev=obj0;
yfx_unlabeled_prev=false(u,1);

while 1
    t=t+1;

    % goal condition: maximum number of iterations
    if t>options.MaxIter, t=t-1; break, end

    % do an exact line search
    [lr,Kalpha,lsi,LKalpha]=linesearch(data,labeled,d,d_b,Kalpha,b,...
                                       gamma_A,gamma_I,Kstep,LKalpha,...
                                       options.Hinge,oc);

    % goal condition: converged to optimal solution
    if lr==0, t=t-1; break, end

    alpha=alpha+lr*d;
    b=b+lr*d_b;
    lsiters(t)=lsi;

    % compute new precgradient and objective
    out=sparse([],[],[],n,1,l);
    out(labeled)=Kalpha(labeled)+b;

    if options.Hinge
        sv=false(n,1);
        sv(labeled)=(data.Y(labeled).*out(labeled)<1);
    else
        sv=labeled;
    end

    g=gamma_A*alpha;
    if gamma_I~=0, g=g+gamma_I*LKalpha; end

    if options.Verbose || options.CgStopType==7 || options.CgStopType==-1
        obj=(Kalpha'*(g)+sum((out(sv)-data.Y(sv)).^2)+oc*b)/2;
        if options.Verbose
            fprintf('[t=%d] obj=%f nev=%d lr=%.4f\n', ...
                     [t full(obj) nnz(sv) lr]);
        end
    end

    g(sv)=g(sv)+(out(sv)-data.Y(sv));

    if options.UseBias, g_b=sum(out(sv)-data.Y(sv))+oc/2; else g_b=0; end

    gn=-g;
    gn_b=-g_b;

    % goal condition: data based
    if mod(t-1,options.CgStopIter)==(options.CgStopIter-1)
        switch options.CgStopType
            case 1 % stability
                yfx_unlabeled=sign(Kalpha(unlabeled)+b);
                diff_unlabeled=1-nnz(yfx_unlabeled==yfx_unlabeled_prev)/u;
                if diff_unlabeled<options.CgStopParam, break, end
                yfx_unlabeled_prev=yfx_unlabeled;
            case 2 % validation
                valerr=1-nnz(sign(data.Kv*alpha+b)==data.Yv)/v;
                if (valerr>(valerr_prev-options.CgStopParam)), break, end
                valerr_prev=valerr;
            case 3 % stability & validation
                valerr=1-nnz(sign(data.Kv*alpha+b)==data.Yv)/v;
                yfx_unlabeled=sign(Kalpha(unlabeled)+b);
                diff_unlabeled=1-nnz(yfx_unlabeled==yfx_unlabeled_prev)/u;
                if (valerr>(valerr_prev-options.CgStopParam(1))) && ...
                    diff_unlabeled<options.CgStopParam(2), break, end
                valerr_prev=valerr;
                yfx_unlabeled_prev=yfx_unlabeled;
            case 4 % gradient norm
                ng=sqrt(sum(gn.^2)+gn_b^2);
                if ng/ng0<options.CgStopParam, break, end
            case 5 % preconditioned gradient norm
                ngp=sqrt(sum(go.^2)+go_b^2);
                if ngp/ngp0<options.CgStopParam, break, end
            case 6 % mixed gradient norm
                ngm=sqrt(Kgo'*go+go_b^2);
                if ngm/ngm0<options.CgStopParam, break, end
            case 7 % relative objective function decrease
                if (obj_prev-obj)<options.CgStopParam, break, end
                obj_prev=obj;
        end
    end

    Kgn=data.K*gn; % multiply by the preconditioner

    % debug
    if options.CgStopType==-1
        ng=sqrt(sum(Kgn.^2)+gn_b^2);
        ngp=sqrt(sum(gn.^2)+gn_b^2);
        ngm=sqrt(Kgn'*gn+gn_b^2);
        stats(t+1,:)=[t,obj/obj0,ng/ng0,ngp/ngp0,ngm/ngm0,alpha',b];
    end

    % Polack-Ribiere update with automatic restart
    be=max(0,(Kgn'*(gn-go)+gn_b*(gn_b-go_b))/(Kgo'*go+go_b^2));

    d=be*d+gn;
    d_b=be*d_b+gn_b;

    Kstep=be*Kstep+Kgn;

    go=gn;
    go_b=gn_b;
    Kgo=Kgn;
end
sec=toc;
if nargout>4, lsiters=lsiters(1:t);
    if nargout>5 && ~isempty(stats), stats=stats(1:t,:); end
end

function [lr,Kalpha,lsi,LKalpha] = linesearch(data,labeled,step,step_b, ...
                                              Kalpha,b,gamma_A,gamma_I, ...
                                              Kstep,LKalpha,hinge,oc)
% {linesearch} does a line search in direction step.

act=step~=0; % the set of points for which alpha change (active)
if isempty(Kstep)
    Kstep=data.K(:,act)*step(act);
end

% precomputations
stepKstep=step(act)'*Kstep(act);
stepKalpha=step(act)'*Kalpha(act);
if gamma_I~=0
    KstepL=Kstep'*data.L;
    KstepLKalpha=KstepL*Kalpha;
    KstepLKstep=KstepL*Kstep;
end

out=Kalpha(labeled)+b;
outstep=Kstep(labeled)+step_b;
out_minus_Y=out-data.Y(labeled);

% breakpoints
if hinge
    sv=(1-out.*data.Y(labeled))>0;
    deltas=-out_minus_Y./outstep;
    deltas(deltas<0)=0;
    [deltas,deltas_map]=sort(deltas);
    lab=length(deltas);
    i=find(deltas>0,1);
    lsi=1;
else
    sv=true(length(out),1);
    lsi=1;
end

% intercepts
if gamma_I~=0
    left=outstep(sv)'*out_minus_Y(sv)+gamma_A*stepKalpha+...
         gamma_I*KstepLKalpha + (oc/2)*step_b;
    right=outstep(sv)'*outstep(sv)+gamma_A*stepKstep+...
          gamma_I*KstepLKstep;
else
    left=outstep(sv)'*out_minus_Y(sv)+gamma_A*stepKalpha + (oc/2)*step_b;
    right=outstep(sv)'*outstep(sv)+gamma_A*stepKstep;
end

% first minimum
zcross=-left/right;
if right<=0, zcross=0; end

if hinge && ~isempty(i)
    if sv(deltas_map(i))==true
        if 0<=zcross && zcross<deltas(i), not_got_it=0;
        else not_got_it=1; end
    else
        if 0<=zcross && zcross<=deltas(i), not_got_it=0;
        else not_got_it=1; end
    end

    while not_got_it
        % updating support vectors
        j=-2*sv(deltas_map(i))+1;
        sv(deltas_map(i))=~sv(deltas_map(i));

        % updating intercepts
        left=left+j*outstep(deltas_map(i))*out_minus_Y(deltas_map(i));
        right=right+j*outstep(deltas_map(i))*outstep(deltas_map(i));

        % computing minimum
        zcross=-left/right;

        % goal conditions
        if i==lab, break, end
        if sv(deltas_map(i))==true
            if sv(deltas_map(i+1))==true
                if deltas(i)<zcross && zcross<deltas(i+1),
                    not_got_it=0; end
            else
                if deltas(i)<zcross && zcross<=deltas(i+1),
                    not_got_it=0; end
            end
        else
            if sv(deltas_map(i+1))==true
                if deltas(i)<=zcross && zcross<deltas(i+1),
                    not_got_it=0; end
            else
                if deltas(i)<=zcross && zcross<=deltas(i+1),
                    not_got_it=0; end
            end
        end

        i=i+1;
        lsi=lsi+1;
    end
end
lr=zcross;

if lr<0, lr=0; return, end % converged
Kalpha=Kalpha+lr*Kstep;
if nargout>3 && gamma_I~=0, LKalpha=LKalpha+lr*KstepL'; end


%% Image Pyramids
%
% In this tutorial:
%
% * We will learn about Image Pyramids
% * We will use Image pyramids to create a new fruit, "Orapple"
% * We will see these functions: |cv.pyrUp|, |cv.pyrDown|, |cv.buildPyramid|
%
% Sources:
%
% * <https://docs.opencv.org/3.1.0/dc/dff/tutorial_py_pyramids.html>
%

%% Theory
% Normally, we used to work with an image of constant size. But on some
% occasions, we need to work with (the same) images in different resolutions.
% For example, while searching for something in an image, like face, we are
% not sure at what size the object will be present in said image. In that
% case, we will need to create a set of the same image with different
% resolutions and search for object in all of them. These set of images with
% different resolutions are called *Image Pyramids* (because when they are
% kept in a stack with the highest resolution image at the bottom and the
% lowest resolution image at top, it looks like a pyramid).
%
% There are two kinds of Image Pyramids:
%
% # Gaussian Pyramid
% # Laplacian Pyramids
%
% Higher level (Low resolution) in a Gaussian Pyramid is formed by removing
% consecutive rows and columns in Lower level (higher resolution) image.
% Then each pixel in higher level is formed by the contribution from 5 pixels
% in underlying level with gaussian weights. By doing so, a |MxN| image
% becomes |M/2xN/2| image. So area reduces to one-fourth of original area. It
% is called an Octave. The same pattern continues as we go upper in pyramid
% (ie, resolution decreases). Similarly while expanding, area becomes 4 times
% in each level. We can find Gaussian pyramids using |cv.pyrDown| and
% |cv.pyrUp| functions.
%
% Below is the 4 levels in an image pyramid:
%

% compute Gaussian pyramid
img = cv.imread(fullfile(mexopencv.root(),'test','fruits.jpg'), 'ReduceScale',2);
p = cv.buildPyramid(img, 'MaxLevel',4);

% combine all levels in one image
[rows,cols,~] = cellfun(@size, p);
dst = zeros([rows(1) cols(1)+cols(2) size(img,3)], class(img));
r_start = [0 0 cumsum(rows(2:end-1))] + 1;
r_end = [rows(1) cumsum(rows(2:end))];
c_start = [0 repmat(cols(1),1,numel(p)-1)] + 1;
c_end = [cols(1) cols(1)+cols(2:end)];
for i=1:numel(p)
    dst(r_start(i):r_end(i), c_start(i):c_end(i), :) = p{i};
end
figure, imshow(dst)

%%
% Now you can go down the image pyramid with |cv.pyrUp| function.
%
% Remember, |higher_reso2| is not equal to |higher_reso|, because once you
% decrease the resolution, you loose the information.
%

higher_reso = img;
lower_reso = cv.pyrDown(higher_reso);
higher_reso2 = cv.pyrUp(lower_reso);
figure, imshow(cat(2, higher_reso, higher_reso2))

%%
% Laplacian Pyramids are formed from the Gaussian Pyramids. There is no
% exclusive function for that. Laplacian pyramid images are like edge images
% only. Most of its elements are zeros. They are used in image compression.
% A level in Laplacian Pyramid is formed by the difference between that level
% in Gaussian Pyramid and expanded version of its upper level in Gaussian
% Pyramid. The three levels of a Laplacian level will look like below
% (contrast is adjusted to enhance the contents):
%

% compute Gaussian pyramid
img = imread(fullfile(mexopencv.root(),'test','butterfly.jpg'));
p = cv.Blender.createLaplacePyr(img, 4);

% combine all levels in one image
[rows,cols,~] = cellfun(@size, p);
dst = zeros([rows(1)+rows(2) cols(1) size(img,3)], class(img));
r_start = [0 repmat(rows(1),1,numel(p)-1)] + 1;
r_end = [rows(1) rows(1)+rows(2:end)];
c_start = [0 0 cumsum(cols(2:end-1))] + 1;
c_end = [cols(1) cumsum(cols(2:end))];
for i=1:numel(p)-1
    dst(r_start(i):r_end(i),  c_start(i):c_end(i), :) = uint8(p{i});
end
figure, imshow(dst)
%imshow(cv.CLAHE(rgb2gray(dst)))

%% Image Blending using Pyramids
% One application of Pyramids is Image Blending. For example, in image
% stitching, you will need to stack two images together, but it may not look
% good due to discontinuities between images. In that case, image blending
% with Pyramids gives you seamless blending without leaving much data in the
% images. One classical example of this is the blending of two fruits, Orange
% and Apple.
%
% Please check first reference in additional resources, it has full
% diagramatic details on image blending, Laplacian Pyramids etc. Simply it is
% done as follows:
%
% * Load the two images of apple and orange
% * Find the Gaussian Pyramids for apple and orange (in this particular
%   example, number of levels is 6)
% * From Gaussian Pyramids, find their Laplacian Pyramids
% * Now join the left half of apple and right half of orange in each levels
%   of Laplacian Pyramids
% * Finally from this joint image pyramids, reconstruct the original image.
%
% Below is the full code. (For sake of simplicity, each step is done
% separately which may take more memory. You can optimize it if you want so).
%

% a pair of images of same size
A = imread(fullfile(mexopencv.root(),'test','apple.jpg'));
B = imread(fullfile(mexopencv.root(),'test','orange.jpg'));

% generate Laplacian Pyramids
pA = cv.Blender.createLaplacePyr(A, 5);
pB = cv.Blender.createLaplacePyr(B, 5);

% add left and right halves of images in each level
p = pB;
for i=1:numel(p)
    idx = round(size(pA{i},2)/2);
    p{i}(:,1:idx,:) = pA{i}(:,1:idx,:);
end
%{
% add top and bottom halves of images in each level
p = pB;
for i=1:numel(p)
    idx = round(size(pA{i},1)/2);
    p{i}(1:idx,:,:) = pA{i}(1:idx,:,:);
end
%}

% reconstruct
C = cv.Blender.restoreImageFromLaplacePyr(p);
C = uint8(C);

% image with direct connecting each half
idx = round(size(A,2)/2);
D = B;
D(:,1:idx,:) = A(:,1:idx,:);

% show images
figure
subplot(221), imshow(A), title('Apple')
subplot(222), imshow(B), title('Orange')
subplot(223), imshow(C), title('Pyramid Blending')
subplot(224), imshow(D), title('Direct Connection')

%% Additional Resources
%
% <http://pages.cs.wisc.edu/~csverma/CS766_09/ImageMosaic/imagemosaic.html Image Blending>
%


function NodeValue = hlp_computeGraphMeasure(varargin)
% NodeValue = hlp_computeGraphMeasure(causality,ch1,selectedvars,graphMetric)
%
% Compute a univariate graph-theoretic measure for a given node of a graph
%
% Inputs:
%
%   causality:      [num_vars_to x num_vars_from x <num_times> x <num_freqs>] causal matrix
%                   obtained from est_mvarConnectivity().
%   srcNodes:       indices of source nodes
%   selectedvars:   indices of target nodes
%   graphMetric:   which graph measure to compute (see below)
%
% Outputs:
%
%   NodeValue:      [srcNodes x <num_times> x <num_freqs>] graph measures for the source nodes
%
%
% See Also: vis_causalBrainMovie3D(), est_mvarConnectivity(),
%           hlp_collapseFrequencies()
%
%
% References:
%
%   Mullen T (2010) The Source Information Flow Toolbox (SIFT):
%   Theoretical Handbook and User Manual. Chapter 6.
%   Available at: http://www.sccn.ucsd.edu/wiki/SIFT
%
% Author: Tim Mullen, 2010-2013, SCCN/INC/UCSD.
% Email:  tim@sccn.ucsd.edu

% This function is part of the Source Information Flow Toolbox (SIFT)
%
% This program is free software; you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation; either version 3 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

arg_define([0 Inf],varargin, ...
    arg_norep({'cmatrix','CausalMatrix'}, mandatory,[],'Causality array. Shape is [num_vars_to x num_vars_from x <num_times> x <num_freqs>]. Last two dims are optional','shape','matrix'), ...
    arg({'srcNodes','SourceNodes'},  [],[],'Indices of source nodes. Default is all nodes'), ...
    arg({'targNodes','TargetNodes'}, [],[],'Indices of target nodes. Default is all nodes'), ...
    arg({'graphMetric','GraphMetric'},'none',{'none','outflow','mag_outflow','inflow', ...
                                              'causalflow','outdegree','indegree',     ...
                                              'causaldegree','asymmetryratio'}, ...
                                              'Graph measure to compute'), ...
    arg({'ignoreSelfConn','IgnoreSelfConn'},true,[],'Ignore self connectivity in graph measure'));

% handle defaults
if isempty(srcNodes)
    srcNodes = 1:size(cmatrix,2);
end
if isempty(targNodes)
    targNodes = 1:size(cmatrix,1);
end

% remove source node from target nodes
if ignoreSelfConn
    cmatrix=hlp_setdiags(cmatrix,0);
%     targNodes = setdiff_bc(targNodes,srcNodes);
end

% compute graph metrics
switch lower(graphMetric)
    case 'none'
        NodeValue = zeros(1,size(cmatrix,3));
    case 'outflow'
        % Compute outflow from srcNodes in each freq
        NodeValue = squeeze(sum(cmatrix(targNodes,srcNodes,:,:),1));
    case 'mag_outflow'
        % Compute outflow from srcNodes in each freq, ignoring sign
        NodeValue = squeeze(sum(abs(cmatrix(targNodes,srcNodes,:,:)),1));
    case 'inflow'
        % Compute inflow to srcNodes in each freq
        NodeValue = squeeze(sum(cmatrix(srcNodes,targNodes,:,:),2));
    case 'causalflow'
        outflow =   squeeze(sum(cmatrix(targNodes,srcNodes,:,:),1));
        inflow  =   squeeze(sum(cmatrix(srcNodes,targNodes,:,:),2));
        NodeValue = outflow - inflow;
    case 'outdegree'
        % Compute number of outgoing edges from srcNodes in each freq
        NodeValue = squeeze(sum(logical(cmatrix(targNodes,srcNodes,:,:)),1));
    case 'indegree'
        % number of incoming edges to srcNodes in each freq
        NodeValue = squeeze(sum(logical(cmatrix(srcNodes,targNodes,:,:)),2));
    case 'causaldegree'
        % outdegree - indegree
        outflow =   squeeze(sum(logical(cmatrix(targNodes,srcNodes,:,:)),1));
        inflow  =   squeeze(sum(logical(cmatrix(srcNodes,targNodes,:,:)),2));
        NodeValue = outflow - inflow;
    case 'asymmetryratio'
        % 1 if all edges are outgoing, -1 if all edges are incoming.
        % 0 if balanced
        outflow =   squeeze(sum(cmatrix(targNodes,srcNodes,:,:),1));
        inflow  =   squeeze(sum(cmatrix(srcNodes,targNodes,:,:),2));
        NodeValue = (outflow - inflow)./(outflow+inflow);
    otherwise
        % user wants to map a different Conn measure to this
        % (e.g., ERSP)
end

% enforce row vector output
if iscolumn(NodeValue)
    NodeValue = NodeValue';
end


function linpack_s_test24 ( )

%*****************************************************************************80
%
%% TEST24 tests SSICO.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    24 June 2009
%
%  Author:
%
%    John Burkardt
%
  n = 100;
  lda = n;

  fprintf ( 1, '\n' );
  fprintf ( 1, 'TEST24\n' );
  fprintf ( 1, '  For a symmetric indefinite matrix,\n' );
  fprintf ( 1, '  SSICO estimates the reciprocal condition number.\n' );
  fprintf ( 1, '  The matrix size is N = %d\n', n );
%
%  Assign values to the matrix A.
%
  a(1:n,1:n) = 0.0;

  for i = 1 : n
    a(i,i) = 2.0;
    if ( i < n )
      a(i,i+1) = -1.0;
    end
  end
%
%  Estimate the condition.
%
  fprintf ( 1, '\n' );
  fprintf ( 1, '  Estimate the condition.\n' );

  [ a, ipvt, rcond, z ] = ssico ( a, lda, n );

  fprintf ( 1, '\n' );
  fprintf ( 1, '  Estimated reciprocal condition = %f\n', rcond );

  return
end


clc
clear all
tic
%***********************User inserted input for the problem***************

nosp = 5;             %reads the number of sales points considered
noft = 4;                %reads the number of factories
nopd = 3;         %number of planning period
trkcap = 70;                    %Capacity of the truck

%***********************inputs from excel file for the formulation ***************

demand = [35 35	60	80	100; 25	100	125	55	20; 150	120	100	60	110; 50	70	40	140	40 ; 90	110	90	50	30;0	110	90	40	50; 60	70	0	110	10; 140	70	90	100	30 ;10	130	130	40	140; 30	90	70	40	0; 70	100	100	50	90; 110	50	50	110	100 ]; %stores the demand in the time horizons for each sales point
lostdemandcost = [30	35	25	40	55] ;  %lost demand per unit for each sales point
transcost = [ 100	400	800	1100	1700 ; 1500	1800	1200	2200	1500 ; 1400	1000	400	1000	700 ; 800	700	900	700	1200] ; %xlsread('greedyinput','Sheet1','O3:S6');        %transportation cost matrix
setupcost = [ 4500	2000	2500	3000 ] ;  %setup costs for individual factories
factorycap = [300	500	250	150]; % xlsread('greedyinput','Sheet2','B4:E4');       %capacity of the factories
factholdcost = [5	2	4	3]; %xlsread('greedyinput','Sheet2','B5:E5');     %holding cost per unit in a factory
for i = 1:nopd
quantityshipped{i} = zeros(noft,nosp);
factoryresiduals{i} = zeros(1,noft);
factoryactivation{i} = zeros(1,noft);
salespointresiduals{i} = zeros(1,nosp);
end

%*****************************counter-check if the inputs are well read*********

disp(demand);
disp(lostdemandcost);
disp(transcost);
disp(setupcost);
disp(factorycap);
disp(factholdcost);
for i = 1:nopd
disp(quantityshipped{i});
end

%****************************starts the calculation of the greedy alg********

initlostsales = calclostsales(demand,lostdemandcost,nopd,nosp); %Calculates the lost demand with the initial assumption that all the demand is lost
disp(initlostsales);
cumsetupcost = 0;
cumtransportcost = 0;
cumholdingcost = 0;
saveddemand = 0;
salespoint{nosp}= [0,0];
for ft = 1:noft                                          %Initialize the factory activation status and residual for each time preiod (zero for the first one)
           factory{ft}(1,1) = 0;
           factory{ft}(1,2) = 0;
end
for time = 1:nopd                                               %iterates for each planning period beginning from the 1st

       for sp = 1:nosp
           salespoint{sp}(1,1) = 0;                             %the array stores the activation status and the demand of a sales point
            salespoint{sp}(1,2) =  demand(time,sp);             %this populates the sales point for a time horizon with demand and filling status
       end

       for ft = 1:noft                                          %Initialize the factory activation status and residual for each time preiod (zero for the first one)
           factory{ft}(1,1) = 0;

       end

   % loadfactory(capacity,prevresidual);

    for points = 1:nosp
        [selectedsp,maxlostsales] = choosesalespt(lostdemandcost,salespoint,nosp);             %Chooses the salespoint with the highest possible lost sales
        disp('----------------------------');
        fprintf('Selected Sales Point: %d\n', selectedsp);
        fprintf('Potential Lost Sales cost: %.2f\n', maxlostsales);

        salespoint{selectedsp}(1,1)=1;

            [costunit,selectedf] = choosefactory(factorycap,factholdcost,transcost,setupcost,factory,trkcap,salespoint,selectedsp); %chooses the preferable factory to deliver to the selected salespoint
            fprintf('selected factory: %d\n',selectedf);
            fprintf('Cost of the selected factory: %.2f\n', costunit);
            quantityshipped{time}(selectedf,selectedsp) = salespoint{selectedsp}(1,2);
            if (factory{selectedf}(1,2)>= salespoint{selectedsp}(1,2))
            factory{selectedf}(1,2) = factory{selectedf}(1,2) - salespoint{selectedsp}(1,2);

            % saveddemand = saveddemand + salespoint{selectedsp}(1,2)*lostdemandcost(selectedsp);                             % accumulates the reduction of the lost demand for each sent item to a destination sales point
            end
            if(factory{selectedf}(1,1)==0 && factory{selectedf}(1,2)<= salespoint{selectedsp}(1,2))
                factory{selectedf}(1,2) = factory{selectedf}(1,2)+factorycap(selectedf)-salespoint{selectedsp}(1,2);
                cumsetupcost = cumsetupcost + setupcost(selectedf);                                                         %accumulates the activation cost for each factory when there is an accompanying activation
               % cumtransportcost = cumtransportcost + ceil(salespoint{selectedsp}(1,2)/trkcap)*transcost(selectedf,selectedsp);
                factory{selectedf}(1,1)=1;
            end
            saveddemand = saveddemand + salespoint{selectedsp}(1,2)*lostdemandcost(selectedsp);
            cumtransportcost = cumtransportcost + ceil(salespoint{selectedsp}(1,2)/trkcap)*transcost(selectedf,selectedsp); %accumulates the transportaton cost for each iteration

            fprintf('Cumulative setupcost: %.2f\n',cumsetupcost);
            fprintf('Cumulative trasportation cost: %.2f\n',cumtransportcost);
    end
    currentholdingcost = 0;
    for i = 1:noft
        currentholdingcost= currentholdingcost + factory{i}(1,2)*factholdcost(i);                                         %the holding costs are accumulatedat each factory at the end of each planning period
        factoryresiduals{time}(i) = factory{i}(1,2);
        factoryactivation{time}(i) = factory{i}(1,1);
        factory{i}(1,1)= 0;
    end
    cumholdingcost = cumholdingcost + currentholdingcost;                                                                  % holding cost of the whole factory at the end of a planning period
    objectivesolution = initlostsales -saveddemand+ cumholdingcost+cumsetupcost+cumtransportcost
    end
%objectivesolution = initlostsales-saveddemand+cumsetupcost+cumtransportcost
%objectivesolution = initlostsales -saveddemand+ cumholdingcost+cumsetupcost+cumtransportcost
disp('====================shipped quantity matrix==================')
for i = 1:nopd
disp(quantityshipped{i});
end
disp('=====================factory residual matrix=================')
for i = 1:nopd
disp(factoryresiduals{i});
end
disp('=====================factory activation status===============')
for i = 1:nopd
disp(factoryactivation{i});
end
objectivesolution
% New_transportation_costs=0;
% F_holding_cost=0;
% Act_cost= 0;
% for t=1:nopd
%     for i=1:noft
%         F_holding_cost = F_holding_cost + factholdcost(i) *factoryresiduals{t}(1,i);
%         Act_cost = Act_cost + factoryactivation{t}(1,i)* setupcost(i);
%
%
%         for j =1:nosp
%     New_transportation_costs= ceil(quantityshipped{t}(i,j)/trkcap) * transcost(i,j)+ New_transportation_costs;
%         end
%     end
% end
% New_transportation_costs
% F_holding_cost
% Act_cost
% tot_cost = New_transportation_costs + F_holding_cost +Act_cost
% saveddemand
% initlostsales
% cumtransportcost
toc


function S = globMatrixPGIFE3DMass(fun,mesh,femI,fem1,fem2)

%% USAGE: generate stiffness global matrix on a 3D mesh
%
% INPUTS:
% fun --- coefficient function
% mesh --- a struct data contains very rich mesh information.
% fem1 --- global DoF for test function space
% fem2 --- global DoF for trial function space
%
% OUTPUTS:
% [IN JN XN] --- triplets of the sparse matrix from regular elements.
% [II JI XI] --- triplets of the sparse matrix from interface elements.

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

%% 0. Initializaiton
if strcmp(fem1.type,'P1')||strcmp(fem1.type,'DGP1')||strcmp(fem1.type,'CR')
    feEvalBas1 = @evalP1Bas3D;
elseif strcmp(fem1.type,'P2')||strcmp(fem1.type,'DGP2')
    feEvalBas1 = @evalP2Bas3D;
end

if strcmp(fem2.type,'P1')||strcmp(fem2.type,'DGP1')||strcmp(fem2.type,'CR')
    feEvalBas2 = @evalP1Bas3D;
elseif strcmp(fem2.type,'P2')||strcmp(fem2.type,'DGP2')
    feEvalBas2 = @evalP2Bas3D;
end

%% 1. Matrix on noninterface elements
dof1 = fem1.ldof; dof2 = fem2.ldof; nloc = dof1*dof2;
ntID = find(mesh.tLoc > 0); ntN = length(ntID);
AN = fem1.area(ntID);
gxN = fem1.gx(ntID,:); gyN = fem1.gy(ntID,:); gzN = fem1.gz(ntID,:); gw = fem1.gw;
XN = zeros(nloc*ntN, 1);

coefN = feval(fun,gxN,gyN,gzN);
Ibas = cell(dof1,1);
Jbas = cell(dof2,1);
for i = 1:dof1
    Ibas{i} = feEvalBas1(fem1.bas(ntID,:,i), gxN, gyN, gzN, [0,0,0]);
end
for j = 1:dof2
    Jbas{j} = feEvalBas2(fem2.bas(ntID,:,j), gxN, gyN, gzN, [0,0,0]);
end

IN = reshape(repmat(fem1.t(ntID,:),4,1),nloc*ntN,1);
JN = repmat(reshape(fem2.t(ntID,:),dof2*ntN,1),4,1);
ind = 0;
for i = 1:dof1
    for j = 1:dof2
        XN(ind+1:ind+ntN) = AN.*(sum(((Ibas{i}.*(coefN.*Jbas{j})).*gw'),2));
        ind = ind + ntN;
    end
end
ID = find(XN~=0);
SN = sparse(IN(ID),JN(ID),XN(ID),size(fem1.p,1),size(fem2.p,1));

%% 2. Matrix on interface elements
AI = femI.area; gw = femI.gw; gxI = femI.gx; gyI = femI.gy; gzI = femI.gz;
ntI = size(femI.t,1); % not number of interface element, but quadrature element
XI = zeros(nloc*ntI, 1);
tID = femI.QelemID;

coefI = feval(fun,gxI,gyI,gzI);
Ibas = cell(dof1,1);
Jbas = cell(dof2,1);
for i = 1:dof1
    Ibas{i} = feEvalBas1(fem1.bas(tID,:,i), gxI, gyI, gzI, [0,0,0]);
end
for j = 1:dof2
    Jbas{j} = feEvalBas2(femI.bas(:,:,j), gxI, gyI, gzI, [0,0,0]);
end

II = reshape(repmat(femI.t,4,1),nloc*ntI,1);
JI = repmat(reshape(femI.t,dof2*ntI,1),4,1);
ind = 0;
for i = 1:dof1
    for j = 1:dof2
        XI(ind+1:ind+ntI) = AI.*(sum(((Ibas{i}.*(coefI.*Jbas{j})).*gw'),2));
        ind = ind + ntI;
    end
end
SI = sparse(II,JI,XI,size(fem1.p,1),size(fem2.p,1));
S = SN + SI;


function [mssim, ssim_map] = ssim(img1, img2, K, window, L)

% ========================================================================
% SSIM Index with automatic downsampling, Version 1.0
% Copyright(c) 2009 Zhou Wang
% All Rights Reserved.
%
% ----------------------------------------------------------------------
% Permission to use, copy, or modify this software and its documentation
% for educational and research purposes only and without fee is hereby
% granted, provided that this copyright notice and the original authors'
% names appear on all copies and supporting documentation. This program
% shall not be used, rewritten, or adapted as the basis of a commercial
% software or hardware product without first obtaining permission of the
% authors. The authors make no representations about the suitability of
% this software for any purpose. It is provided "as is" without express
% or implied warranty.
%----------------------------------------------------------------------
%
% This is an implementation of the algorithm for calculating the
% Structural SIMilarity (SSIM) index between two images
%
% Please refer to the following paper and the website with suggested usage
%
% Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli, "Image
% quality assessment: From error visibility to structural similarity,"
% IEEE Transactios on Image Processing, vol. 13, no. 4, pp. 600-612,
% Apr. 2004.
%
% http://www.ece.uwaterloo.ca/~z70wang/research/ssim/
%
% Note: This program is different from ssim_index.m, where no automatic
% downsampling is performed. (downsampling was done in the above paper
% and was described as suggested usage in the above website.)
%
% Kindly report any suggestions or corrections to zhouwang@ieee.org
%
%----------------------------------------------------------------------
%
%Input : (1) img1: the first image being compared
%        (2) img2: the second image being compared
%        (3) K: constants in the SSIM index formula (see the above
%            reference). defualt value: K = [0.01 0.03]
%        (4) window: local window for statistics (see the above
%            reference). default widnow is Gaussian given by
%            window = fspecial('gaussian', 11, 1.5);
%        (5) L: dynamic range of the images. default: L = 255
%
%Output: (1) mssim: the mean SSIM index value between 2 images.
%            If one of the images being compared is regarded as
%            perfect quality, then mssim can be considered as the
%            quality measure of the other image.
%            If img1 = img2, then mssim = 1.
%        (2) ssim_map: the SSIM index map of the test image. The map
%            has a smaller size than the input images. The actual size
%            depends on the window size and the downsampling factor.
%
%Basic Usage:
%   Given 2 test images img1 and img2, whose dynamic range is 0-255
%
%   [mssim, ssim_map] = ssim(img1, img2);
%
%Advanced Usage:
%   User defined parameters. For example
%
%   K = [0.05 0.05];
%   window = ones(8);
%   L = 100;
%   [mssim, ssim_map] = ssim(img1, img2, K, window, L);
%
%Visualize the results:
%
%   mssim                        %Gives the mssim value
%   imshow(max(0, ssim_map).^4)  %Shows the SSIM index map
%========================================================================

if (nargin < 2 || nargin > 5)
   mssim = -Inf;
   ssim_map = -Inf;
   return;
end

if (size(img1) ~= size(img2))
   mssim = -Inf;
   ssim_map = -Inf;
   return;
end

[M N] = size(img1);

if (nargin == 2)
   if ((M < 11) || (N < 11))
	   mssim = -Inf;
	   ssim_map = -Inf;
      return
   end
   %window = fspecial('gaussian', 11, 1.5);	%
   window = Gaussian_window(11,1.5);
   K(1) = 0.01;					% default settings
   K(2) = 0.03;					%
   L = 255;                                     %
end

if (nargin == 3)
   if ((M < 11) || (N < 11))
	   mssim = -Inf;
	   ssim_map = -Inf;
      return
   end
   %window = fspecial('gaussian', 11, 1.5);
   window = Gaussian_window(11,1.5);
   L = 255;
   if (length(K) == 2)
      if (K(1) < 0 || K(2) < 0)
		   mssim = -Inf;
   		ssim_map = -Inf;
	   	return;
      end
   else
	   mssim = -Inf;
   	ssim_map = -Inf;
	   return;
   end
end

if (nargin == 4)
   [H W] = size(window);
   if ((H*W) < 4 || (H > M) || (W > N))
	   mssim = -Inf;
	   ssim_map = -Inf;
      return
   end
   L = 255;
   if (length(K) == 2)
      if (K(1) < 0 || K(2) < 0)
		   mssim = -Inf;
   		ssim_map = -Inf;
	   	return;
      end
   else
	   mssim = -Inf;
   	ssim_map = -Inf;
	   return;
   end
end

if (nargin == 5)
   [H W] = size(window);
   if ((H*W) < 4 || (H > M) || (W > N))
	   mssim = -Inf;
	   ssim_map = -Inf;
      return
   end
   if (length(K) == 2)
      if (K(1) < 0 || K(2) < 0)
		   mssim = -Inf;
   		ssim_map = -Inf;
	   	return;
      end
   else
	   mssim = -Inf;
   	ssim_map = -Inf;
	   return;
   end
end

img1 = double(img1);
img2 = double(img2);

% automatic downsampling
f = max(1,round(min(M,N)/256));
%downsampling by f
%use a simple low-pass filter
if(f>1)
    lpf = ones(f,f);
    lpf = lpf/sum(lpf(:));
    img1 = imfilter(img1,lpf,'symmetric','same');
    img2 = imfilter(img2,lpf,'symmetric','same');

    img1 = img1(1:f:end,1:f:end);
    img2 = img2(1:f:end,1:f:end);
end

C1 = (K(1)*L)^2;
C2 = (K(2)*L)^2;
window = window/sum(sum(window));

mu1   = filter2(window, img1, 'valid');
mu2   = filter2(window, img2, 'valid');
mu1_sq = mu1.*mu1;
mu2_sq = mu2.*mu2;
mu1_mu2 = mu1.*mu2;
sigma1_sq = filter2(window, img1.*img1, 'valid') - mu1_sq;
sigma2_sq = filter2(window, img2.*img2, 'valid') - mu2_sq;
sigma12 = filter2(window, img1.*img2, 'valid') - mu1_mu2;

if (C1 > 0 && C2 > 0)
   ssim_map = ((2*mu1_mu2 + C1).*(2*sigma12 + C2))./((mu1_sq + mu2_sq + C1).*(sigma1_sq + sigma2_sq + C2));
else
   numerator1 = 2*mu1_mu2 + C1;
   numerator2 = 2*sigma12 + C2;
	denominator1 = mu1_sq + mu2_sq + C1;
   denominator2 = sigma1_sq + sigma2_sq + C2;
   ssim_map = ones(size(mu1));
   index = (denominator1.*denominator2 > 0);
   ssim_map(index) = (numerator1(index).*numerator2(index))./(denominator1(index).*denominator2(index));
   index = (denominator1 ~= 0) & (denominator2 == 0);
   ssim_map(index) = numerator1(index)./denominator1(index);
end

mssim = mean(ssim_map(:));

return


%% Solves the Krusell and Smith (1998)
% Solving Krusell-Smith model using perturbation method involves simple 5
% step process given by
%
% # Solve for Steady State
% # Linearize Model Equations
% # Solve out Static Constraint (optional reduce model)
% # Solve Linear System
% # Compute Impulse Response Functions
% # Check for Internal Consistency (if model reduction is used)
% where each step just requires about one function call.
%
% In this example file, the entire problem will be solved total of 3 times.
% First, KS model will be solved without any model reduction. Second, KS
% model will be solved with model reduction, to compare how well the model
% reduction works. Lastly, different parts of observables will be
% introduced to compare how the model reduction works with different
% requirements. We have updates to Internal Consistency check coming in the
% future.
%
% Estimated Runtime: 1.5 seconds
%
% REFERENCES:
%
% * Ahn, SeHyoun, Greg Kaplan, Benjamin Moll, Thomas Winberry, and
% Christian Wolf. "When Inequality Matters for Macro and Macro Matters for
% Inequality."
% * Krusell, Per, and Anthony A. Smith, Jr. "Income and wealth heterogeneity in the macroeconomy." Journal of political Economy 106.5 (1998): 867-896.
%
% REQUIRES:
%
% * auto diff toolbox: <https://github.com/sehyoun/MATLABAutoDiff>
% * phact toolbox: <https://github.com/gregkaplan/phact>
% * compute_steady_state.m
% * equilibrium_conditions.m
% * plot_steady_state.m
% * plot_IRFs.m
% Relevant files except for auto diff toolbox, can be found in
%    /examples/KrusellSmith folder of phact toolbox provided on github
%
%%

%% Setup the toolbox
% Just need to include folders containing the files in the path.

addpath('/home/sehyoun/Dropbox/0Packages/AutoDiff');
addpath('/home/sehyoun/Dropbox/0Packages/PHACT');

%% Set options for this example run
% example_case : 1   run full and compare with only value function reduction
%                2   run full and compare with only state space reduction
%                3   run full and compare with both reductions
example_case = 3;
reduceDist_hor = 50;        % m of K_m(A,b)

% initialize shocks for simulation
T = 200;
N = 2000;
vAggregateShock = zeros(1,N);
vAggregateShock(1,1) = 1;
%vAggregateShock = randn(1,N);     % Uncomment for one realization instead
                                  % of IRFs

%% Step 0: Set Parameters
% The script sets up parameters relevant for the model
set_parameters;

%% Step 1: Solve for Steady State
% Non-stochastic steady state can be found using any methods. In
%    particular, example codes can be found at
%    <http://www.princeton.edu/~moll/HACTproject.htm>.

tStart = tic;
fprintf('Computing steady state...\n')
global IfSS IbSS I0SS varsSS A

[rSS,wSS,KSS,ASS,uSS,cSS,VSS,gSS,dVUSS,dVfSS,dVbSS,IfSS,IbSS,I0SS] = ...
    compute_steady_state();

fprintf('Time to compute steady state: %.3g seconds\n\n\n',toc(tStart));

% Store steady state values
varsSS = zeros(nVars,1);
varsSS(1:2*I,1) = reshape(VSS,2*I,1);
ggSS = reshape(gSS,2*I,1);
varsSS(2*I+1:4*I-1,1) = ggSS(1:2*I-1);
varsSS(4*I,1) = 0;
varsSS(4*I+1,1) = KSS;
varsSS(4*I+2,1) = rSS;
varsSS(4*I+3,1) = wSS;
varsSS(4*I+4,1) = (KSS ^ aalpha) * (zAvg ^ (1 - aalpha));
CSS = sum(cSS(:) .* gSS(:) * da);
varsSS(4*I+5,1) = CSS;
varsSS(4*I+6,1) = ddelta * KSS;

% plot steady state results
plot_steady_state;

%% Step 2: Linearize Model Equations
% For computing derivatives, the codes written for solving for the
%    steady-state can be used almost verbatim using automatic
%    differentiation toolbox as long as only the functions supported by
%    automatic differentation are used. For list of supported functions and
%    documentation of relevant syntax check
%    <https://github.com/sehyoun/MATLABAutoDiff>. Example usage/syntax of
%    automatic differentiation can be found at
%    <http://sehyoun.com/EXAMPLE_AutoDiff_Syntax.html>
fprintf('Taking derivatives of equilibrium conditions...\n')
t0 = tic;

% Prepare automatic differentiation
vars = zeros(2*nVars+nEErrors+1,1);
vars = myAD(vars);

% Evaluate derivatives
derivativesIntermediate = equilibrium_conditions(vars);

% Extract out derivative values
derivs = getderivs(derivativesIntermediate);

tDerivs = toc(t0);
fprintf('...Done!\n')
fprintf('Time to compute derivatives: %2.4f seconds\n\n\n',tDerivs)
if tDerivs > 1
    warning('If you do not compile mex files for automatics differentiation, matrix vector multiplication will be slow');
    disp('Press any key to continue...');
    pause();
end

% Unpackage derivatives
mVarsDerivs = derivs(:,1:nVars);
mVarsDotDerivs = derivs(:,nVars+1:2*nVars);
mEErrorsDerivs = derivs(:,2*nVars+1:2*nVars+nEErrors);
mShocksDerivs = derivs(:,2*nVars+nEErrors+1);

%% Step 3 (full): Solve out Static Constraints and/or Reduce Models
% For the first run, a unreduced model will be solved

% rename derivative matrix to match notation in paper
g0 = mVarsDotDerivs;
g1 = -mVarsDerivs;
c = sparse(nVars,1);
psi = -mShocksDerivs;
pi = -mEErrorsDerivs;

% Solve out static constratins
[~,inv_state_red,g0,g1,constant,pi,psi] = clean_G0_sparse(g0,g1,c,pi,psi);

%% Step 4: Solve Linear System
t0 = tic;
fprintf('Solving linear system...\n')

[G1,~,impact,eu] = schur_solver(g0,g1,c,psi,pi,1,1,1);

fprintf('...Done!\n')
fprintf('Existence and uniqueness? %2.0f and %2.0f\n',eu);
fprintf('Time to solve full linear system: %2.4f seconds\n\n\n',toc(t0))

%% Step 5: Simulate Impulse Response Functions
t0 = tic;

[simulated,~] = simulate(G1,impact,T,N,vAggregateShock,'implicit',inv_state_red);%,4*I:4*I+6);

fprintf('...Done!\n')
fprintf('Time to simulate model: %2.4f seconds\n\n\n',toc(t0))

big_simul = simulated + varsSS;
simulated = simulated(4*I:4*I+6,:);

varsSS_small = varsSS(4*I:4*I+6,1);
% Add state-states back in to get values in levels
vAggregateTFP			= simulated(1,:) + varsSS_small(1);
vAggregateOutput		= simulated(5,:) + varsSS_small(5);
vAggregateConsumption	= simulated(6,:) + varsSS_small(6);
vAggregateInvestment 	= simulated(7,:) + varsSS_small(7);

% Compute log differences for plotting
vAggregateTFP_full		= vAggregateTFP;
vAggregateOutput_full	= log(vAggregateOutput) - log(varsSS_small(5));
vAggregateConsumption_full	= log(vAggregateConsumption) - log(varsSS_small(6));
vAggregateInvestment_full	= log(vAggregateInvestment) - log(varsSS_small(7));

%% Step 3 (reduced): Solve Out Static Constraints and/or Reduce Models

% Set relevant parameters for example cases
if (example_case == 1)
    reduceV = 1;
    reduceDistribution = 0;
elseif (example_case == 2)
    reduceV = 0;
    reduceDistribution = 1;
else
    reduceV = 1;
    reduceDistribution = 1;
end

% rename derivatives to match notation in paper
g0 = mVarsDotDerivs;
g1 = -mVarsDerivs;
c = sparse(nVars,1);
psi = -mShocksDerivs;
pi = -mEErrorsDerivs;

t0 = tic;
fprintf('Model Reduction ...\n')

%%
% *State space reduction using Krylov subspace method*
if reduceDistribution == 1
    % State space reduction
    [state_red,inv_state_red,n_g_red] = krylov_reduction(g0,g1,n_v,n_g,reduceDist_hor);
    [g1,psi,pi,c,g0] = change_basis(state_red,inv_state_red,g1,psi,pi,c,g0);
else
    % Clean G0
    [state_red,inv_state_red,g0,g1,c,pi,psi] = clean_G0_sparse(g0,g1,c,pi,psi);
    n_g_red = n_g;
end

%%
% *Value function reduction using spline inspired bases*
if reduceV == 1
    % Create knot points for spline (the knot points are not uniformly spaced)
    knots = linspace(amin,amax,n_knots-1)';
    knots = (amax-amin)/(2^c_power-1)*((knots-amin)/(amax-amin)+1).^c_power+amin-(amax-amin)/(2^c_power-1);
    % Function calls to create basis reduction
    [from_spline, to_spline] = oneDquad_spline(x,knots);
    [from_spline, to_spline] = extend_to_nd(from_spline,to_spline,n_prior,n_post);
    n_splined = size(from_spline,2);
    [from_spline, to_spline] = projection_for_subset(from_spline,to_spline,0,n_g_red);

    % Reduce the decision vector
    [g1,psi,~,c,g0] = change_basis(to_spline,from_spline,g1,psi,pi,c,g0);
    pi = to_spline * pi * from_spline(1:n_v,1:n_splined);
elseif reduceV == 0
    from_spline = speye(n_g_red + n_v);
    to_spline = speye(n_g_red + n_v);
    n_splined = n_v;
end
fprintf('...Done!\n')
fprintf('Time to reduce dimensionality: %2.4f seconds\n\n\n',toc(t0))

%% Step 4: Solve Linear System
t0 = tic;
fprintf('Solving reduced linear system...\n')

[G1,~,impact,eu,F] = schur_solver(g0,g1,c,psi,pi,1,1,1);

fprintf('...Done!\n')
fprintf('Existence and uniqueness? %2.0f and %2.0f\n',eu);
fprintf('Time to solve linear system: %2.4f seconds\n\n\n',toc(t0))

%% Step 5: Simulate Impulse Response Functions
fprintf('Simulating Model...\n')
t0 = tic;

trans_mat = inv_state_red*from_spline;
[simulated,vTime] = simulate(G1,impact,T,N,vAggregateShock,'implicit',trans_mat,4*I:4*I+6);

fprintf('...Done!\n')
fprintf('Time to simulate model: %2.4f seconds\n\n\n',toc(t0))

% Add state-states back in to get values in levels
vAggregateTFP = simulated(1,:) + varsSS_small(1);
vAggregateOuput = simulated(5,:) + varsSS_small(5);
vAggregateConsumption = simulated(6,:) + varsSS_small(6);
vAggregateInvestment = simulated(7,:) + varsSS_small(7);

% Compute log differences for plotting
vAggregateTFP_reduced = vAggregateTFP;
vAggregateOutput_reduced = log(vAggregateOutput) - log(varsSS_small(5));
vAggregateConsumption_reduced = log(vAggregateConsumption) - log(varsSS_small(6));
vAggregateInvestment_reduced = log(vAggregateInvestment) - log(varsSS_small(7));

%% Step 6: Internal Consistency Check
% For large problem, we are forced to take explicit update, so it can take
% awhile to run, but the good thing is that this only needs to be run at
% the end. We are currently working on this, so speed for this part might
% improve in the future. (For small problems, the speed is not an issue). A
% different consistency check is in works, and will be updated in the
% future.

g1 = -mVarsDerivs;
psi = -mShocksDerivs;
from_red = inv_state_red * from_spline;
to_red = to_spline * state_red;
[epsilon] = internal_consistency_check(G1,impact,n_g_red,from_red,to_red,g1,psi,F,n_v,n_g,1000,varsSS,1,0);

%% (optional) Step 7: Plot relevant values
% Plot impulse response functions
plot_IRFs;

%% State Space Reduction with More Observables
% It is possible to include other variables into observables. In this part,
%    different parts of the distribution will be included in the model
%    reduction to show how the Krylov subspace based model reduction
%    behaves. To that regard, the reduced model is solved while including
%    different parts of distribution as part of the observable. In this
%    case, a very strong requirement of knowing the probability density
%    value at different grid points of the distribution g is taken. In the
%    first case, 80th through 89th grid points were included in the parts
%    that we consider to be observable. For the second example, 85th
%    through 94th grid points were included in observable.

% rename derivative matrix to match notation in paper
g0 = mVarsDotDerivs;
g1 = -mVarsDerivs;
c = sparse(nVars,1);
psi = -mShocksDerivs;
pi = -mEErrorsDerivs;

% Solve full model
[~,inv_state_red,g0_full,g1_full,c_full,pi_full,psi_full] = clean_G0_sparse(g0,g1,c,pi,psi);
[G1,~,impact,eu] = schur_solver(g0_full,g1_full,c_full,psi_full,pi_full,1,1,1);
[simulated,vTime] = simulate(G1,impact,T,N,vAggregateShock,'implicit',inv_state_red);
simulated_full = simulated + varsSS;

% Solve reduced without extra observables
[state_red,inv_state_red,n_g_red] = krylov_reduction(g0,g1,n_v,n_g,20);
[g1_red,psi_red,pi_red,c_red,g0_red] = change_basis(state_red,inv_state_red,g1,psi,pi,c,g0);
[G1,~,impact,eu] = schur_solver(g0_red,g1_red,c_red,psi_red,pi_red,1,1,1);
[simulated,vTime] = simulate(G1,impact,T,N,vAggregateShock,'implicit',inv_state_red);
simulated_red = simulated + varsSS;

%%
% Solve reduced with extra observables for 80th through 89th grid point

[state_red,inv_state_red,n_g_red] = krylov_reduction(g0,g1,n_v,n_g,20,spdiags(ones(n_g,1),79,10,n_g));
[g1_80,psi_80,pi_80,c_80,g0_80] = change_basis(state_red,inv_state_red,g1,psi,pi,c,g0);
[G1,~,impact,eu] = schur_solver(g0_80,g1_80,c_80,psi_80,pi_80,1,1,1);
[simulated,vTime] = simulate(G1,impact,T,N,vAggregateShock,'implicit',inv_state_red);
simulated_obs = simulated + varsSS;
diff_red = (abs(simulated_red(n_v+1:end,:) - simulated_full(n_v+1:end,:)))./simulated_full(n_v+1:end,:);
diff_obs80 = (abs(simulated_obs(n_v+1:end,:) - simulated_full(n_v+1:end,:)))./simulated_full(n_v+1:end,:);

figure;
plot(70:99,log(diff_red(70:99,11)));
hold on;
plot(70:99,log(diff_obs80(70:99,11)));
grid on
legend('Location','northwest');
legend('without','with grid 80~89 included');
title(['Error Comparison at time ',num2str(vTime(11))],'interpreter','latex','fontsize',14)
ylabel('$\log$(relative error)','interpreter','latex')
xlabel('Distribution Location','interpreter','latex')

%%
% Solve reduced with extra observables for 85th through 94th grid point

[state_red,inv_state_red,n_g_red] = krylov_reduction(g0,g1,n_v,n_g,20,spdiags(ones(n_g,1),84,10,n_g));
[g1,psi,pi,c,g0] = change_basis(state_red,inv_state_red,g1,psi,pi,c,g0);
[G1,~,impact,eu] = schur_solver(g0,g1,c,psi,pi,1,1,1);
[simulated,vTime] = simulate(G1,impact,T,N,vAggregateShock,'implicit',inv_state_red);
simulated_obs = simulated + varsSS;
diff_obs = (abs(simulated_obs(n_v+1:end,:) - simulated_full(n_v+1:end,:)))./simulated_full(n_v+1:end,:);

figure;
plot(70:99,log(diff_red(70:99,11)));
hold on;
plot(70:99,log(diff_obs(70:99,11)));
grid on
legend('Location','northwest');
legend('without','with grid 85~94 included');
title(['Error Comparison at time ',num2str(vTime(11))],'interpreter','latex','fontsize',14)
ylabel('$\log$(relative error)','interpreter','latex')
xlabel('Distribution Location','interpreter','latex')

%%
% Note how the errors decreases only for the parts that we required to be
% observable while the irrelevant parts were reduced away. This is the
% magic of Krylov subspace based reduction in action. Though we could solve
% the full model for the simple Krusell-Smith case, this ability to reduce
% away irrelevant part will be essential for bigger problems.
%

%%
% Lastly, some of the functions have lengthy parameters, you can always
% type "help function_name;" to see documentation. All functions contain the
% function call syntax as the last line, and thoese can be copied and
% pasted into the program file. For example, to use the
% krylov_reduction function, you can call
help krylov_reduction;

%%
% and just copy the lastline of the documentation string. Alternately, you
% can call "doc function_name" to read the doc string in the MATLAB help
% browser, and "edit function_name" to see the function codes.


clear;
disp('Program started');
% control mode (0: impedance control
%               1: admittance control(velocity command)
%               2: admittance control(position command))
control_mode = 0;

% reference trajectory mode (0: horizontal line
%                            1: circle)
trajectory_mode = 0;

% dyanmic params(the same as those used in vrep)
I1 = 1;
I2 = 1;
l1 = 0.3;
l2 = 0.3;
m1 = 0.5;
m2 = 0.5;
a1 = 0.15;
a2 = 0.15;
g = 9.81;
% impedance model params
kpx = 150;
kpy = 150;
kdx = 70;
kdy = 70;
mx = 1;
my = 1;

% for moving avarage force filter
windowSize = 5;
b = (1/windowSize)*ones(1,windowSize);
a = 1;
raw_externalforces = zeros(2,windowSize);
filtered_externalforce = zeros(2,1);

vrep=remApi('remoteApi'); % using the prototype file (remoteApiProto.m)
vrep.simxFinish(-1);      % just in case, close all opened connections
id=vrep.simxStart('127.0.0.1',19997,true,true,5000,5); % connect to vrep server
if (id>-1)
    disp('Connected to remote API server');
    % timestep
    dt = 0.01;
    vrep.simxSetFloatingParameter(id,vrep.sim_floatparam_simulation_time_step,dt,vrep.simx_opmode_oneshot_wait);
    % set sychronous mode
    vrep.simxSynchronous(id,true);
    % start the simulation:
    vrep.simxStartSimulation(id,vrep.simx_opmode_oneshot_wait);
    % trigger one step simulation and pause a short time to clear dirty
    % data in vrep buffer
    vrep.simxSynchronousTrigger(id);
%     pause(0.3);

    handles = two_link_workcell_init(vrep,id);
    armJoints = handles.armJoints;
    forcesensor = handles.forcesensor;
    x0 = [0;0;0;0];
    tau = [0;0];

    % set joint control mode in vrep
    if control_mode == 2
        for i=1:2
          vrep.simxSetObjectIntParameter(id,armJoints(i),2001,1,vrep.simx_opmode_oneshot_wait);
        end
    else
        for i=1:2
          vrep.simxSetObjectIntParameter(id,armJoints(i),2001,0,vrep.simx_opmode_oneshot_wait);
        end
    end
    % get init joint states
    for i=1:2
        [res,x0(i)] = vrep.simxGetJointPosition(id,armJoints(i),vrep.simx_opmode_buffer);vrchk(vrep, res, true);
        [res,x0(i+2)] = vrep.simxGetObjectFloatParameter(id,armJoints(i),2012,vrep.simx_opmode_buffer);vrchk(vrep, res, true);
        [res,tau(i)] = vrep.simxGetJointForce(id,armJoints(i),vrep.simx_opmode_buffer);vrchk(vrep, res, true);
    end
    % get external force in force sensor frame
    [res,state,fexternal,tauexternal] = vrep.simxReadForceSensor(id,forcesensor,vrep.simx_opmode_buffer);vrchk(vrep, res, true);
    % transform external force into the robot base frame(here the robot base frame coincides with the world frame in vrep)
    [res,eulerAngles]=vrep.simxGetObjectOrientation(id,forcesensor,-1,vrep.simx_opmode_buffer);vrchk(vrep, res, true);
    forcesensorrotation = rotx(double(eulerAngles(1)*180/pi))*roty(double(eulerAngles(2)*180/pi))*rotz(double(eulerAngles(3)*180/pi));
    fexternal_inertial = forcesensorrotation*[fexternal(1);fexternal(2);fexternal(3)];
    raw_externalforces(1,:) = fexternal_inertial(1)*ones(1,windowSize);
    raw_externalforces(2,:) = fexternal_inertial(3)*ones(1,windowSize);
    % record data
    recordData.x = [];
    recordData.tau = [];
    recordData.fmea = [];
    recordData.u = [];
    recordData.t = [];
    % max simulation duration
    maxsimtime = 20;
    % current time
    current_time = 0;
    % reference joint trajectory
    qref = x0(1:2);
    qdotref = [0;0];
    qdotdotref = [0;0];
    % reference ee trajectory
    eeposref = Direct_Kinematics(x0(1:2),[l1;l2]);
    eevref = [0;0];
    eearef = [0;0];
    % for impedance model
    xe_admittance = [0;0];
    xedot_admittance = [0;0];
    xedotdot_admittance = [0;0];
    q_admittance = [0;0];
    qdot_admittance = [0;0];
    qdotdot_admittance = [0;0];

    while vrep.simxGetConnectionId(id)~=-1
        % get robot joint states
        for i=1:2
            [res,x0(i)] = vrep.simxGetJointPosition(id,armJoints(i),vrep.simx_opmode_buffer);vrchk(vrep, res, true);
            [res,x0(i+2)] = vrep.simxGetObjectFloatParameter(id,armJoints(i),2012,vrep.simx_opmode_buffer);vrchk(vrep, res, true);
            [res,tau(i)] = vrep.simxGetJointForce(id,armJoints(i),vrep.simx_opmode_buffer);vrchk(vrep, res, true);
        end
        % get external force
        [res,state,fexternal,tauexternal] = vrep.simxReadForceSensor(id,forcesensor,vrep.simx_opmode_buffer);vrchk(vrep, res, true);

        % transfor external force form force sensor frame to robot base frame
        [res,eulerAngles]=vrep.simxGetObjectOrientation(id,forcesensor,-1,vrep.simx_opmode_buffer);vrchk(vrep, res, true);
        forcesensorrotation = rotx(double(eulerAngles(1)*180/pi))*roty(double(eulerAngles(2)*180/pi))*rotz(double(eulerAngles(3)*180/pi));
        fexternal_inertial = forcesensorrotation*[fexternal(1);fexternal(2);fexternal(3)];

        % shift raw external forces array
        for i=1:windowSize-1
            raw_externalforces(:,i) = raw_externalforces(:,i+1);
        end
        raw_externalforces(1,windowSize) = fexternal_inertial(1);
        raw_externalforces(2,windowSize) = fexternal_inertial(3);
        % force filtering
        filtered_externalforces_x = filter(b,a,raw_externalforces(1,:));
        filtered_externalforces_y = filter(b,a,raw_externalforces(2,:));
        filtered_externalforce = [filtered_externalforces_x(windowSize);filtered_externalforces_y(windowSize)];

        % feedbacked joint force in v-rep is the force sufferred by joint,
        % so change the sign to get the joint drive torque
        tau = -tau;

        if (current_time > maxsimtime)
            break;
        end

        % get reference trajectory point
        [eeposref,eevref,eearef] = ReferenceTrajectory(current_time,trajectory_mode);

       %% impedance control
        if control_mode == 0
            J =Geometric_Jacobian(x0(1:2),[l1;l2]);
            dJ_dt = Geometric_Jacobian_Derivative(x0,[l1;l2]);
            qdotdotref = J\(eearef - dJ_dt*x0(3:4));
            Mass = Mass_Matrix(x0(1:2),[I1 I2 l1 l2 m1 m2 a1 a2 g]');
            u = Impedance_Controller(x0,[l1;l2],[kpx;kpy;kdx;kdy],eeposref,eevref)...
                + Mass*qdotdotref...
                + Coriolis_Centrifugal_Torque(x0,[l1;l2;m1;m2;a1;a2;g])...
                + Gravity_Torque(x0(1:2),[l1;l2;m1;m2;a1;a2;g]);
            % limit joint tourque
            for i=1:2
                if u(i) > 25
                    u(i) = 25;
                elseif u(i) < -25
                    u(i) = -25;
                end
            end
            % send torque command to vrep
            for i=1:2
                if u(i) < 0
                    res = vrep.simxSetJointTargetVelocity(id,armJoints(i),-99999,vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
                    res = vrep.simxSetJointForce(id,armJoints(i),-u(i),vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
                else
                    res = vrep.simxSetJointTargetVelocity(id,armJoints(i),99999,vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
                    res = vrep.simxSetJointForce(id,armJoints(i),u(i),vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
                end
            end
       %% admittance control
        elseif control_mode == 1
            % --- velocity command --- (qdot_d + qdot_admittance)
            % compute qdotdot corresponds to impedance model
            qdotdot_admittance =  Admittance_Controller(x0,filtered_externalforce,[I1 I2 l1 l2 m1 m2 a1 a2 g]',[mx;my;kpx;kpy;kdx;kdy],eeposref,eevref);
            J =Geometric_Jacobian(x0(1:2),[l1;l2]);
            dJ_dt = Geometric_Jacobian_Derivative(x0,[l1;l2]);
            qdotdot_tracking =  J\(eearef - dJ_dt*x0(3:4));
            qdot_admittance = qdot_admittance + (qdotdot_admittance)*dt;
            qdotref  = J\eevref + qdot_admittance;
            u = qdotref;
            % send velocity command to vrep
            for i=1:2
                vrep.simxSetJointTargetVelocity(id,armJoints(i),qdotref(i),vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
            end
        elseif control_mode == 2
            % --- position command ---  qd + q_admittance
            %compute qdotdot corresponds to impedance model
            qdotdot_admittance =  Admittance_Controller(x0,filtered_externalforce,[I1 I2 l1 l2 m1 m2 a1 a2 g]',[mx;my;kpx;kpy;kdx;kdy],eeposref-eevref*dt,eevref-eearef*dt);%-eevref*dt -eearef*dt
            J =Geometric_Jacobian(x0(1:2),[l1;l2]);
            dJ_dt = Geometric_Jacobian_Derivative(x0,[l1;l2]);
            qdotdot_tracking =  J\(eearef - dJ_dt*x0(3:4));
            q_admittance = q_admittance + qdot_admittance*dt +  1/2*qdotdot_admittance*dt^2;
            qdot_admittance = qdot_admittance + qdotdot_admittance*dt;
            qref = Inverse_Kinematics(x0(1:2),eeposref,[l1 l2]) + q_admittance;
            u = qref;
            % send position command to vrep
            for i=1:2
                vrep.simxSetJointTargetPosition(id,armJoints(i),u(i),vrep.simx_opmode_oneshot);vrchk(vrep, res, true);
            end
        end

        %% record data
        recordData.x = [recordData.x,x0];
        recordData.tau = [recordData.tau,tau];
        recordData.fmea = [recordData.fmea,raw_externalforces(:,end)];
        recordData.u = [recordData.u,u];
        recordData.t = [recordData.t,current_time];

        % update current_time
        current_time = current_time + dt;
        vrep.simxSynchronousTrigger(id);
    end
    % display result
    titles = {'q1','q2','dq1','dq2'};
    figure('name','state');
    for i=1:4
        subplot(2,2,i);
        plot(recordData.t,recordData.x(i,:));
        title(titles{i});
        xlabel('t/s');
    end

    titles = {'u1','u2'};
    figure('name','input');
    for i=1:2
        subplot(1,2,i);
        plot(recordData.t,recordData.u(i,:));
        title(titles{i});
        xlabel('t/s');
    end

    figure('name','external force');
    subplot(1,2,1);
    plot(recordData.t,recordData.fmea(1,:));
    title('fextx');
    xlabel('t/s');
    hold off

    subplot(1,2,2);
    plot(recordData.t,recordData.fmea(2,:));
    title('fexty');
    xlabel('t/s');
    hold off
else
    disp('Failed connecting to remote API server1');
end
vrep.simxStopSimulation(id,vrep.simx_opmode_oneshot_wait);
vrep.simxFinish(id);
vrep.delete();