function [model,parms] = loreta_similarity_gal(data,class_labels,parms,init_model)
% [model,parms] = loreta_similarity(data,class_labels,parms,init_model)
%
% Input:   
% -- data              d X N matrix (d = instance dimension, N = number of 
%                      instances)
% -- class_labels      N X 1 vector
% -- init_model        structure with init_model.A, init_model.B being d X k 
%                      matrices where k is the model rank. 
%                      The model is W = A*B'.
%
% Parameters:
% -- parms.step_size   (default 1)
% -- parms.rseed       random seed (default 1)
% -- parms.num_steps   user defined
% -- parms.inner_loop  size of inner loop for counting losses and reporting 
%                      progress (default 1000)
% -- parms.verbose     report once every parms.inner_loop steps (default 0)
% -- parms.use_psd     use PSD version of Loreta instead of general low-rank. 
%                      Uses only init_model.A. The model is W=A*A'
%
% in case init_model is not given as input:       
% -- parms.init_mode   'rand' for random initialization, 'eye' for  zero-padded 
%                      identity initialization (default 'eye')
% -- parms.loreta_lowdim        rank of model; no default, this must be user specified
%
% Output:
% -- model             structure with model.A, model.B being d X k matrices
%                      with k the model rank. The model is W = A*B', unless 
%                      parms.use_psd = 1 in which case the model is W=A*A'.
%                      model.loss_steps is an indicator vector containing 
%                      all loss events = all model update events
% -- parms

    if exist('init_model','var') 
        [init_model, parms, num_samples] = check_arguments(data,parms, ...
						  init_model);
    else
        [init_model, parms, num_samples] = check_arguments(data,parms);
    end
       
    % Rearrange data
    [unused_values, inds] = sort(class_labels); %#ok
    class_labels = class_labels(inds);
    data = data(:,inds);
    classes = sort(unique(class_labels));
    num_classes = length(classes);
  
    % Translate class labels to serial numbers 1,2,...
    new_class_labels = zeros(size(class_labels));
    for i=1:length(classes)
      new_class_labels(class_labels==classes(i)) = i;
    end
    class_labels = new_class_labels;
    class_sizes = zeros(num_classes,1);
    class_start = zeros(num_classes,1);
    for k=1:num_classes
      class_sizes(k) = sum(class_labels==k);
      class_start(k) = find(class_labels==k, 1, 'first');
    end    
    num_without_class = -class_sizes + sum(class_sizes);

    A = init_model.A;
    if parms.use_psd == 0
        B = init_model.B;
    end    
    clear init_model;
    
    % Initialize pseudo inverses - assumes A and B are of rank k
    Apinv = (A'*A)\A';
    if parms.use_psd == 0
        Bpinv = (B'*B)\B';
    end

    max_warp_samples = parms.max_warp_samples;
    num_steps = parms.num_steps;
    step_size = parms.step_size;
    inner_loop_size = parms.inner_loop;
    outer_loss = zeros(1,floor(num_steps/inner_loop_size));
    inner_loss = zeros(1,inner_loop_size);
    inner_l = 0;
    ctime = clock;
    t0 = cputime;

    i = NaN;
    
    n_negative = 0;
    for i_step = 1:num_steps
       if mod(i_step, 100) == 0, fprintf(' %d ', n_negative); end
       if mod(i_step, 1000) == 0,
	 fprintf('[%d/%d]', ceil(i_step/1000), ceil(num_steps/1000)); 
       end    
       if mod(i_step, 10000) == 0, fprintf('[%s]', datestr(now)); end

        % Sample a query/positive/negative trio
        q_ind = floor(rand*num_samples)+1;
	q = data(:,q_ind);
	q_proj = q'*A;
        q_label = class_labels(q_ind);
        pos_ind = class_start(q_label)+floor(rand* class_sizes(q_label));
	
	if max_warp_samples < 1
	  max_warp_samples = length(num_without_class);
	end

	loss = -1;
	n_negative = 0;
	while n_negative < max_warp_samples && loss <= 0
	  n_negative = n_negative + 1;
	  
	  neg_ind = floor(rand*num_samples)+1;        
	  while class_labels(neg_ind) == q_label
            neg_ind = ceil(rand(1)*num_samples);
	  end
	  delta_p = data(:,pos_ind)-data(:,neg_ind); 
	  delta_proj = B'*delta_p;
	  loss = 1 - q_proj* delta_proj;
	end
	
	if parms.use_psd == 0
	  warped_step_size  = step_size * sum(1./ (1: ...
						   floor(num_without_class(q_label)/ n_negative)));
	  
	  pred = 1-loss;    
	  if pred<1
	    inner_loss(i_step-inner_l*inner_loop_size) = 1;
	    [A, B, Apinv, Bpinv] = online_step_loreta1(A, B, Apinv, ...
						       Bpinv, q, ...
						       delta_p, 1, ...
						       warped_step_size);
	  end
	else
            pred = q_proj * (A'*delta_p);    
            if pred<1
                inner_loss(i_step-inner_l*inner_loop_size) = 1;
                [A, Apinv] = online_step_loreta_psd(A, Apinv, q, ...
						    delta_p, 1, ...
						    warped_step_size);
            end
        end           

        if mod(i_step, inner_loop_size) == 0              
            inner_l = inner_l+1;
            outer_loss(inner_l) = mean(inner_loss);
            inner_loss = zeros(1,inner_loop_size);
            if parms.verbose == 1       
                c = clock;
                fprintf('iteration %d out of %d, time elapsed %g sec \n', i_step,...
                    num_steps,etime(c,ctime));
            end
        end
    end    

    % Set output
    model.A = A;
    if parms.use_psd == 0
        model.B = B;
    end
    model.loss_steps = outer_loss;
    model.cputime = cputime-t0;
end

% ============================================================
function [init_model,parms,num_samples] = ...
    check_arguments(data,parms,init_model)

    [dim,num_samples] = size(data);

    if ~exist('init_model','var')
        if ~isfield(parms,'init_model')
            if ~isfield(parms,'loreta_lowdim')
                error(['error: please input initial model or specify' ...
                'initial model rank']);
            else
                parms.init_model = 'eye';
                fprintf(['no initial model specified, using zero-padded'...
                    'identity matrix \n']);
            end
        end            
        switch parms.init_model	
          case 'rand',
              init_model.A = randn(dim,parms.loreta_lowdim);              
              init_model.B = randn(dim,parms.loreta_lowdim);
          case 'eye' ,
              init_model.A = [eye(parms.loreta_lowdim); ...
                  zeros(dim-parms.loreta_lowdim,parms.loreta_lowdim)];
              init_model.B = init_model.A;
          otherwise
              error(['parms.init_model = %s is an invalid choice.' ...
                  'Please set to ''eye'' or ''rand'''],parms.init_model);
        end
    end
  
    if ~isfield(parms,'step_size')
        parms.step_size = 1;
        fprintf('step size value set to 1 \n')
    end

    [dimA,rankA] = size(init_model.A);
    [dimB,rankB] = size(init_model.B);

    if dimA~=dim 
        error('data dim is %d while model.A dim is %d',dim,dimA);
    elseif dimB~=dim
        error('data dim is %d while model.B dim is %d',dim,dimB);
    end

    if rankA~=rankB
        error('model.A and model.B are of incompatible sizes %d and %d',...
            rankA,rankB);
    end

    if ~isfield(parms,'rseed')
        parms.rseed = 1;
    end

    rand('twister',parms.rseed);
    if ~isfield(parms,'inner_loop')
        parms.inner_loop = 1000;
    end

    if ~isfield(parms,'verbose')
        parms.verbose = 0;
    end

    if ~isfield(parms,'use_psd')
        parms.use_psd = 0;
        fprintf('Using Loreta-1 low-rank model \n');
    else
        if parms.use_psd == 0
            fprintf('Using Loreta-1 low-rank model \n');
        else
            fprintf('Using Loreta PSD model \n');              
        end
    end    

    if ~isfield(parms,'num_steps')       
        error('no number of steps provided, please set parms.num_steps \n');
    end
    
end

% =================================================================================    
function [Z1,Z2,Apinv,Bpinv] = online_step_loreta1(A, B, Apinv, Bpinv, ...
						   x1, x2, y, t)
% [Z1,Z2,Apinv,Bpinv] = online_step_loreta1(A,B,Apinv,Bpinv,x1,x2,y,t)
%
% This function implements Loreta-1 from the paper:
% "Online Learning in the Manifold of Low-rank Matrices", Uri Shalit,
% Daphna Weinshall and Gal Chechik, NIPS 2010%
%
% Do retraction step given the gradient: Z1*Z2' = R_{AB}(t*y*A*B')
% The gradient is represented as y*x1*x2' (a matrix)
% 
% Inputs: 
% A,B          - The low-rank factors of the current model. 
%                A is (n X k), B is (m X k)
% Apinv, Bpinv - The pseudo-inverses of A and B
% x1, x2       - The factors of the rank-1 gradient matrix. x1 is (n X 1), 
%                x2 is (m X 1)  
% t            - The step size
% y            - The sign of the step, either -1 or 1
%
% Outputs:
% Z1, Z2       - The low-rank factors of the model after the retractions step
% Apinv, Bpinv - The pseudo-inverses of Z1, Z2 respectively

    x2temp = t*y*x2;
    temp1 = x2temp'*(Bpinv'*(Apinv*x1)); %this is a scalar if x1, x2 are rank-1

    UUTx1 = A*(Apinv*x1);
    cA =  UUTx1*(-(1/2)+(3/8)*(temp1) ) + x1*(1-(1/2)*temp1) ; %nx1
    dA =  Bpinv*x2temp; %kx1
    Z1 = A + cA*dA';                 

    x2VVT = (x2temp'*B)*Bpinv;
    cB = Apinv*x1; %kx1
    dB = ((-1/2)+(3/8)*temp1)*x2VVT' + (1-(1/2)*temp1)*x2temp; %nx1
    Z2 = B + dB*cB';

    Apinv = rank1_pinv_update(A,Apinv,cA,dA);
    Bpinv = rank1_pinv_update(B,Bpinv,dB,cB);
    
end
   
% ============================================================
function [Y1,Ypinv] = online_step_loreta_psd(Y,Ypinv,G1,G2,y,t)
% [Y1,Ypinv] = online_step_loreta_psd(Y,Ypinv,G1,G2,y,t)
%
% This function implements a PSD version of Loreta
%
%
%
% Do retraction step given the gradient: Z1*Z2' = R_{AB}(t*y*Y*Y')
% The gradient is represented as y*x1*x2' (a matrix)
% 
% Inputs: 
% Y          - The low-rank factor of the current model. Y is (n X k)
% Ypinv      - The pseudo-inverse of Y
% G1, G2     - The factors of the rank-1 gradient matrix. x1, x2 are (n X 1)
% t          - The step size
% y          - The sign of the step, either -1 or 1
%
% Outputs:
% Y         - The low-rank factor of the model after the retraction step
% Ypinv      - The pseudo-inverses of Y

    G1 = t*y*G1;
    g1 = Ypinv*G1;
    g2 = Ypinv*G2;
    norm_g1 = g1'*g1;
    norm_g2 = g2'*g2;
    g11 = Y*g1;
    g22 = Y*g2;
    q = g2'*g1;

    L1 = ( (-0.25+q*0.09375)*g11 + (0.5-0.125*q)*G1 + ... 
        norm_g1*0.09375*g22 - 0.125*norm_g1*G2 );
    L2 = ( (-0.25+q*0.09375)*g22 + (0.5-0.125*q)*G2 + ...
        norm_g2*0.09375*g11 - 0.125*norm_g2*G1 );

    outer1 = L1*g2';
    outer2 = L2*g1';
    Y1 = Y + outer1+outer2;

    Ypinv_intermediate = rank1_pinv_update(Y,Ypinv,L1,g2);
    
    Ypinv = rank1_pinv_update(Y+outer1,Ypinv_intermediate,L2,g1);

end

% ============================================================
function [Apinv_new] = rank1_pinv_update(A,Apinv,c,d)
%
% A is nxk, c is nx1 and d is kx1
% Apinv_new is the pseudoinverse of A+c*d'
% ref.: Carl D Mayer, "Generalized inversion of modified matrices"

    beta = 1+d'*(Apinv*c);
    v = Apinv*c;
    n = Apinv'*d;
    w = c-A*(Apinv*c);
    % m = d-A'*(Apinv'*d);   we deal only with full column rank matrices, 
    % therefore norm_m should be always zero

    norm_w = w'*w;
    %norm_m = m'*m;
    norm_m = 0; %we deal only with full column rank matrices, 
                %therefore norm_m should be always zero
    norm_v = v'*v;
    norm_n = n'*n;

    if abs(beta)>eps && norm_m<eps 
        G = (Apinv*n)/beta;
        G = G*w';
        scalar = (beta/(norm_w*norm_n+beta^2));
        temp = (norm_w/beta)*(Apinv*n);
        temp = scalar*(temp+v);    
        temp = temp*(norm_n*w/beta + n)';
        G = G-temp;
    elseif norm_w>eps &&  abs(beta)<eps && norm_m<eps
       % G = -(Apinv*n)*n'/norm_n-v*w'/norm_w;
       G = -Apinv*(n/norm_n);
       G = G*n';
       temp = v*(w'/norm_w);
       G = G-temp;
    elseif norm_w>eps && norm_m>eps
        G = -v*w'/norm_w-m*n'/norm(m)+beta*m*w'/(norm_w*norm_m);
    elseif norm_w<eps && norm_m>eps && abs(beta)<eps
        G = -v*(v'*Apinv)/norm_v-m*n'/norm(m);
    elseif norm_w<eps && abs(beta)>eps
        G = m*(v'*Apinv)/beta - (beta/(norm_v*norm_m+beta^2))* ...
            (norm_v*m/beta+v)*( (norm_m/beta)*(Apinv'*v)+n)';
    elseif norm_w<eps && norm_m<eps && abs(beta)<eps
        G = -v*(v'*Apinv)/norm_v - n*(n'*Apinv)/norm_n + ...
            ((v'*(Apinv*n))/(norm_v*norm_n))*v*n';
    else
        error(['something is wrong: norm_w=%g, norm_v=%g, norm_m=%g,' ...
            'norm_n=%g, and beta=%g'],norm_w,norm_v,norm_m,norm_n,beta);
    end

    Apinv_new = Apinv+G;
    
end
    