function [Out1,Out2]=qsynth(Action,In1,In2)
% Dick Benson
% Design single terminated elliptic passive LC low pass filters.
% Copyright 2001-2013 The MathWorks, Inc.
% [L,C] = qsynth('getLC');

    if nargin ==0
       Action = 'init';
    end;

    if ~strcmp(Action,'init')
      Hqsynth_ = get(gcf,'userdata');
    end

    % handle index definitions
    i_f1      = 1;
    i_topology   = 2;
    i_order   = 3;
    i_synthpb = 4;
    i_ripl    = 5;
    i_ripv    = 6;
    i_stopl   = 7;
    i_stopv   = 8;
    i_ax1     = 9;
    i_pl1     = 10;
    i_pl2     = 11;
    i_fcl     = 12;
    i_fcv     = 13;
    i_lv      = 14;
    i_cv      = 15;
    i_lcl     = 16;
    i_testpb  = 17;

    if strcmp(Action,'init')

        Hqsynth_(i_f1)= figure('Color',[0 0 0] ,'Name','Qsynth: Single Terminated Ladder Elliptic Filter',...
                          'Position',[20 20 640 440],...
                          'Resize','off',...
                          'Color',[0 0 1],...
                          'NumberTitle','off','visible','on');
                          % ###### userdata contains vector of handles if desired

        Hqsynth_(i_topology)=uicontrol('Style','popup','visible','on',...
                                    'String','RIN=1,RL=inf|RIN=0,RL=1',...
                                    'Position',[10 400 105 20],...
                                    'BackgroundColor',[0 0 0],...
                                    'ForeGroundColor',[0 1 1],...
                                    'Value',1,...
                                    'callback','qsynth(''clear'')',...
                                    'HorizontalAlignment','Left');

        Hqsynth_(i_order)=uicontrol('Style','popup','visible','on',...
                                'String','Order:3|Order:5|Order:7|Order:9|Order:11',...
                                'Position',[10 380 105 20],...
                                'BackgroundColor',[0 0 0],...
                                'ForeGroundColor',[0 1 1],...
                                'Value',4,...
                                'HorizontalAlignment','Left');

        Hqsynth_(i_ripl) =  uicontrol('Style','text','visible','on',...
                                'String','Ripple (dB):',...
                                'Position',[10,355,60,18],...
                                'BackGroundColor',[1 1 1]*0.8,...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left');

         Hqsynth_(i_ripv) = uicontrol('Style','edit','visible','on',...
                                'Position',[70,355,45,18],...
                                'String','1',...
                                'BackgroundColor',[1 1 0],...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','center',...
                                'Max',1);

         Hqsynth_(i_stopl) =  uicontrol('Style','text','visible','on',...
                                'String','Stop (dB):',...
                                'Position',[10,330,60,18],...
                                'BackGroundColor',[1 1 1]*0.8,...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left');

         Hqsynth_(i_stopv) = uicontrol('Style','edit','visible','on',...
                                'Position',[70,330,45,18],...
                                'String','60',...
                                'BackgroundColor',[1 1 0],...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','center',...
                                'Max',1);

         Hqsynth_(i_fcl) =  uicontrol('Style','text','visible','on',...
                                'String','Cutoff (Hz):',...
                                'Position',[10,305,60,18],...
                                'BackGroundColor',[1 1 1]*0.8,...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left');

         Hqsynth_(i_fcv) = uicontrol('Style','edit','visible','on',...
                                'Position',[70,305,45,18],...
                                'String','1.0e6',...
                                'BackgroundColor',[1 1 0],...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','center',...
                                'Max',1);

          Hqsynth_(i_lcl) =  uicontrol('Style','text','visible','on',...
                                'String','              L                         C',...
                                'Position',[10,280,165,18],...
                                'BackGroundColor',[1 1 1]*0.8,...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left');

        Hqsynth_(i_lv)  =  uicontrol('Style','edit','visible','on',...
                                'Position',[10,35,85,240],...   % was 80
                                'String','',...
                                'BackgroundColor',[1 1 0],...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left',...
                                'Max',14);

        Hqsynth_(i_cv)  =  uicontrol('Style','edit','visible','on',...
                                'Position',[100,35,85,240],...  % was 95,35,80
                                'String','',...
                                'BackgroundColor',[1 1 0],...
                                'ForeGroundColor',[0 0 0],...
                                'HorizontalAlignment','left',...
                                'Max',14);

        Hqsynth_(i_ax1)=axes('Units','pixels','Position',[240,80,370,300],...
                                'Box','on',...
                                'visible','on',...
                                'NextPlot','add',...
                                'DrawMode','fast',...
                                'Color',[0 0 0],...
                                'TickDir','out',...
                                'YlimMode','auto',...
                                'XlimMode','auto',...
                                'Xcolor',[1 1 1],...
                                'Ycolor',[1 1 1],...
                                'FontSize',12);

        title('Singly Terminated Elliptic Filter ','color',[1 1 1],'fontsize',9);
        xlabel('Frequency in Hz','fontsize',9);
        ylabel('Magnitude in dB','fontsize',9)

        Hqsynth_(i_pl1) = plot(Hqsynth_(i_ax1),0,0,'clipping','on',...
                                  'Color',[0 1 0],...
                                  'erasemode','xor','visible','on');

        Hqsynth_(i_pl2) = plot(Hqsynth_(i_ax1),0,0,'clipping','on',...
                                  'Color',[1 0 0],...
                                  'erasemode','xor','visible','on');

        Hqsynth_(i_synthpb) = uicontrol('Style','Pushbutton',...
                              'Position',[10,10,105,20],...
                              'String','Synthesize',...
                              'Callback','qsynth(''synth'');');

        Hqsynth_(i_testpb) = uicontrol('Style','Pushbutton',...
                              'Position',[120,10,105,20],...
                              'String','Test',...
                              'Callback','qsynth(''test'');');

        set(Hqsynth_(i_f1),'userdata',Hqsynth_);
        zoom on

    elseif strcmp(Action,'synth')
       fmtstr = '%8.5g';

       n  = get(Hqsynth_(i_order),'value')*2+1;          % only odd orders supported
       rp = str2num(get(Hqsynth_(i_ripv),'string'));
       rs = str2num(get(Hqsynth_(i_stopv),'string'));
       wn = 2*pi*str2num(get(Hqsynth_(i_fcv),'string'));

       [b,a]=ellip(n,rp,rs,wn,'s');
       save lpf_coefs b a
       w = 0:(wn/25):10*wn;
       h = freqs(b,a,w);

       set(Hqsynth_(i_ax1),'Ylim',[-rs-20,5],'Xlim',[0,10*wn/(2*pi)]);
       set(Hqsynth_(i_pl1),'xdata',w/(2*pi),'ydata',20*log10(h));

       topology = get(Hqsynth_(i_topology),'value');
       if topology ==1
          %  Rin = 1, Rl = inf , Vsource with 1 ohm Z0, no load
          % must be odd order
          m     = b(2:length(b));
          zeros = roots(m);
          m2 = [];
          n2 = [];
          for i=1:length(a)
              if rem(i,2)
                 n2=[n2,a(i),0];
              else
                 m2=[m2,a(i),0];
              end;
          end;
          m2=m2(1:(length(m2)-1));

          % remove admittance pole @ s=inf
          % z=m2/n2  y = n2/m2 yp = n2/m2-c1s

          C(1) = n2(1)/m2(1);  % shunt capacitor
          n2 = n2 - conv(m2,[C(1),0]);
          % now it gets tricky ....
          for i = 1:(n-1)/2
              % need an xmission zero @ w=wx
              wx = j*abs(zeros(2*i)); % pick an xmission zero....
              % do a partial removal of the impedance pole @ inf
              % first evaluate impedance @ w=wx
              zmag = polyval(m2,wx)/polyval(n2,wx);
              L(2*i) = zmag/wx;

              if i==1
                 n2=n2(3:length(n2));
              end;
              % this partially removes the pole
              m2 = m2-conv(n2,[L(2*i),0]);

              % now remaining admittance has a pole at w=wx
              % remove it by shunting a series LC

              % first must determine L & C values
              [q,r1]= deconv(m2,[1 0 abs(wx)^2]);  % r should be zero
              k    = n2(1);
              n2p  = n2/k;
              n2p  = n2p(1:length(n2p)-1);
              L(2*i+1)=1/(k*(polyval(n2p,wx)/polyval(q,wx)));
              C(2*i+1)= 1/(L(2*i+1)*(abs(wx)^2));

              % then remove the addmittance.... not obvious!
              [n2pp,r2] = deconv(( n2p-q/(k*L(2*i+1)) )*k,[1 0 1/(C(2*i+1)*L(2*i+1))]);
              n2 = conv(n2pp,[1 0]);
              m2 = q;
          end;
          L=L';
          C=C';

       elseif topology==2
          % Rin=0, Rl=1 (Vsource input, 1 ohm load)
          %
          m     = b(2:length(b)); % drop leading 0.0 term
          zeros = roots(m);       %
          m2 = [];
          n2 = [];
          for i=1:length(a)
              if rem(i,2)
                 n2=[n2,a(i),0];
              else
                 m2=[m2,a(i),0];
              end;
          end;

          m2=m2(1:(length(m2)-1));
          for i = 1:(n-1)/2
              % need an xmission zero @ w=wx
              wx = j*abs(zeros(2*i)); % pick an xmission zero....

              % do a partial removal of the impedance pole @ inf
              % first evaluate impedance @ w=wx
              zmag = polyval(n2,wx)/polyval(m2,wx);
              L(2*i-1) = zmag/wx;
              % this partially removes the pole

              n2 = n2-conv(m2,[L(2*i-1),0]);
              % now remaining admittance has a pole at w=wx

              % remove it by shunting a series LC
              % first must determine L & C values
              [q,r1] = deconv(n2,[1 0 abs(wx)^2]);  % r should be zero
              k = polyval(m2,wx)/polyval(conv(q,[1 0]),wx);
              L(2*i) = 1/k;
              C(2*i) = 1/(L(2*i)*(abs(wx)^2));
              [p,r2] = deconv( m2-k*conv(q,[1 0])   ,  [1 0 abs(wx)^2]);
              m2=p;
              n2=q;
          end;
          L(n)= n2(1)/m2;     % final series L .... amazing .... absolutly amazing!
          L = fliplr(L)';     % need to reverse this network
          C = [0,fliplr(C)]'; %

       end;
       set(Hqsynth_(i_ripl),'userdata',L);
       set(Hqsynth_(i_ripv),'userdata',C);

       s='';
       for i=1:(length(L)-1)
           s=[s,sprintf([fmtstr,'\n'],L(i))];
       end;
       s=[s,sprintf(fmtstr,L(length(L)))];
       set(Hqsynth_(i_lv),'string',s);
       s='';
       for i=1:(length(C)-1)
           s=[s,sprintf([fmtstr,'\n'],C(i))];
       end;
       s=[s,sprintf(fmtstr,C(length(C)))];
       set(Hqsynth_(i_cv),'string',s);

    elseif strcmp(Action,'getLC');
      % Out1 = get(Hqsynth_(i_ripl),'userdata'); % inductors
      % Out2 = get(Hqsynth_(i_ripv),'userdata'); % capacitors
      % huh? ,  should get it from strings in edit fields
       caps = get(Hqsynth_(i_cv),'string');
       inds = get(Hqsynth_(i_lv),'string');

       for i=1:length(caps(:,1))
          Out2(i) = str2num(caps(i,:));
       end;

       for i=1:length(inds(:,1))
          Out1(i) = str2num(inds(i,:));
       end;

    elseif strcmp(Action,'test')
       topology = get(Hqsynth_(i_topology),'value');   % network topology
       wn = 2*pi*str2num(get(Hqsynth_(i_fcv),'string'));
       w  = 0:(wn/25):10*wn;
       [L,C]=qsynth('getLC');
       if topology ==1
          % test the synthesis
          % begin with source resistor
          [a,b,c,d,q]=pabcd([],[],[],[],[],'rs',1);
          % then the shunt capacitor
          [a,b,c,d,q]=pabcd(a,b,c,d,q,'cp',C(1));
          for i = 1:(length(L)-1)/2
              % series L
              [a,b,c,d,q]=pabcd(a,b,c,d,q,'ls',L((i*2)));
              % shunt series LC
              [a,b,c,d,q]=pabcd(a,b,c,d,q,'lcs',L((i*2)+1),C((i*2)+1));
          end;

       elseif topology ==2
          a = []; b = []; c = []; d = []; q = [];
          for i = 1:(length(L)-1)/2
              % series L
              [a,b,c,d,q]=pabcd(a,b,c,d,q,'ls',L((i*2-1)));
              % shunt series LC
              [a,b,c,d,q]=pabcd(a,b,c,d,q,'lcs',L((i*2)),C((i*2)));
          end;
          % finish with a series L
          [a,b,c,d,q]=pabcd(a,b,c,d,q,'ls',L(length(L)));
          % and a 1 ohm shunt R load resistor
          [a,b,c,d,q]=pabcd(a,b,c,d,q,'rp',1);
       end;

       % [mag,phase]   =   bode(q,a,w);
       h = freqs(q,a,w);
       set(Hqsynth_(i_pl2),'xdata',w/(2*pi),'ydata',20*log10(h));

    elseif strcmp(Action,'clear')
       set(Hqsynth_(i_ripl),'userdata',[]);
       set(Hqsynth_(i_ripv),'userdata',[]);
       set(Hqsynth_(i_cv),'string','');
       set(Hqsynth_(i_lv),'string','');
       set(Hqsynth_(i_pl2),'xdata',[],'ydata',[]);
       set(Hqsynth_(i_pl1),'xdata',[],'ydata',[]);

    else
       disp([Action,' not reconized in qsynth'])
    end;

% end qsynth

function [aout,bout,cout,dout,qout]=pabcd(ain,bin,cin,din,qin,element,val1,val2,val3)
% [aout,bout,cout,dout,qout]=pabcd(ain,bin,cin,din,qin,element,val1,val2)
% Uses chain matrix parameters to cascades abcd with denominator q (in)
% with an element that has val1 (val2 for LC).
% Useful for computing functions & parameters of ladder networks.
% element code   description
%     rs       = series resistor
%     lcs      = shunt series lc val1=l val2=c
%     ls       = series l
%     cp       = shunt c
%     rp       = shunt r
%     lcp      = series chunt lc val1=l val2=c val3 = r (in series with l for loss)
%
%     a*e2-b*i2 = e1
%     c*e2-d*i2 = i1
%
% Dick Benson
%

      if strcmp(element,'rs')
           % rs  = series resistor
           a = 1;
           b = val1;
           c = 0;
           d = 1;
           q = 1;
      elseif strcmp(element,'lcs')
           % lcs = shunt series lc
           % l=val1 c=val2
           q = [val1*val2, 0 , 1];
           a = q;
           b = 0;
           c = [val2 0];
           d = q;

      elseif strcmp(element,'lcp')
          % l=val1 c=val2
           q = [val1*val2 val2*val3 1];
           a = q;
           b = [val1 val3];
           c = 0;
           d = q;
      elseif strcmp(element,'ls')
           % ls = series l
           a = 1;
           b = [val1 0];
           c = 0;
           d = 1;
           q = 1;
      elseif strcmp(element,'cp')
           % cp = shunt c
           q = 1;
           a = 1;
           b = 0;
           c = [val1 0];
           d = 1;
      elseif strcmp(element,'rp')
           % rp = shunt r
           q = 1;
           a = 1;
           b = 0;
           c = 1/val1;
           d = 1;
      else
         disp([element,' not supported in pabcd '])
      end;

      if isempty(ain)
         aout=a;
         bout=b;
         cout=c;
         dout=d;
         qout=q;
      else
         aout = padd(conv(ain,a),conv(bin,c));
         bout = padd(conv(ain,b),conv(bin,d));
         cout = padd(conv(cin,a),conv(din,c));
         dout = padd(conv(cin,b),conv(din,d));
         qout = conv(qin,q);
      end;

% end function

function pout = padd(p1,p2)
    lp1 = length(p1);
    lp2 = length(p2);
    if lp1==lp2
       pout = p1+p2;
    elseif lp2>lp1
       pout = p2;
       pout((lp2-lp1+1):lp2)=pout((lp2-lp1+1):lp2)+p1;
    else
       pout = p1;
       pout((lp1-lp2+1):lp1)=pout((lp1-lp2+1):lp1)+p2;
    end;
% end padd


function [x, infos] = mu_conv_nmf(V, rank, t, in_options)
% Multiplicative update (MU) based convolutive non-negative matrix factorization (MU-Conv-NMF).
%
% The problem of interest is defined as
%
%
% Given a non-negative matrix V, factorized non-negative matrices {W, H} are calculated.
%
%
% Inputs:
%       matrix      V
%       rank        rank
%
% Output:
%       w           solution of w
%       infos       information
%
% References:
%
%
% This file is part of NMFLibrary
%
% This file has been ported by H.Kasai from
% convNMF_MM1.m and convNMF_MM2.m at https://github.com/lyn202206/ADMM-Convolutive-NMF.
%
% Ported by H.Kasai on June 29, 2022
%
% Change log:
%
%

    % set dimensions and samples
    [m, n] = size(V);

    % set local options
    local_options = [];
    local_options.metric_type = 'beta-div';
    local_options.d_beta = 2;
    local_options.alg = 'type1';

    % check input options
    if ~exist('in_options', 'var') || isempty(in_options)
        in_options = struct();
    end
    % merge options
    options = mergeOptions(get_nmf_default_options(), local_options);
    options = mergeOptions(options, in_options);

    % initialize factors
    init_options = options;
    if ~isfield(options, 'x_init')
        W = zeros(m, rank, t);
        for i = 1 : t
            [init_factors, ~] = generate_init_factors(V, rank, init_options);
            W(:, :, i) = init_factors.W;
        end
        H = init_factors.H;
    else
        W = init_options.x_init.W;
        H = init_options.x_init.H;
    end

    % initialize
    epoch = 0;
    grad_calc_count = 0;

    options = check_divergence(options);
    sub_mode = sprintf('beta=%.1f', options.d_beta);
    if ~strcmp(options.metric_type, 'beta-div')
        sub_mode = options.metric_type;
    end
    method_name = sprintf('MU-Conv (%s)', sub_mode);

    if options.verbose > 0
        fprintf('# %s: started ...\n', method_name);
    end

    % initialize for this algorithm
    V_hat = zeros(m, n);
    for i = 0 : t-1
        V_hat = V_hat + W(:, :, i+1) * shift_t(H, i);
    end

    % store initial info
    clear infos;

    [Wcon, Hcon] = reconstruct_wh(W, H, t);
    [infos, f_val, optgap] = store_nmf_info(V, Wcon, Hcon, [], options, [], epoch, grad_calc_count, 0);

    if options.verbose > 1
        fprintf('MU-Conv (%s): Epoch = 0000, cost = %.16e, optgap = %.4e\n', sub_mode, f_val, optgap);
    end

    % set start time
    start_time = tic();

    % main loop
    while true

        % check stop condition
        [stop_flag, reason, max_reached_flag] = check_stop_condition(epoch, infos, options);
        if stop_flag
            display_stop_reason(epoch, infos, options, method_name, reason, max_reached_flag);
            break;
        end

        % update H sequentially
        for i = 1 : n

            num = zeros(rank, 1);
            denom = num;
            if i <= n-t+1
                for n_prime = i : i+t-1
                    num = num + W(:, :, n_prime-i+1)' * ((V(:, n_prime) + eps) .* (V_hat(:, n_prime) + eps) .^ (options.d_beta - 2));
                    denom = denom + W(:, :, n_prime-i+1)'*(V_hat(:, n_prime)+eps).^(options.d_beta-1);
                end
            else
                for n_prime = i : n
                    num = num + W(:, :, n_prime-i+1)' * ((V(:, n_prime) + eps) .* (V_hat(:, n_prime) + eps) .^ (options.d_beta - 2));
                    denom = denom + W(:, :, n_prime-i+1)'*(V_hat(:, n_prime) + eps) .^ (options.d_beta-1);
                end
            end

            if strcmp(options.alg, 'type1')
                h_n_old = H(:, i);
                H(:, i) = H(:, i) .* (num./denom) .^ gamma_beta(options.d_beta);

                if i <= n-t+1
                    for n_prime = i:i+t-1
                        V_hat(:, n_prime) = max(V_hat(:, n_prime) + W(:, :, n_prime - i + 1) * (H(:, i) - h_n_old), 0);
                        % max(.,0) ensures the nonnegativity
                    end
                else
                    for n_prime = i:n
                        V_hat(:, n_prime) = max(V_hat(:, n_prime) + W(:, :, n_prime - i + 1) * (H(:, i) - h_n_old), 0);
                    end
                end
            else
                H(:, i) = H(:, i) .* (num./denom) .^ gamma_beta(options.d_beta);
            end

        end

        if strcmp(options.alg, 'type2')

            V_hat = zeros(m, n);
            for i = 0 : t-1
                V_hat = V_hat + W(:, :, i+1) * shift_t(H, i);
            end
        end

        % update W
        for i = 0 : t - 1
            W_t_old = W(:, :, i+1);
            H_shift_t = shift_t(H, i);
            W(:, :, i+1) = W(:, :, i+1) .* (((((V + eps) .* (V_hat + eps) .^ (options.d_beta - 2)) * H_shift_t') ...
                ./ ((V_hat + eps).^(options.d_beta - 1) * H_shift_t'))) .^ gamma_beta(options.d_beta);
            V_hat = max(V_hat + (W(:, :, i+1) - W_t_old) * H_shift_t, 0);
            % max(.,0) ensures nonnegativity
        end

        [W, H] = renormalize_convNMF(W,H);

        if strcmp(options.alg, 'type1')
            % recalate the V_hat
            X_hat = zeros(size(V));
            for i=0:t-1
                tW = W(:, :, i+1);
                tH = shift_t(H, i);
                X_hat = X_hat + tW * tH;
            end
            V_hat = max(X_hat,0);
        end

        % measure gradient calc count
        grad_calc_count = grad_calc_count + m*n;

        % measure elapsed time
        elapsed_time = toc(start_time);

        % update epoch
        epoch = epoch + 1;

        % store info
        [Wcon, Hcon] = reconstruct_wh(W, H, t);
        infos = store_nmf_info(V, Wcon, Hcon, [], options, infos, epoch, grad_calc_count, elapsed_time);

        % display info
        display_info(method_name, epoch, infos, options);

    end

    x.W = W;
    x.H = H;

end

function [W_concat, H_concat] = reconstruct_wh(W, H, t)

    W_concat = [];
    H_concat = [];
    for j = 1 : t
        W_concat = [W_concat W(:, :,j)];
        H_concat = [H_concat; shift_t(H, j-1)];
    end
end


%% FUNCTION Logistic_Trace
%   Trace-Norm Regularized Learning with Logistic Loss.
%
%% OBJECTIVE
%   argmin_{W,C} { sum_i^t (- sum(log (1./ (1+ exp(-X{i}*W(:, i) - Y{i} .* C(i)))))/length(Y{i}))
%            + rho1 * \|W\|_*}
%  where \|W\|_* = sum(svd(W, 0)) is the trace norm
%
%% INPUT
%   X: {n * d} * t - input matrix
%   Y: {n * 1} * t - output matrix
%   rho1: trace norm regularization parameter
%
%% OUTPUT
%   W: model: d * t
%   C: model: 1 * t
%   funcVal: function value vector.
%
%% LICENSE
%   This program is free software: you can redistribute it and/or modify
%   it under the terms of the GNU General Public License as published by
%   the Free Software Foundation, either version 3 of the License, or
%   (at your option) any later version.
%
%   This program is distributed in the hope that it will be useful,
%   but WITHOUT ANY WARRANTY; without even the implied warranty of
%   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%   GNU General Public License for more details.
%
%   You should have received a copy of the GNU General Public License
%   along with this program.  If not, see <http://www.gnu.org/licenses/>.
%
%   Copyright (C) 2011 - 2012 Jiayu Zhou and Jieping Ye
%
%   You are suggested to first read the Manual.
%   For any problem, please contact with Jiayu Zhou via jiayu.zhou@asu.edu
%
%   Last modified on June 3, 2012.
%
%% RELATED PAPERS
%
%   [1] Ji, S. and Ye, J. An Accelerated Gradient Method for Trace Norm Minimization, ICML 2009
%
%% RELATED FUNCTIONS
%  Least_Trace, init_opts

%% Code starts here
function [W, C, funcVal] = Logistic_Trace(X, Y, rho1, opts)

if nargin <3
    error('\n Inputs: X, Y, and rho1 should be specified!\n');
end
X = multi_transpose(X);

if nargin <4
    opts = [];
end

% initialize options.
opts=init_opts(opts);

if isfield(opts, 'rho_L2')
    rho_L2 = opts.rho_L2;
else
    rho_L2 = 0;
end

task_num  = length (X);
dimension = size(X{1}, 1);
funcVal = [];

%initialize a starting point
C0_prep = zeros(1, task_num);
for t_idx = 1: task_num
    m1 = nnz(Y{t_idx} == 1);
    m2 = nnz(Y{t_idx} == -1);
    if ( m1==0 || m2==0 )
        C0_prep(t_idx) = 0;
    else
        C0_prep(t_idx) = log(m1/m2);
    end
end

if opts.init==2
    W0 = zeros(dimension, task_num);
    C0 = zeros(1, task_num);
elseif opts.init== 0
    W0 = randn(dimension, task_num);
    C0 = C0_prep;
else
    if isfield(opts,'W0')
        W0=opts.W0;
        if (nnz(size(W0)-[dimension, task_num]))
            error('\n Check the input .W0');
        end
    else
        W0 = zeros(dimension, task_num);
    end
    if isfield(opts,'C0')
        C0=opts.C0;
    else
        C0=C0_prep;
    end
end

bFlag=0; % this flag tests whether the gradient step only changes a little

Wz= W0;
Cz= C0;
Wz_old = W0;
Cz_old = C0;

t = 1;
t_old = 0;
iter = 0;
gamma = 1;
gamma_inc = 2;

while iter < opts.maxIter
    alpha = (t_old - 1) /t;

    Ws = (1 + alpha) * Wz - alpha * Wz_old;
    Cs = (1 + alpha) * Cz - alpha * Cz_old;

    % compute function value and gradients of the search point
    [gWs, gCs, Fs ]  = gradVal_eval(Ws, Cs);

    % the Armijo Goldstein line search scheme
    while true
        [Wzp Wzp_tn] = trace_projection(Ws - gWs/gamma, 2 * rho1 / gamma);
        Czp = Cs - gCs/gamma;
        Fzp = funVal_eval  (Wzp, Czp);

        delta_Wzp = Wzp - Ws;
        delta_Czp = Czp - Cs;
        nrm_delta_Wzp = norm(delta_Wzp, 'fro')^2;
        nrm_delta_Czp = norm(delta_Czp, 'fro')^2;
        r_sum = (nrm_delta_Wzp+nrm_delta_Czp)/2;

        %         Fzp_gamma = Fs + trace(delta_Wzp' * gWs)...
        %             + trace(delta_Czp' * gCs)...
        %             + gamma/2 * nrm_delta_Wzp ...
        %             + gamma/2 * nrm_delta_Czp;
        Fzp_gamma = Fs + sum(sum(delta_Wzp .* gWs))...
            + sum(sum(delta_Czp .* gCs))...
            + gamma/2 * nrm_delta_Wzp ...
            + gamma/2 * nrm_delta_Czp;

        if (r_sum <=1e-20)
            bFlag=1; % this shows that, the gradient step makes little improvement
            break;
        end

        if (Fzp <= Fzp_gamma)
            break;
        else
            gamma = gamma * gamma_inc;
        end
    end

    Wz_old = Wz;
    Cz_old = Cz;
    Wz = Wzp;
    Cz = Czp;

    %funcVal = cat(1, funcVal, Fzp + rho1 * sum( svd(Wzp, 0) ));
    funcVal = cat(1, funcVal, Fzp + rho1 * Wzp_tn);

    if (bFlag)
        % fprintf('\n The program terminates as the gradient step changes the solution very small.');
        break;
    end

    % test stop condition.
    switch(opts.tFlag)
        case 0
            if iter>=2
                if (abs( funcVal(end) - funcVal(end-1) ) <= opts.tol)
                    break;
                end
            end
        case 1
            if iter>=2
                if (abs( funcVal(end) - funcVal(end-1) ) <=...
                        opts.tol* funcVal(end-1))
                    break;
                end
            end
        case 2
            if ( funcVal(end)<= opts.tol)
                break;
            end
        case 3
            if iter>=opts.maxIter
                break;
            end
    end

    iter = iter + 1;
    t_old = t;
    t = 0.5 * (1 + (1+ 4 * t^2)^0.5);

end

W = Wzp;
C = Czp;

% private functions
    function [grad_W, grad_C, funcVal] = gradVal_eval(W, C)
        grad_W = zeros(dimension, task_num);
        grad_C = zeros(1, task_num);
        lossValVect = zeros (1 , task_num);
        if opts.pFlag
            parfor i = 1:task_num
                [ grad_W(:, i), grad_C(:, i), lossValVect(:, i)] = unit_grad_eval( W(:, i), C(i), X{i}, Y{i});
            end
        else
            for i = 1:task_num
                [ grad_W(:, i), grad_C(:, i), lossValVect(:, i)] = unit_grad_eval( W(:, i), C(i), X{i}, Y{i});
            end
        end
        grad_W = grad_W+ rho_L2 * 2 * W;
        % here when computing function value we do not include
        % l1 norm.
        funcVal = sum(lossValVect) + rho_L2 * norm(W, 'fro')^2;
    end

    function [funcVal] = funVal_eval (W, C )
        funcVal = 0;
        if opts.pFlag
            parfor i = 1: task_num
                funcVal = funcVal + unit_funcVal_eval( W(:, i), C(i), X{i}, Y{i});
            end
        else
            for i = 1: task_num
                funcVal = funcVal + unit_funcVal_eval( W(:, i), C(i), X{i}, Y{i});
            end
        end
        % here when computing function value we do not include
        % l1 norm.
        funcVal = funcVal + rho_L2 * norm(W, 'fro')^2;
    end

end

function [ grad_w, grad_c, funcVal ] = unit_grad_eval( w, c, x, y)
%gradient and logistic evaluation for each task
m = length(y);
weight = ones(m, 1)/m;
weighty = weight.* y;
aa = -y.*(x'*w + c);
bb = max( aa, 0);
funcVal = weight'* ( log( exp(-bb) +  exp(aa-bb) ) + bb );
pp = 1./ (1+exp(aa));
b = -weighty.*(1-pp);
grad_c = sum(b);
grad_w = x * b;
end

function [ funcVal ] = unit_funcVal_eval( w, c, x, y)
%function value evaluation for each task
m = length(y);
weight = ones(m, 1)/m;
aa = -y.*(x'*w + c);
bb = max( aa, 0);
funcVal = weight'* ( log( exp(-bb) +  exp(aa-bb) ) + bb );
end


% Ryan Steindl based on Robotics Toolbox for MATLAB (v6 and v9)
%
% Copyright (C) 1993-2011, by Peter I. Corke
%
% This file is part of The Robotics Toolbox for MATLAB (RTB).
%
% RTB is free software: you can redistribute it and/or modify
% it under the terms of the GNU Lesser General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% RTB 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 Lesser General Public License for more details.
%
% You should have received a copy of the GNU Leser General Public License
% along with RTB.  If not, see <http://www.gnu.org/licenses/>.
%
% http://www.petercorke.com
function t = q2tr(q)

    q = double(q);
    s = q(1);
    x = q(2);
    y = q(3);
    z = q(4);

    r = [   1-2*(y^2+z^2)   2*(x*y-s*z) 2*(x*z+s*y)
        2*(x*y+s*z) 1-2*(x^2+z^2)   2*(y*z-s*x)
        2*(x*z-s*y) 2*(y*z+s*x) 1-2*(x^2+y^2)   ];
    t = eye(4,4);
    t(1:3,1:3) = r;
    t(4,4) = 1;
endfunction


% SP_BSPLINE_FLUID: Construct different pair of B-Splines spaces on the physical domain for fluid problems.
%
%   [spv, spp] = sp_bspline_fluid (elem_name, knots, nsub, ...
%                                               degreep, regularity, msh)
%
% INPUTS:
%
%   elem_name:  the name of the element. Right now 'TH' (Taylor-Hood),
%                'NDL' (Nedelec, 2nd family), 'RT' (Raviart-Thomas) and
%                'SG' (SubGrid) are supported.
%   knots:      knot vector of the coarse geometry.
%   nsub:       number of subdivisions of each interval.
%   degreep:    degree of the pressure space along each parametric direction
%   regularity: continuity of the pressure space along each parametric direction
%   msh:        msh object containing (in the field msh.qn) the points
%                along each parametric direction in the parametric
%                domain at which to evaluate, i.e. quadrature points
%                or points for visualization (see msh_cartesian).
%
% OUTPUT:
%
%   spv: object representing the discrete velocity function space (see sp_vector)
%   spp: object representing the discrete pressure function space (see sp_scalar)
%
%   For more details, see:
%      A. Buffa, C. de Falco, G. Sangalli,
%      IsoGeometric Analysis: Stable elements for the 2D Stokes equation
%      Internat. J. Numer. Methods Fluids, 2011
%
%      A. Bressan, G. Sangalli,
%      Isogeometric discretizations of the Stokes problem: stability
%       analysis by the macroelement technique
%      IMA J. Numer. Anal., 2013.
%
% Copyright (C) 2009, 2010, 2011 Carlo de Falco
% Copyright (C) 2011, 2015 Rafael Vazquez
%
%    This program is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.

%    This program is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with this program.  If not, see <http://www.gnu.org/licenses/>.

function [spv, spp] = sp_bspline_fluid (element_name, ...
                   knots, nsub_p, degree_p, regularity_p, msh)

% Construction of the knot vectors and the discrete space for the velocity
switch (lower (element_name))
  case {'th'}
    knotsp = kntrefine (knots, nsub_p-1, degree_p, regularity_p);
    spp = sp_bspline (knotsp, degree_p, msh);

    degree_v = degree_p + 1;
    regularity_v = regularity_p;
    nsub_v = nsub_p;
    knots_v = kntrefine (knots, nsub_v-1, degree_v, regularity_v);
    scalar_space = sp_bspline (knots_v, degree_v, msh);
    for idim = 1:msh.ndim
      scalar_spaces{idim} = scalar_space;
    end
    spv = sp_vector (scalar_spaces, msh);
    clear scalar_spaces scalar_space

  case {'sg'}
    knotsp = kntrefine (knots, nsub_p-1, degree_p, regularity_p);
    spp = sp_bspline (knotsp, degree_p, msh);

    degree_v = degree_p + 1;
    regularity_v = regularity_p+1;
    nsub_v = 2*nsub_p;
    knots_v = kntrefine (knots, nsub_v-1, degree_v, regularity_v);
    scalar_space = sp_bspline (knots_v, degree_v, msh);
    for idim = 1:msh.ndim
      scalar_spaces{idim} = scalar_space;
    end
    spv = sp_vector (scalar_spaces, msh);
    clear scalar_spaces scalar_space

  case {'ndl'}
% In this case the regularity is assigned first in the velocity space
    degree_h1 = degree_p + 1;
    regularity_h1 = regularity_p + 1;
    knots_h1 = kntrefine (knots, nsub_p-1, degree_h1, regularity_h1);
    [knots_hdiv, degree_hdiv] = knt_derham (knots_h1, degree_h1, 'Hdiv');

    degree_v = degree_h1;
    for idim = 1:msh.ndim
      knots_v{idim} = knots_hdiv{idim}{idim};
      for jdim = setdiff (1:msh.ndim, idim)
        knots_v{jdim} = sort ([knots_hdiv{idim}{jdim}, unique(knots_hdiv{idim}{jdim})]);
      end
      scalar_spaces{idim} = sp_bspline (knots_v, degree_v, msh);
    end
    spv = sp_vector (scalar_spaces, msh, 'div-preserving');
    clear scalar_spaces

    [knotsp, degp] = knt_derham (knots_h1, degree_h1, 'L2');
    spp = sp_bspline (knotsp, degp, msh, 'integral-preserving');

  case {'rt'}
% In this case the regularity is assigned first in the velocity space
    degree_h1 = degree_p + 1;
    regularity_h1 = regularity_p + 1;
    knots_h1 = kntrefine (knots, nsub_p-1, degree_h1, regularity_h1);

    [knots_v, degree_v] = knt_derham (knots_h1, degree_h1, 'Hdiv');
    for idim = 1:msh.ndim
      scalar_spaces{idim} = sp_bspline (knots_v{idim}, degree_v{idim}, msh);
    end
    spv = sp_vector (scalar_spaces, msh, 'div-preserving');
    clear scalar_spaces

    [knotsp, degp] = knt_derham (knots_h1, degree_h1, 'L2');
    spp = sp_bspline (knotsp, degp, msh, 'integral-preserving');

%     if (nargout == 3)
%       PI = b2nst__ (spp, knotsp, degree_p, msh);
%     end
  otherwise
    error ('sp_bspline_fluid: unknown element type')
end

end


 function [Gauss,terms,logp,comptime] = laplacedegenerate_ep(y,X,K,varargin)

% specialized version of our fast variant of
% Expectation Propagation for Logistic Regression (2 classes) with a (correlated) Laplace prior
% in the degenerate case, i.e., when the number of features >> number of samples
% (also works when number of samples > number of features, but then less stable!!!)
%
% input:
% labels    = an N x 1 vector of class labels [1,2]
% examples  = an N x M matrix of input data
% K         = the prior precision matrix of size M x M
%
% Note: the bias term should be explicitly added to K and the examples!
%
% regression parameters refers to betas
% auxiliary variables refers to u and v whose precision matrix auxK = inv(Lambda) couples the features.
%
% we have a precision matrix of the form
%
% | K_beta         |
% |        K_u     |
% |            K_v |
%
% priorGauss: struct with fields
%
%     hatK       diagonal of precision matrix (number of samples x 1);
%                   (initially zero)
%     diagK      diagonal of precision matrix (number of features x 1)
%                   (initially zero)
%
%     precision matrix of regression parameters K_beta = A' hatK A + diagK
%
%     h          canonical mean (number of features x 1)
%                   (initially zero)
%
%     auxK       precision matrix of auxiliary variables (number of features x number of features; sparse)
%                   (contains the covariance structure of interest)
%
%     A          feature matrix (number of samples x number of features)
%
% terms: struct with fields
%
%     hatK       as priorGauss
%     diagK      as priorGauss
%     hath       canonical mean (number of samples x 1)
%     h          canonical mean (number of features x 1)
%                                canonical mean of regression parameters = h + A' hath
%     auxK       as priorGauss, but then only diagonal elements (number of features x 1)
%
% opt: struct with fields (all optional, defaults in brackets)
%
%     maxstepsize   maximum step size [1]
%     fraction      fraction or power for fractional/power EP [1]
%     niter         maximum number of iterations [100]
%     tol           convergence criterion [1e-5]
%     nweights      number of points for numerical integration [20]
%     temperature   temperature for simulated annealing [1]
%     verbose       give output during run [1]
%
% Gauss: struct with fields as in priorGauss plus
%
%     hatB       diagonal of projected covariance (number of samples x 1)
%                                                               hatB = A Covariance A'
%     hatn       projected mean (number of samples x 1)
%                                                               hatn = A m
%     diagC      diagonal of covariance matrix of regression parameters (number of features x 1)
%     auxC       diagonal of covariance matrix of auxiliary variables (number of features x 1)
%
%
% logp: estimated marginal loglikelihood (implementation doubtful...)
%
% comptime: computation time

%% initialization

% parse opt
opt = [];
for i=1:2:length(varargin)
  opt.(varargin{i}) = varargin{i+1};
end

if ~isfield(opt,'maxstepsize'), opt.maxstepsize = 1;    end
if ~isfield(opt,'fraction'),    opt.fraction    = 0.95; end
if ~isfield(opt,'niter'),       opt.niter       = 1000; end
if ~isfield(opt,'tol'),         opt.tol         = 1e-5; end
if ~isfield(opt,'nweights'),    opt.nweights    = 50;   end
if ~isfield(opt,'temperature'), opt.temperature = 1;    end
if ~isfield(opt,'lambda'),      opt.lambda      = 1;    end
if ~isfield(opt,'verbose'),     opt.verbose      = 1;    end

if opt.verbose
  fprintf('starting EP\n');
end

tic

[nsamples,nfeatures] = size(X);

%% create priorGauss and terms

A = X.*repmat(y,1,nfeatures);

% construct Gaussian representation
priorGauss.A     = A;
priorGauss.hatK  = zeros(nsamples,1);
priorGauss.h     = zeros(nfeatures,1);
priorGauss.diagK = zeros(nfeatures,1);
priorGauss.auxK  = K;

% compute additional terms for the EP free energy
[cholK,dummy,S]  = chol(K,'lower');
LogPriorRestAux2 = (2*sum(log(full(diag(cholK))))); % -nfeatures*log(2*pi)); % redundant MvG

% construct term representation
terms.hatK  = ones(nsamples,1)/10;
terms.hath  = zeros(nsamples,1);
terms.diagK = ones(nfeatures,1)./(10*opt.lambda);
terms.auxK  = zeros(nfeatures,1);
terms.h     = zeros(nfeatures,1);

%% precompute points and weights for numerical integration

[xhermite,whermite] = gausshermite(opt.nweights);
xhermite = xhermite(:);       % nhermite x 1
whermite = whermite(:);       % nhermite x 1

[xlaguerre,wlaguerre] = gausslaguerre(opt.nweights);
xlaguerre = xlaguerre(:);     % nlaguerre x 1
wlaguerre = wlaguerre(:);     % nlaguerre x 1

% divide all canonical parameters by the temperature
if opt.temperature ~= 1,
   [priorGauss,terms] = correct_for_temperature(priorGauss,terms,opt.temperature);
end

%% build initial Gauss

% add terms to prior
Gauss        = update_Gauss(priorGauss,terms);
Gauss        = canonical_to_moments(Gauss);
[myGauss,ok] = project_all(Gauss,terms,opt.fraction);
if ~ok,
   error('improper cavity distributions\n');
end

prior = Gauss; % save prior Gauss

%% enter the iterations

logp     = 0;
logpold  = 2*opt.tol;
change   = 0;
teller   = 0;
stepsize = opt.maxstepsize;

while abs(logp-logpold) > stepsize*opt.tol && teller < opt.niter,

   teller    = teller+1;
   logpold   = logp;
   oldchange = change;

   % compute the new term approximation by applying an adf update on all the cavity approximations
   [fullterms,logreglogz,crosslogz] = ...
      adfupdate_all(myGauss,terms,xhermite,whermite,xlaguerre,wlaguerre,opt.fraction,opt.temperature);

   ok  = 0;
   ok1 = 1;
   ok2 = 1;

   while ~ok,

      % try to replace the old term approximations by the new term approximations and check whether the new Gauss is still fine

     [newGauss,newterms,ok1] = try_update(Gauss,fullterms,terms,stepsize);
     %[newGauss,newterms,ok1] = try_update(Gauss,fullterms,terms,0.5);

      if ok1,
         % compute all cavity approximations needed for the next EP updates and check whether they are all fine
         [newGauss,myGauss,logZappx,ok2] = try_project(newGauss,newterms,opt.fraction);
      end

      ok = (ok1 & ok2);
      if ok,    % accept

        terms = newterms;
        Gauss = newGauss;
        stepsize = min(opt.maxstepsize,stepsize*1.9);

      else  % try with smaller stepsize

        stepsize = stepsize/2;
         if ~ok1,
            fprintf('improper full covariance: lowering stepsize to %g\n',stepsize');
         elseif ~ok2,
            fprintf('improper cavity covariance: lowering stepsize to %g\n',stepsize');
         end
         if stepsize < 1e-10,
            warning('Cannot find an update that leads to proper cavity approximations');
            teller = opt.niter;
            break;
         end

      end

   end

   % compute marginal moments

   if ok,

      % compute marginal loglikelihood

      %logp = sum(logz) + sum(crosslogz) + logdet/2 + contribQ;
%      CorrTermCross1 = (myGauss.m.^2)./myGauss.diagC - (newGauss.m.^2)./newGauss.diagC + log(myGauss.diagC/newGauss.diagC);
%	   CorrTermCross2 = log(myGauss.auxC/newGauss.auxC);
%      CorrTerm  = 0.5*CorrTermCross1 + 0.5*(CorrTermCross2 + CorrTermCross2);

      logp      = (sum(crosslogz) + sum(logreglogz))./opt.fraction + logZappx + LogPriorRestAux2;

      if opt.verbose
        fprintf('%d: %g (stepsize: %g)\n',teller,logp,stepsize);
      end

      % check whether marginal loglikelihood is going up and down, if so, lower stepsize

      change = logp-logpold;
      if change*oldchange < 0,   % possibly cycling
         stepsize = stepsize/2;
      end
      oldchange = change;
   end

end

comptime = toc;

if opt.verbose
  fprintf('EP finished in %s seconds\n',num2str(comptime));
end

Gauss.prior = prior;

%%% END MAIN

%%%%%%%%%
%
% compute the cavity approximations that result when subtracting a fraction of the term approximations

function [myGauss,ok] = project_all(Gauss,terms,fraction)

if nargin < 3,
   fraction = 1;
end

% take out and project in moment form

% (1) regression parameters

[myGauss.hatB,myGauss.hatn] = ...
   rank_one_update(Gauss.hatB,Gauss.hatn,-fraction*terms.hatK,-fraction*terms.hath);

% (2) cross terms between regression parameters and auxiliary parameters

[myGauss.diagC,myGauss.m] = ...
   rank_one_update(Gauss.diagC,Gauss.m,-fraction*terms.diagK,-fraction*terms.h);

 myGauss.auxC = rank_one_update(Gauss.auxC,[],-fraction*terms.auxK);

% check whether all precision matrices are strictly positive definite

if nargout > 1,
   ok = (all(myGauss.hatB > 0) & all(myGauss.diagC > 0) & all(myGauss.auxC > 0));
end

%%%%%%%%%
%
% compute the new term approximation by applying an adf update on all the cavity approximations

function [fullterms,logreglogz,crosslogz] = ...
   adfupdate_all(myGauss,terms,xhermite,whermite,xlaguerre,wlaguerre,fraction,temperature)

if nargin < 8,
   temperature = 1;
end
if nargin < 7,
   fraction = 1;
end

%% (1) regression parameters

oldm     = myGauss.hatn;
oldC     = myGauss.hatB;
sqrtC    = sqrt(oldC);
nsamples = length(oldm);
nhermite = length(whermite);

% translate and scale the sample points to get the correct mean and variance

x = repmat(oldm,1,nhermite) + sqrtC*xhermite';

% compute the terms at the sample points

g = logist(x);   % returns - log (1 + exp(-x)) with special attention for very small and very large x

% correct for fraction and temperature and incorporate the sample weights

g           = fraction*g/temperature + log(repmat(whermite',nsamples,1));
maxg        = max(g,[],2);
g           = g-repmat(maxg,1,nhermite);
expg        = exp(g);
denominator = sum(expg,2);
neww        = expg./repmat(denominator,1,nhermite);

% compute the moments

Ex   = sum(x.*neww,2);
Exx  = sum(x.^2.*neww,2);
newm = Ex;
newC = Exx-Ex.^2;

% derive the term approximation from the change in mean and variance

[fullterms.hatK,fullterms.hath,logzextra] = compute_termproxy(newC,newm,oldC,oldm,fraction);

% contributions to marginal loglikelihood
logreglogz = maxg + log(denominator) + logzextra;

%% (2) cross terms between regression parameters and auxiliary variables

oldm      = myGauss.m;
oldC      = myGauss.diagC;
oldlambda = myGauss.auxC;
nfeatures = length(oldm);
nlaguerre = length(wlaguerre);

% this part heavily relies on the accompanying note
% basic idea:
% - the cavity approximation on U is an exponential distribution
% - we have analytical formulas for the moments of x conditioned upon U
% - marginal moments can then be computed through numerical integration with Gauss-Laguerre

% translate and scale the sample points to get the correct mean

U  = 2*oldlambda*xlaguerre';     % nfeatures x nlaguerre

mm = repmat(oldm,1,nlaguerre);
CC = repmat(oldC,1,nlaguerre);

% compute the partition function (integral over x) given U and turn this into weights required for computing the marginal moments

g           = (1-fraction)*log(U)/2 - fraction*mm.^2./(U + fraction*CC)/2 - log(U+ fraction*CC)/2;

g           = bsxfun(@plus,g,log(wlaguerre'));
maxg        = max(g,[],2);
g           = bsxfun(@minus,g,maxg);
expg        = exp(g);

denominator = sum(expg,2);

neww        = bsxfun(@rdivide,expg,denominator);

% compute the marginal moments through numerical integration

ExgU = mm.*U./(U + fraction*CC);
Ex   = sum(ExgU.*neww,2);

ExxgU = ExgU.^2 + CC.*U./(U+fraction*CC);
Exx   = sum(ExxgU.*neww,2);
EU    = sum(U.*neww,2);

newm      = Ex;
newC      = Exx-Ex.^2;
newlambda = EU/2;

% derive the term approximation from the change in mean and variance

[fullterms.diagK,fullterms.h,logzextra1] = compute_termproxy(newC,newm,oldC,oldm,fraction);
[fullterms.auxK,dummy,logzextra2] = compute_termproxy(newlambda,zeros(nfeatures,1),oldlambda,zeros(nfeatures,1),fraction);

crosslogz = maxg + log(denominator) + logzextra1 + 2*logzextra2;
% multiplied the last term by two

%%%%%%%%%%
%
% compute the moments corresponding to the canonical parameters

function [Gauss,logp] = canonical_to_moments(Gauss)

[nsamples,nfeatures] = size(Gauss.A);

%% (1) regression parameters

if nsamples > nfeatures,   % in the non-degenerate case, this direct route is more stable and faster

   scaledA     = Gauss.A.*(repmat(Gauss.hatK,1,nfeatures));
   K           = Gauss.A'*scaledA + diag(Gauss.diagK);
   [C,logdet1] = invert_chol(K);
   Gauss.m     = C*Gauss.h;
   Gauss.hatB  = zeros(nsamples,1);          % only need diagonal
   for k=1:nsamples,
      Gauss.hatB(k) = Gauss.A(k,:)*C*Gauss.A(k,:)';
   end
   Gauss.diagC = diag(C);
   Gauss.hatn  = Gauss.A*Gauss.m;

else

   % this part heavily relies on the appendix of the accompanying note
   % basic idea:
   % - the precision matrix K is of the form A' hatK A + diagK, where both hatK and diagK are diagonal matrices
   % - apply Woodbury's formula to replace inverting an (nfeat x nfeat) matrix by an (nsample x nsample) alternative
   % - projections of the covariance matrix and the mean onto the feature matrix then follow immediately

   scaledA     = bsxfun(@rdivide,Gauss.A,Gauss.diagK');
   W           = Gauss.A*scaledA';
   W           = (W + W')/2;    % make symmetric
   [Q,logdet1] = invert_chol(diag(1./Gauss.hatK) + W);

   Gauss.hatB = zeros(nsamples,1);
   for k=1:nsamples,
      Gauss.hatB(k) = W(k,k) - W(k,:)*Q*W(:,k);
   end

   Gauss.m     = Gauss.h./Gauss.diagK - scaledA'*(Q*(scaledA*Gauss.h));
   Gauss.hatn  = Gauss.A*Gauss.m;

   Gauss.diagC = 1./Gauss.diagK;

%     for i=1:nfeatures,
%        Gauss.diagC(i) = Gauss.diagC(i) - scaledA(:,i)'*Q*scaledA(:,i);
%     end

   % adriana's recipe
   z = scaledA' * Q; for i=1:size(z), Gauss.diagC(i) = Gauss.diagC(i) - z(i,:) * scaledA(:,i); end

   logdet1 = logdet1 + sum(log(Gauss.diagK)) + sum(log(Gauss.hatK));

end

% compute quadratic term (BC)

qterm = sum(Gauss.m .* Gauss.diagK .* Gauss.m); % = m' * diagK * m
qterm = qterm + sum(Gauss.hatn .* Gauss.hatK .* Gauss.hatn);

logp1 = 0.5*(qterm - logdet1);

%% (2) auxiliary variables; i.e., wrt scale mixture representation of
% Laplace prior

% this is (by far) the most expensive step when nsamples << nfeatures
% and the precision matrix of the auxiliary variables is non-diagonal

[auxC,logdet2] = invert_chol(Gauss.auxK); % only need diagonal terms
Gauss.auxC     = full(diag(auxC));        % turn into full vector
logp2          = 0.5*( - logdet2);
logp           = logp1 + 2*logp2;

%%%%%%%%%%
%
% take out the old term proxies and add the new termproxies and check whether the resulting Gaussian is still normalizable

function [newGauss,newterms,ok] = try_update(Gauss,fullterms,terms,stepsize)

if nargin < 4,
   stepsize = 1;
end

% take out the old term proxies

newGauss = update_Gauss(Gauss,terms,-1);

% compute the new term proxies as a weighted combi of the old ones and the "full" (stepsize 1) term proxies

newterms = combine_terms(fullterms,terms,stepsize);

% add the new term proxies

newGauss = update_Gauss(newGauss,newterms,1);

[L,check,dummy] = chol(newGauss.auxK,'lower');   % check whether full covariance matrix is ok
                % note that this is bit inefficient, since we redo the Cholesky later when everything is fine

ok = (check == 0 & all(newGauss.hatK > 0) & all(newGauss.diagK > 0));  % perhaps a bit too strong???

%%%%%%%%%%%%
%
% compute the moment form of the current Gauss and all cavity approximations and check whether they are fine

function [Gauss,myGauss,logdet,ok] = try_project(Gauss,terms,fraction)

if nargin < 3,
   fraction = 1;
end

[Gauss,logdet] = canonical_to_moments(Gauss);
[myGauss,ok]   = project_all(Gauss,terms,fraction);

%%%%%%%%%%
%
% if we use a temperature < 1, to get closer to the MAP solution, we have to change the prior and initial term proxies accordingly

function [Gauss,terms] = correct_for_temperature(Gauss,terms,temperature)

% note: choose temperature small to implement MAP-like behavior

Gauss.hatK  = Gauss.hatK/temperature;
Gauss.h     = Gauss.h/temperature;
Gauss.auxK  = Gauss.auxK/temperature;
Gauss.diagK = Gauss.diagK/temperature;

terms.hatK  = terms.hatK/temperature;
terms.hath  = terms.hath/temperature;
terms.diagK = terms.diagK/temperature;
terms.auxK  = terms.auxK/temperature;
terms.h     = terms.h/temperature;

%%%%%%%%%%
%
% invert a positive definite matrix using Cholesky factorization

function [invA,logdet] = invert_chol(A)

if issparse(A)

  if 0 % matlab version; slower but useful in case of mex problems

    [L,dummy,S] = chol(sparse(A),'lower');   % now A = S*(L*L')*S' and (L*L') = S'*A*S

    n           = length(L);
    invdiagL2   = 1./spdiags(L,0).^2;

    invA = A;
    for i=n:-1:1,
      I      = i+find(L(i+1:n,i));
      invA(I,i) = -(invA(I,I)*L(I,i))/L(i,i);
      invA(i,I) = invA(I,i)';
      invA(i,i) = invdiagL2(i) - (invA(i,I)*L(I,i))/L(i,i);
    end

    invA = S*invA*S';

  else

    [L,dummy,S] = chol(sparse(A),'lower');   % now A = S*(L*L')*S' and (L*L') = S'*A*S

    if dummy
      error('matrix is not p.d.');
    end

    invA = fastinvc(L);
    invA = S*invA*S';
  end

else

  [L,dummy] = chol(A,'lower');

  if dummy
    error('matrix is not p.d.');
  end

  invA = inv(A);

end

if nargout > 1,
  logdet = 2*sum(log(full(diag(L))));
end

%%%%%%%%%%
%
% compute the term proxy when [oldC,oldm] changes to [newC,newm]

function [K,h,logz] = compute_termproxy(newC,newm,oldC,oldm,fraction)

if nargin < 5,
   fraction = 1;
end

K    = (1./newC - 1./oldC)/fraction;
h    = (newm./newC - oldm./oldC)/fraction;
logz = oldm.^2./oldC/2 - newm.^2./newC/2 + log(full(oldC./newC))/2 ;

%%%%%%%%%%
%
% Sherman-Morrison formula to compute the change from [oldC,oldm] to [newC,newm] when we add [K,h] to the corresponding canonical parameters

function [newC,newm] = rank_one_update(oldC,oldm,K,h)

dummy         = K.*oldC;
oneminusdelta = 1./(1+dummy);
newC          = oneminusdelta.*oldC;

if nargout > 1,
   newm = oneminusdelta.*(oldm + h.*oldC);
end

%%%%%%%%%%%
%
% general procedure for a weighted combi of the fields of two structures

function terms = combine_terms(terms1,terms2,stepsize)

names1 = fieldnames(terms1);
names2 = fieldnames(terms2);
names  = intersect(names1,names2);

terms = struct;
for i=1:length(names)
  terms.(names{i}) = stepsize*terms1.(names{i}) + (1-stepsize)*terms2.(names{i});
end

%%%%%%%%%%%
%
% updates the Gaussian representation with new term proxies

function Gauss = update_Gauss(Gauss,terms,const)

if nargin < 3,
   const = 1;
end

Gauss.h     = Gauss.h     + const*Gauss.A'*terms.hath + const*terms.h;
Gauss.hatK  = Gauss.hatK  + const*terms.hatK;
Gauss.diagK = Gauss.diagK + const*terms.diagK;

% get diagonal elements
diagidx             = 1:(size(Gauss.auxK,1)+1):numel(Gauss.auxK);
Gauss.auxK(diagidx) = Gauss.auxK(diagidx) + const*terms.auxK';


classdef TREE3 < PROBLEM
% <multi> <real> <large/none> <constrained> <expensive/none>
% The time-varying ratio error estimation problem
% T --- 1000 --- Length of data (related to the number of variables)

%------------------------------- Reference --------------------------------
% C. He, R. Cheng, C. Zhang, Y. Tian, Q. Chen, and X. Yao, Evolutionary
% large-scale multiobjective optimization for ratio error estimation of
% voltage transformers, IEEE Transactions on Evolutionary Computation,
% 2020, 24(5): 868-881.
%------------------------------- Copyright --------------------------------
% Copyright (c) 2023 BIMK Group. You are free to use the PlatEMO for
% research purposes. All publications which use this platform or any code
% in the platform should acknowledge the use of "PlatEMO" and reference "Ye
% Tian, Ran Cheng, Xingyi Zhang, and Yaochu Jin, PlatEMO: A MATLAB platform
% for evolutionary multi-objective optimization [educational forum], IEEE
% Computational Intelligence Magazine, 2017, 12(4): 73-87".
%--------------------------------------------------------------------------

    properties(Access = private)
        Data;   % Dataset
        Mean;   % Mean values of the dataset
    end
    methods
        %% Default settings of the problem
        function Setting(obj)
            % Load data
            T = obj.ParameterSet(1000);
            CallStack = dbstack('-completenames');
            load(fullfile(fileparts(CallStack(1).file),'Dataset_TREE.mat'),'Dataset');
            obj.Data = Dataset.TREE3(1:min(max(T,10),end),:);
            % Set the numbers of objectives and decision variables
            K = 6;
            obj.M = 2;
            obj.D = T*K;
            % Set the upper and lower boundaries
            Lower    = zeros(T,K);
            Upper    = zeros(T,K);
            obj.Mean = zeros(T,K);
            for i = 1 : 3
                Lower(:,i)      = min(obj.Data(:,i:3:end*0.5),[],2)/1.01;
                Upper(:,i)      = max(obj.Data(:,i:3:end*0.5),[],2)/0.99;
                Lower(:,i+3)    = min(obj.Data(:,end*0.5+i:3:end),[],2)/1.01;
                Upper(:,i+3)    = max(obj.Data(:,end*0.5+i:3:end),[],2)/0.99;
                obj.Mean(:,i)   = mean(obj.Data(:,i:3:end*0.5),2);
                obj.Mean(:,i+3) = mean(obj.Data(:,end*0.5+i:3:end),2);
            end
            % The decision variables are the offset of the mean values of
            % the dataset
            obj.lower    = reshape(Lower-obj.Mean,1,[]);
            obj.upper    = reshape(Upper-obj.Mean,1,[]);
            obj.encoding = ones(1,obj.D);
        end
        %% Generate initial solutions
        function Population = Initialization(obj,N)
            if nargin < 2; N = obj.N; end
            PopDec = obj.Mean(:)'.*(rand(N,obj.D)*0.008-0.004);
            PopDec = min(max(PopDec,repmat(obj.lower,N,1)),repmat(obj.upper,N,1));
            Population = obj.Evaluation(PopDec);
        end
        %% Calculate objective values
        function PopObj = CalObj(obj,PopDec)
            N      = size(PopDec,1);
            KP     = size(obj.Data,2);
            [T,K]  = size(obj.Mean);
            PopDec = reshape(PopDec,N,T,K) + repmat(reshape(obj.Mean,1,T,K),N,1,1);
            PopObj = zeros(N,obj.M);
            % First objective
            eA1 = abs(repmat(reshape(obj.Data(:,1:end*0.5),1,T,KP/2),N,1,1)./repmat(PopDec(:,:,1:3),1,1,KP/6)-1);
            eA2 = abs(repmat(reshape(obj.Data(:,end*0.5+1:end),1,T,KP/2),N,1,1)./repmat(PopDec(:,:,4:6),1,1,KP/6)-1);
            eA  = cat(3,eA1,eA2);
            PopObj(:,1) = sum(sum(eA,3),2);
            % Second objective
            Delta = std(eA(:,2:end,:)-eA(:,1:end-1,:),0,2);
            PopObj(:,2) = sum(reshape(Delta,N,[]),2);
        end
        %% Calculate constraint violations
        function PopCon = CalCon(obj,PopDec)
            N      = size(PopDec,1);
            [T,K]  = size(obj.Mean);
            PopDec = reshape(PopDec,N,T,K) + repmat(reshape(obj.Mean,1,T,K),N,1,1);
            PopCon = TREE_CalCon(PopDec,2);
        end
        %% Generate a point for hypervolume calculation
        function R = GetOptimum(obj,~)
            X = zeros(1,obj.D);
            X(1:2:end) = obj.lower(1:2:end);
            X(2:2:end) = obj.upper(2:2:end);
            R = obj.CalObj(X);
        end
    end
end


function preconfun = preconhessiansolve(problem, options)
% Preconditioner based on the inverse Hessian, by solving linear systems.
%
% function preconfun = preconhessiansolve(problem)
% function preconfun = preconhessiansolve(problem, options)
%
% Input:
%
% A Manopt problem structure (already containing the manifold and enough
% information to compute the Hessian of the cost) and an options structure
% (optional, currently ignored). Notice that if the Hessian is not positive
% definite, then its inverse is not positive definite either and this
% preconditioner is not suitable.
%
% If the Hessian cannot be computed on 'problem', a warning is issued. An
% approximation of the Hessian will be used instead, and the present
% preconditioner will attempt to invert that (although it may not be a
% linear operator). If no approximate Hessian is provided either, a generic
% approximation is used. Behavior is unspecified.
%
% Output:
%
% Returns a function handle, encapsulating a generic preconditioner of the
% Hessian based on solving linear systems of the form:
%   Hessian(x)[preconfun(x, xdot)] = xdot,
% where x is the point on the manifold, xdot is the input to the
% preconditioner (a tangent vector) and preconfun(x, xdot) is returned
% (also a tangent vector). The solve may be approximate.
%
% The returned preconfun has this calling pattern:
%
%   function precxdot = preconfun(x, xdot)
%   function precxdot = preconfun(x, xdot, storedb)
%   function precxdot = preconfun(x, xdot, storedb, key)
%
% x is a point on the manifold problem.M, xdot is a tangent vector to that
% manifold at x, storedb is a StoreDB object, and key is the StoreDB key to
% point x.
%
% Usage:
%
% Typically, the user will set problem.M and other fields to define the
% cost, the gradient and the Hessian (typically, problem.cost, problem.grad
% and problem.hess, or problem.egrad and problem.ehess). Then, to use this
% generic purpose Hessian preconditioner:
%
%   problem.precon = preconhessiansolve(problem, options);
%
% Passing that problem structure to the conjugategradients solver
% (which uses preconditioning) configured in steepest descent mode results
% in a type of Riemannian Newton method.
%
% See also: conjugategradients

% This file is part of Manopt: www.manopt.org.
% Original author: Nicolas Boumal, April 9, 2015.
% Contributors:
% Change log:

    % Check availability of the Hessian, or at least of an approximation.
    if ~canGetHessian(problem) && ~canGetApproxHessian(problem)
        % Note: we do not give a warning if an approximate Hessian is
        % explicitly given in the problem description, as in that case the
        % user seems to be aware of the issue.
        warning('manopt:getHessian:approx', ...
               ['No Hessian provided. Using an FD approximation instead.\n' ...
                'To disable this warning: warning(''off'', ''manopt:getHessian:approx'')']);
        problem.approxhess = approxhessianFD(problem);
    end

    % Set local defaults here, and merge with user options, if any.
    localdefaults = struct();
    if ~exist('options', 'var') || isempty(options)
        options = struct();
    end
    options = mergeOptions(localdefaults, options);

    % Build and return the function handle here. This extra construct via
    % funhandle makes it possible to make storedb and key optional.
    preconfun = @funhandle;
    function precxdot = funhandle(x, xdot, storedb, key)
        % Allow omission of the key, and even of storedb.
        if ~exist('key', 'var')
            if ~exist('storedb', 'var')
                storedb = StoreDB();
            end
            key = storedb.getNewKey();
        end
        precxdot = hessiansolvehelper(options, problem, x, xdot, ...
                                      storedb, key);
    end

end

function precxdot = hessiansolvehelper(options, problem, x, xdot, storedb, key)
% This function does the actual work.

    % Exclude the case where xdot is zero
    norm_xdot = problem.M.norm(x, xdot);
    if norm_xdot < eps
        precxdot = problem.M.zerovec(x);
        return;
    end

    % Get a shorthand for the Hessian of the cost on M at x.
    hessian = @(u) getHessian(problem, x, u, storedb, key);

    % Setup an optimization problem on the tangent space to problem.M at x.
    M = problem.M;
    tgtspace = tangentspacefactory(M, x);
    prblm.M = tgtspace;
    prblm.cost = @cost;
    prblm.grad = @grad;
    prblm.hess = @(u, udot) 2*hessian(hessian(udot))/norm_xdot;

    function [f, store] = cost(u, store)
        if ~isfield(store, 'residue')
            Hu = hessian(u);
            store.residue = M.lincomb(x, 1, Hu, -1, xdot);
        end
        f = M.norm(x, store.residue).^2 / norm_xdot;
    end
    function [g, store] = grad(u, store)
        if ~isfield(store, 'residue')
            Hu = hessian(u);
            store.residue = M.lincomb(x, 1, Hu, -1, xdot);
        end
        g = 2 * hessian(store.residue) / norm_xdot;
    end

    % checkgradient(prblm); pause;
    % checkhessian(prblm); pause;

    localdefaults.solver = @trustregions;
    localdefaults.verbosity = 0;
    % Merge local defaults with user options, if any.
    if ~exist('options', 'var') || isempty(options)
        options = struct();
    end
    options = mergeOptions(localdefaults, options);

    % Solve the linear system by solving the optimization problem.
    precxdot = manoptsolve(prblm, M.zerovec(), options);

end


function a = sec(a)
%SEC          Gradient secant sec(a)
%

% written  10/16/98     S.M. Rump
% modified 10/14/00     S.M. Rump  use Tony's trick
% modified 03/22/04     S.M. Rump  improved performance
% modified 04/04/04     S.M. Rump  set round to nearest for safety
%                                    accelaration for sparse input
% modified 04/06/05     S.M. Rump  rounding unchanged
% modified 09/28/08     S.M. Rump  check for rounding to nearest improved
% modified 10/08/08     S.M. Rump  improved sparse multiplication: not using intval data type
% modified 08/26/12     S.M. Rump  global variables removed
%

  e = 1e-30;
  if 1+e==1-e                           % fast check for rounding to nearest
    rndold = 0;
  else
    rndold = getround;
    setround(0)
  end

  N = getappdata(0,'INTLAB_GRADIENT_NUMVAR');

  ax = tan(a.x(:));
  a.x = sec(a.x);
  ax = a.x(:) .* ax;
  if issparse(a.dx)
    sizeax = size(a.dx,1);
    [ia,ja,sa] = find(a.dx);
    % take care of scalar a.x: cures Matlab V6.0 bug
    % a=7; i=[1 1]; x=a(i), b=sparse(a); y=b(i)  yields row vector x but column vector y
    ax = ax(ia);
    if isa(a.x,'intval')
      adx = times(ax(:),sa(:),0);
      if adx.complex
        a.dx = intval( sparse(ia,ja,adx.mid,sizeax,N) , sparse(ia,ja,adx.rad,sizeax,N) , 'midrad' );
       else
        a.dx = intval( sparse(ia,ja,adx.inf,sizeax,N) , sparse(ia,ja,adx.sup,sizeax,N) , 'infsup' );
      end
    else
      a.dx = sparse(ia,ja,ax(:).*sa(:),sizeax,N);
    end
  else
    a.dx = a.dx .* ax(:,ones(1,N));
  end

  if rndold
    setround(rndold)
  end


function [samples, logp, diagn] = hmc_nuts(f, theta0, opt)
%HMC_NUTS No-U-Turn Sampler (NUTS)
%
%  Description
%    [SAMPLES, LOGP, DIAGN] = HMC_NUTS(f, theta0, opt)
%    Implements the No-U-Turn Sampler (NUTS), specifically,
%    algorithm 6 from the NUTS paper (Hoffman & Gelman, 2011). Runs
%    opt.Madapt steps of burn-in, during which it adapts the step
%    size parameter epsilon, then starts generating samples to
%    return.
%
%    f(theta) should be a function that returns the log probability its
%    gradient evaluated at theta. I.e., you should be able to call
%    [logp grad] = f(theta).
%
%    opt.epsilon is a step size parameter.
%    opt.M       is the number of samples to generate.
%    opt.Madapt  is the number of steps of burn-in/how long to run
%                the dual averaging algorithm to fit the step size
%                epsilon. Note that there is no need to provide
%                opt.epsilon if doing adaptation.
%    opt.theta0  is a 1-by-D vector with the desired initial setting
%                of the parameters.
%    opt.delta   should be between 0 and 1, and is a target HMC
%                acceptance probability. Defaults to 0.8 if
%                unspecified.
%
%
%    The returned variable "samples" is an (M+Madapt)-by-D matrix
%    of samples generated by NUTS, including burn-in samples.
%
%    Note that when used from gp_mc, opt.M and opt.Madapt are both 0 or
%    1 (hmc_nuts returns only one sample to gp_mc). Number of epsilon
%    adaptations should be set in hmc options structure hmc_opt.nadapt, in
%    gp_mc(... ,'hmc_opt', hmc_opt).
%
%    The returned structure diagn includes step-size vector
%    epsilon, number of rejected samples and dual averaging
%    parameters so its possible to continue adapting step-size
%    parameter.

%      Copyright (c) 2011, Matthew D. Hoffman
%      Copyright (c) 2012, Ville Tolvanen
%      All rights reserved.
%

% Redistribution and use in source and binary forms, with or
% without modification, are permitted provided that the following
% conditions are met:
%
% Redistributions of source code must retain the above copyright
% notice, this list of conditions and the following disclaimer.
%
% Redistributions in binary form must reproduce the above copyright
% notice, this list of conditions and the following disclaimer in
% the documentation and/or other materials provided with the
% distribution.
%
% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
% CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
% INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
% MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
% DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
% CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
% SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
% LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
% USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
% AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
% LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
% ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
% POSSIBILITY OF SUCH DAMAGE.

global nfevals;
nfevals = 0;

if ~isfield(opt, 'delta')
  delta = 0.8;
else
  delta = opt.delta;
end
if ~isfield(opt, 'M')
  M = 1;
else
  M = opt.M;
end
if ~isfield(opt, 'Madapt')
  Madapt = 0;
else
  Madapt = opt.Madapt;
end

diagn.rej = 0;
assert(size(theta0, 1) == 1);

D = length(theta0);
samples = zeros(M+Madapt, D);

[logp grad] = f(theta0);
samples(1, :) = theta0;

% Parameters to the dual averaging algorithm.
gamma = 0.05;
t0 = 10;
kappa = 0.75;
% Initialize dual averaging algorithm.
epsilonbar = 1;
Hbar = 0;

if isfield(opt, 'epsilon') && ~isempty(opt.epsilon)
  epsilon = opt.epsilon(end);

  % Hbar & epsilonbar are needed when doing adaptation of step-length
  if isfield(opt, 'Hbar') && ~isempty(opt.Hbar)
    Hbar = opt.Hbar;
  end
  if isfield(opt, 'epsilonbar') && ~isempty(opt.epsilonbar)
    epsilonbar=opt.epsilonbar;
  else
    epsilonbar=opt.epsilon;
  end
  mu = log(10*opt.epsilon(1));
else
  % Choose a reasonable first epsilon by a simple heuristic.
  epsilon = find_reasonable_epsilon(theta0, grad, logp, f);
  mu = log(10*epsilon);

  opt.epsilon = epsilon;
  opt.epsilonbar = epsilonbar;
  opt.Hbar = Hbar;
end

for m = 2:M+Madapt+1,
%     m
    % Resample momenta.
    r0 = randn(1, D);
    % Joint log-probability of theta and momenta r.
    joint = logp - 0.5 * (r0 * r0');
    % Resample u ~ uniform([0, exp(joint)]).
    % Equivalent to (log(u) - joint) ~ exponential(1).
    logu = joint - exprnd(1);
    % Initialize tree.
    thetaminus = samples(m-1, :);
    thetaplus = samples(m-1, :);
    rminus = r0;
    rplus = r0;
    gradminus = grad;
    gradplus = grad;
    % Initial height j = 0.
    j = 0;
    % If all else fails, the next sample is the previous sample.
    samples(m, :) = samples(m-1, :);
    % Initially the only valid point is the initial point.
    n = 1;
    rej = 0;
    % Main loop---keep going until the criterion s == 0.
    s = 1;
    while (s == 1)
        % Choose a direction. -1=backwards, 1=forwards.
        v = 2*(rand() < 0.5)-1;
        % Double the size of the tree.
        if (v == -1)
            [thetaminus, rminus, gradminus, tmp, tmp, tmp, thetaprime, gradprime, logpprime, nprime, sprime, alpha, nalpha] = ...
                build_tree(thetaminus, rminus, gradminus, logu, v, j, epsilon, f, joint);
        else
            [tmp, tmp, tmp, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alpha, nalpha] = ...
                build_tree(thetaplus, rplus, gradplus, logu, v, j, epsilon, f, joint);
        end
        % Use Metropolis-Hastings to decide whether or not to move to a
        % point from the half-tree we just generated.
        if ((sprime == 1) && (rand() < nprime/n))
          samples(m, :) = thetaprime;
          logp = logpprime;
          grad = gradprime;
        else
          rej = rej + 1;
        end
        % Update number of valid points we've seen.
        n = n + nprime;
        % Decide if it's time to stop.
        s = sprime && stop_criterion(thetaminus, thetaplus, rminus, rplus);
        % Increment depth.
        j = j + 1;
    end

    % Do adaptation of epsilon if we're still doing burn-in.
    eta = 1 / (length(opt.epsilon) + t0);
    Hbar = (1 - eta) * Hbar + eta * (delta - alpha / nalpha);
    if (m <= Madapt+1)
        epsilon = exp(mu - sqrt(m-1)/gamma * Hbar);
        eta = (length(opt.epsilon))^-kappa;
        epsilonbar = exp((1 - eta) * log(epsilonbar) + eta * log(epsilon));
    else
        epsilon = epsilonbar;
    end
    opt.epsilon(end+1) = epsilon;
    opt.epsilonbar = epsilonbar;
    opt.Hbar = Hbar;
    diagn.rej = diagn.rej + rej;
end

diagn.opt = opt;
end

function [thetaprime, rprime, gradprime, logpprime] = leapfrog(theta, r, grad, epsilon, f)
rprime = r + 0.5 * epsilon * grad;
thetaprime = theta + epsilon * rprime;
[logpprime, gradprime] = f(thetaprime);
rprime = rprime + 0.5 * epsilon * gradprime;
global nfevals;
nfevals = nfevals + 1;
end

function criterion = stop_criterion(thetaminus, thetaplus, rminus, rplus)
thetavec = thetaplus - thetaminus;
criterion = (thetavec * rminus' >= 0) && (thetavec * rplus' >= 0);
end

% The main recursion.
function [thetaminus, rminus, gradminus, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alphaprime, nalphaprime] = ...
                build_tree(theta, r, grad, logu, v, j, epsilon, f, joint0)
if (j == 0)
    % Base case: Take a single leapfrog step in the direction v.
    [thetaprime, rprime, gradprime, logpprime] = leapfrog(theta, r, grad, v*epsilon, f);
    joint = logpprime - 0.5 * (rprime * rprime');
    % Is the new point in the slice?
    nprime = logu < joint;
    % Is the simulation wildly inaccurate?
    sprime = logu - 1000 < joint;
    % Set the return values---minus=plus for all things here, since the
    % "tree" is of depth 0.
    thetaminus = thetaprime;
    thetaplus = thetaprime;
    rminus = rprime;
    rplus = rprime;
    gradminus = gradprime;
    gradplus = gradprime;
    % Compute the acceptance probability.
    alphaprime = min(1, exp(logpprime - 0.5 * (rprime * rprime') - joint0));
    nalphaprime = 1;
else
    % Recursion: Implicitly build the height j-1 left and right subtrees.
    [thetaminus, rminus, gradminus, thetaplus, rplus, gradplus, thetaprime, gradprime, logpprime, nprime, sprime, alphaprime, nalphaprime] = ...
                build_tree(theta, r, grad, logu, v, j-1, epsilon, f, joint0);
    % No need to keep going if the stopping criteria were met in the first
    % subtree.
    if (sprime == 1)
        if (v == -1)
            [thetaminus, rminus, gradminus, tmp, tmp, tmp, thetaprime2, gradprime2, logpprime2, nprime2, sprime2, alphaprime2, nalphaprime2] = ...
                build_tree(thetaminus, rminus, gradminus, logu, v, j-1, epsilon, f, joint0);
        else
            [tmp, tmp, tmp, thetaplus, rplus, gradplus, thetaprime2, gradprime2, logpprime2, nprime2, sprime2, alphaprime2, nalphaprime2] = ...
                build_tree(thetaplus, rplus, gradplus, logu, v, j-1, epsilon, f, joint0);
        end
        % Choose which subtree to propagate a sample up from.
        if (rand() < nprime2 / (nprime + nprime2))
            thetaprime = thetaprime2;
            gradprime = gradprime2;
            logpprime = logpprime2;
        end
        % Update the number of valid points.
        nprime = nprime + nprime2;
        % Update the stopping criterion.
        sprime = sprime && sprime2 && stop_criterion(thetaminus, thetaplus, rminus, rplus);
        % Update the acceptance probability statistics.
        alphaprime = alphaprime + alphaprime2;
        nalphaprime = nalphaprime + nalphaprime2;
    end
end
end

function epsilon = find_reasonable_epsilon(theta0, grad0, logp0, f)
epsilon = 0.1;
r0 = randn(1, length(theta0));
% Figure out what direction we should be moving epsilon.
[tmp, rprime, tmp, logpprime] = leapfrog(theta0, r0, grad0, epsilon, f);
acceptprob = exp(logpprime - logp0 - 0.5 * (rprime * rprime' - r0 * r0'));

% Here we presume that energy function returns NaN, if energy cannot be
% evaluated at the suggested hyperparameters so that we need smalled epsilon
if isnan(acceptprob)
  acceptprob=0;
end
a = 2 * (acceptprob > 0.5) - 1;
% Keep moving epsilon in that direction until acceptprob crosses 0.5.
while (acceptprob^a > 2^(-a))
    epsilon = epsilon * 2^a;
    [tmp, rprime, tmp, logpprime] = leapfrog(theta0, r0, grad0, epsilon, f);
    acceptprob = exp(logpprime - logp0 - 0.5 * (rprime * rprime' - r0 * r0'));
end
end


function [groups, orphans, R, C] = connectedComponents(model, type, figures, files)
% Assuming two reactions are connected if they share metabolites, calculates the connected components
% in the stoichiometric matrix, that is, the sets of reactions that share a set of metabolites
%
% USAGE:
%
%    [groups, orphans, R, C] = connectedComponents(model, type, figures)
%
% INPUT:
%    model:
%
% OPTIONAL INPUTS:
%    type:       {('allComponents'), 'largestComponent'}
%    figures:    Will generate plots of the grouping algorithm as it creates block diagonal
%                groups in from top left to bottom right in `W`.
%    files:      Indicator, whether several files containing indicator
%                matrices are generated.
% OUTPUTS:
%    groups:     a structure array (the number of distinct groups is length(groups)) with fields:
%
%                  * `groups(i).num_els` - number of reactions in group `i`.
%                  * `groups(i).block` - sub-block identity of group `i`.
%                  * `groups(i).elements` - reactions of W that are in group `i`.
%                  * `groups(i).degrees` - degrees of connection for each reaction in group `i`.
%    orphans:    elements of W that were not in any group, becasue they did not meet the constraints.
%    R:          reaction adjacency
%    C:          compound adjacency
%
% All components require:
% Connected Component Analysis on an Undirected Graph by Tristan Ursell
% http://www.mathworks.com/matlabcentral/fileexchange/35442-connected-component-analysis-on-an-undirected-graph.
%
% Largest component requires:
% gaimc : Graph Algorithms In Matlab Code by David Gleich
% http://www.mathworks.com/matlabcentral/fileexchange/24134-gaimc-graph-algorithms-in-matlab-code.
%
%  .. Author:
%        - Ronan Fleming, 2012
%        - Thomas Pfau May 2017, Speedup and addition of files indicator

if ~exist('type','var')
    type='allComponents';
end
if ~exist('figures','var')
    figures=0;
end
if ~exist('files','var')
    files=0;
end

model=findSExRxnInd(model);

%stoichiometric matrix
S=model.S;
%dont include exchange reactions
S(:,~model.SIntRxnBool)=0;

[m,n]=size(S);

%binary form
B=sparse(m,n);
B(S~=0)=1;

%Compound adjacency
C1=B*B';

%number of reactions a species participates in
nReactionsSpeciesParticipatesIn=diag(C1,0);

%take out connections by cofactors
[nReactionsSpeciesParticipatesInSorted,IX] = sort(nReactionsSpeciesParticipatesIn,'descend');
%model.mets(IX(1:80))

% %omit reactions connected by cofactors
% omitMet=false(m,1);
% for i=1:m
%     metAbbr=model.mets{i};
%     if strcmp(metAbbr(1:2),'h[')
%         omitMet(i)=1;
%         continue;
%     end
%     if strcmp(metAbbr(1:3),'k[')
%         omitMet(i)=1;
%         continue;
%     end
%     if length(metAbbr)>3
%         if strcmp(metAbbr(1:3),'pi[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:3),'cl[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:3),'o2[')
%             omitMet(i)=1;
%             continue;
%         end
%     end
%
%     if length(metAbbr)>4
%         if strcmp(metAbbr(1:4),'na1[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'h2o[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'co2[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'atp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'adp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'utp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'gtp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'gdp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'amp[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'nad[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'fad[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'coa[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'ppi[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'nh4[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'ACP[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'thf[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:4),'crn[')
%             omitMet(i)=1;
%             continue;
%         end
%     end
%
%     if length(metAbbr)>5
%         if strcmp(metAbbr(1:5),'nadh[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:5),'fadh[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:5),'nadp[')
%             omitMet(i)=1;
%             continue;
%         end
%     end
%
%     if length(metAbbr)>6
%         if strcmp(metAbbr(1:6),'nadph[')
%             omitMet(i)=1;
%             continue;
%         end
%         if strcmp(metAbbr(1:6),'accoa[')
%             omitMet(i)=1;
%             continue;
%         end
%     end
%
% end
%
% %omit these metabolites
% B(omitMet,:)=0;

%Reaction adjacency
R1=B'*B;

%number of species in a reaction
nMolecularSpeciesInReaction=diag(R1,0);

R=R1;
%Set the diagonal to 0
R(logical(eye(size(R1,1)))) = 0;

if files
    R2=triu(R);
    fid=fopen('reactionAdjacencyOtherThanCofactors.txt','w');
    for j=1:n
        fprintf(fid,'%s\t',model.rxns{j});
    end
    fprintf(fid,'\n');
    for j=1:n
        fprintf(fid,'%s\t',model.rxns{j});
        for k=1:n
            fprintf(fid,'%u\t',full(R2(j,k)));
        end
        fprintf(fid,'\n');
    end
    fclose(fid);
end

C=C1;
%Set the diagonal to 0
C(logical(eye(size(C1,1)))) = 0;

if strcmp(type,'largestComponent')
    if ~exist('largest_component')
        error('Install gamic and add it to your path. (http://www.mathworks.com/matlabcentral/fileexchange/24134-gaimc-graph-algorithms-in-matlab-code)')
    end
    [Acc,p] = largest_component(R);
    degree=sum(Acc);
    groups(1).num_els=nnz(degree);
    groups(1).block='largest';
    groups(1).elements=find(p);
    groups(1).degrees=degree;
    orphans=[];
else
    if figures==1
        [groups,orphans]=graph_analysis(R,'plot',1);
    else
        [groups,orphans]=graph_analysis(R);
    end

end

if files
    fid=fopen('reactionsNotConnectedByAnything.txt','w');
    bool=model.SIntRxnBool;
    bool(groups.elements) = 0;
    for j=1:n
        if bool(j)
            fprintf(fid,'%s\n',model.rxns{j});
        end
    end
    fclose(fid);
end


function hf_out = lhs_operation_gpu(hf, samplesf, reg_filter, sample_weights)

% This is the left-hand-side operation in Conjugate Gradient

% Get sizes
num_features = length(hf);
filter_sz = zeros(num_features,2);
for k = 1:num_features
    filter_sz(k,:) = [size(hf{k},1), size(hf{k},2)];
end
[~, k1] = max(filter_sz(:,1));  % Index for the feature block with the largest spatial size
block_inds = 1:num_features;
block_inds(k1) = [];
output_sz = [size(hf{k1},1), 2*size(hf{k1},2)-1];

% Compute the operation corresponding to the data term in the optimization
% (blockwise matrix multiplications)
%implements: A' diag(sample_weights) A f

% sum over all features and feature blocks
sh = sum(bsxfun(@times, samplesf{k1}, hf{k1}), 3);    % assumes the feature with the highest resolution is first
pad_sz = cell(1,1,num_features);
for k = block_inds
    pad_sz{k} = (output_sz - [size(hf{k},1), 2*size(hf{k},2)-1]) / 2;

    sh(1+pad_sz{k}(1):end-pad_sz{k}(1), 1+pad_sz{k}(2):end,1,:) = ...
        sh(1+pad_sz{k}(1):end-pad_sz{k}(1), 1+pad_sz{k}(2):end,1,:) + sum(bsxfun(@times, samplesf{k}, hf{k}), 3);
end

% weight all the samples and take conjugate
sh = conj(bsxfun(@times,sample_weights,sh));

% multiply with the transpose
hf_out = cell(1,1,num_features);
hf_out{k1} = conj(sum(bsxfun(@times, sh, samplesf{k1}), 4));
for k = block_inds
    hf_out{k} = conj(sum(bsxfun(@times, sh(1+pad_sz{k}(1):end-pad_sz{k}(1), 1+pad_sz{k}(2):end,1,:), samplesf{k}), 4));
end

% compute the operation corresponding to the regularization term (convolve
% each feature dimension with the DFT of w, and the tramsposed operation)
% add the regularization part
% hf_conv = cell(1,1,num_features);
for k = 1:num_features
    reg_pad = min(size(reg_filter{k},2)-1, size(hf{k},2)-1);

    % add part needed for convolution
    hf_conv = cat(2, hf{k}, conj(rot90(hf{k}(:, end-reg_pad:end-1, :), 2)));

    % do first convolution
    hf_conv = convn(hf_conv, reg_filter{k});

    % do final convolution and put toghether result
    hf_out{k} = hf_out{k} + convn(hf_conv(:,1:end-reg_pad,:), reg_filter{k}, 'valid');
end

end


function a = sqrt(a)
%SQRT         Gradient square root sqrt(a)
%

% written  10/16/98     S.M. Rump
% modified 10/14/00     S.M. Rump  use Tony's trick
% modified 04/04/04     S.M. Rump  set round to nearest for safety
%                                    accelaration for sparse input
% modified 04/06/05     S.M. Rump  rounding unchanged
% modified 09/28/08     S.M. Rump  check for rounding to nearest improved
% modified 10/08/08     S.M. Rump  improved sparse multiplication: not using intval data type
% modified 08/26/12     S.M. Rump  global variables removed
%

  e = 1e-30;
  if 1+e==1-e                           % fast check for rounding to nearest
    rndold = 0;
  else
    rndold = getround;
    setround(0)
  end

  N = getappdata(0,'INTLAB_GRADIENT_NUMVAR');

  wng = warning;
  warning off

  a.x = sqrt(a.x);
  % use full(a.x(:)): cures Matlab V6.0 bug
  % a=7; i=[1 1]; x=a(i), b=sparse(a); y=b(i)  yields row vector x but column vector y
  % ax is full anyway
  ax = 1./(2*full(a.x(:)));
  if issparse(a.dx)
    sizeax = size(a.dx,1);
    [ia,ja,sa] = find(a.dx);
    if isa(a.x,'intval')
      adx = times(ax(ia),sa,0);
      if adx.complex
        a.dx = intval( sparse(ia,ja,adx.mid,sizeax,N) , sparse(ia,ja,adx.rad,sizeax,N) , 'midrad' );
      else
        a.dx = intval( sparse(ia,ja,adx.inf,sizeax,N) , sparse(ia,ja,adx.sup,sizeax,N) , 'infsup' );
      end
    else
      a.dx = sparse(ia,ja,ax(ia).*sa,sizeax,N);
    end
  else
    a.dx = a.dx .* ax(:,ones(1,N));
  end

  warning(wng)

  if rndold
    setround(rndold)
  end


function [pkx_final, cluster_final, LogL_final] = LCGMM(X, k, W, options)
% Local Consistent Gaussian Mixture Model (LCGMM)
%
% where
%   X
% Notation:
% X ... (nSmp x mFea)  observed data matrix
%       nSmp  ... number of samples
%       mFea  ... number of features
%
% K     ... number of clusters
% W     ... weight matrix of the affinity graph
%
% options ... Structure holding all settings
%
% You only need to provide the above four inputs.
%
% pkx ... P(z|x)
% R ... covariance matrix
% mu ... mean
%
%
% References:
% [1] Jialu Liu, Deng Cai, Xiaofei He, "Gaussian Mixture Model with
% Local Consistency", AAAI 2010.
% [2] Jialu Liu, "Notes on Local Consistent Gaussian Mixture Model(LCGMM)",
% Online available at http://relau.com/jialuliu/technical_notes/LCGMM.pdf
% [accessed 27-Dec-2011].
%
%   version 2.0 --Dec/2011
%   version 1.0 --Dec/2009
%
%   Written by Jialu Liu (remenberl AT gmail.com)
%              Deng Cai (dengcai AT gmail.com)
%
ZERO_OFFSET = 1e-200;

differror = 1e-7;
if isfield(options,'error')
    differror = options.error;
end

lambda = 0.1;
if isfield(options,'lambda')
    lambda = options.lambda;
end

nRepeat = 5;
if isfield(options,'nRepeat')
    nRepeat = options.nRepeat;
end

maxIter = 100;
if isfield(options,'maxIter')
    maxIter = options.maxIter;
end

minIterOrig = 5;
if isfield(options,'minIter')
    minIterOrig = options.minIter;
end
minIter = minIterOrig-1;

meanFitRatio = 0.1;
if isfield(options,'meanFitRatio')
    meanFitRatio = options.meanFitRatio;
end
meanFitControl = 1;
if isfield(options,'meanFitControl')
    meanFitControl = options.meanFitControl;
end

if ~isfield(options,'InitWay')
    options.InitWay = 'kmeans';
end

show = 0;
if isfield(options,'show')
    show = options.show;
end

debug = 0;
if isfield(options,'debug')
    debug = options.debug;
end

% init mixture
[nSmp mFea] = size(X);

if lambda > 0
    DCol = full(sum(W,2));
    D = spdiags(DCol,0,nSmp,nSmp);
    L = D - W;
end

[cluster,pkx,LogL] = GMM_init(X,k,options);
meanFit = LogL/10;

tryNo = 0;
selectInit = 1;
nIter = 0;
while tryNo < nRepeat
    tryNo = tryNo+1;
    maxErr = 1;
    retry = 0;
    while(maxErr > differror || maxErr==-1)
        % EM iteration
        alertFlag = 0;
        for kidx=1:k
            % compute pi
            cluster(kidx).pb = sum(pkx(:,kidx));
            if cluster(kidx).pb < 1e-20
                retry = 1;
                break;
            end
            % compute Tk
            if lambda > 0
                Tk = 1 - lambda * sum(bsxfun(@minus, pkx(:,kidx), pkx(:,kidx)') .* W, 2) ./ (pkx(:,kidx) + ZERO_OFFSET);
            else
                Tk = ones(nSmp, 1);
            end

            if min(Tk) < 0 && alertFlag == 0 && debug == 1
                alertFlag = 1;
                disp('The covariance matrix might not be positive semidefinate since lambda is too big such that some value of Tik is negative.');
            end

            % compute mean
            cluster(kidx).mu = ((pkx(:,kidx) .* Tk)' * X) / cluster(kidx).pb;

            % compute covariance matrix
            Y1 = X-repmat(cluster(kidx).mu,nSmp,1);
            Y2 = bsxfun(@times,sqrt(pkx(:,kidx) .* Tk), Y1);
            R = (Y2'*Y2) /cluster(kidx).pb;
            %for elemetns in Tk which are negative
            Y3 = bsxfun(@times,sqrt(pkx(:,kidx) .* (-(Tk<0) .* Tk)), Y1);
            R = R - 2 * (Y3'*Y3) /cluster(kidx).pb;

            %
            clear Y2;
            R = max(R,R');

            detR = det(R);
            if detR <= 0
                retry = 1;
                break;
            end

            cluster(kidx).cov = R;
            const = -(mFea*log(2*pi) + log(detR))/2;

            Y2 = R\Y1';
            Y2 = -Y2'/2;
            pkx(:,kidx) = dot(Y1,Y2,2)+const;
            clear Y1 Y2 R;
        end

        if retry
            break;
        end

        llmax=max(pkx,[],2);
        pkx =exp( pkx-repmat(llmax,1,k) );
        pkx = pkx.*repmat([cluster(:).pb],nSmp,1);
        ss = sum(pkx,2);
        llnew = sum(log(ss)+llmax);
        pkx = pkx./(repmat(ss,1,k));

        % compute new likelihood
        if lambda > 0
            llnew = llnew - sum(sum((log(pkx' + ZERO_OFFSET) * L).* pkx')) * lambda / 2;
        end

        LogL = [LogL llnew];
        %
        nIter = nIter + 1;

        meanFit = meanFitRatio*meanFit + (1-meanFitRatio)*llnew;
        maxErr = (llnew-meanFit)/meanFit;
        if show
            if length(LogL) > 1
                disp(['tryNo: ',num2str(tryNo),' Iteration: ',num2str(nIter),' LogL: ',num2str(LogL(end)),' deltaLogL: ',num2str(LogL(end)-LogL(end-1)),' maxErr:',num2str(maxErr)]);
            else
                disp(['tryNo: ',num2str(tryNo),' Iteration: ',num2str(nIter),' LogL: ',num2str(LogL(end)),' maxErr:',num2str(maxErr)]);
            end
        end
        if nRepeat > 1 && selectInit
            maxErr = 1;
        end
        if ~meanFitControl
            maxErr = 1;
        end
        if nIter > minIter
            if selectInit
                maxErr = 0;
            else
                if nIter >= maxIter
                    maxErr = 0;
                end
            end
        end
    end

    if retry && ~(tryNo == nRepeat && nIter >= nIter_final)
        tryNo = tryNo - 1;
        [cluster,pkx,LogL] = GMM_init(X,k,options);
        meanFit = LogL/10;
        nIter = 0;
        continue;
    end

    if tryNo == 1
        pkx_final = pkx;
        cluster_final = cluster;
        LogL_final = LogL;
        nIter_final = nIter;
    else
        if LogL(end) > LogL_final(end)
            pkx_final = pkx;
            cluster_final = cluster;
            LogL_final = LogL;
            nIter_final = nIter;
        end
    end

    if selectInit
        if tryNo < nRepeat
            [cluster,pkx,LogL] = GMM_init(X,k,options);
            meanFit = LogL/10;
            nIter = 0;
        else
            tryNo = tryNo - 1;
            selectInit = 0;
            pkx = pkx_final;
            cluster = cluster_final;
            LogL = LogL_final;
            nIter = nIter_final;
            meanFit = LogL(end)/10;
        end
    end
end

function [cluster,pkx,llnew] = GMM_init(X,k,options)
    ZERO_OFFSET = 1e-200;
    [nSmp, mFea] = size(X);
    if strcmpi(options.InitWay,'kmeans')
        kmeansres = litekmeans(X,k,'maxIter',10);
        residx = unique(kmeansres);
        for kidx=1:k
            smpidx = kmeansres==residx(kidx);
            cluster(kidx).mu = mean(X(smpidx,:),1);
            cluster(kidx).pb = 1/k;
        end
    else
        permutation = randperm(nSmp);
        nSmpClass = floor(nSmp/k);
        for kidx=1:k
            cluster(kidx).mu = mean(X(permutation((kidx-1)*nSmpClass+1:kidx*nSmpClass),:),1);
            cluster(kidx).pb = 1/k;
        end
    end

    R = (nSmp-1)*cov(X)/nSmp;
    R = max(R,R');

    % EM iteration

    pkx=zeros(nSmp,k);
    detR = det(R);
    if detR <= 0
        error('The covariance matrix is not positive definite. Use PCA to reduce the dimensions first!');
    end
    const = -(mFea*log(2*pi) +log(detR))/2;
    for kidx=1:k
        Y1=X-repmat(cluster(kidx).mu,nSmp,1);
        Y2 = R\Y1';
        Y2 = -Y2'/2;
        pkx(:,kidx) = dot(Y1,Y2,2)+const;
        clear Y1 Y2;
    end
    clear R;
    llmax=max(pkx,[],2);
    pkx =exp( pkx-repmat(llmax,1,k) );
    pkx = pkx.*repmat([cluster(:).pb],nSmp,1);
    ss = sum(pkx,2);
    llnew = sum(log(ss)+llmax);
    pkx = pkx./(repmat(ss,1,k));