function model = mktl_train(K, P, Y, model)
% multi kernel transfer learning training function 
%   -- modified from the OBSCURE code in the DOGMA library


% This code is part of the supplementary material to the ICCV 2011 paper
% "Multiclass Transfer Learning from Unconstrained Priors", L. Jie*, T.
% Tommasi*, B. Caputo (* equal authorship). 

% Copyright (C) 2010-2011, Tatiana Tommasi
%
% 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/>.
%
% Contact the authors: ttommasi [at] idiap.ch , jluo [at] idiap.ch


timerstart = cputime;

n        = length(Y);   % number of training samples
n_kernel = size(K,3);   % number of kernels
n_prior  = size(P,1);   % number of priors, equalivent to n_prior*n_cls kernels

if isfield(model,'lambda')==0
  if isfield(model, 'C')==0  
     model.lambda = 1/numel(Y);
  else
     model.lambda = 1/(numel(Y)*model.C);
  end
end

if isfield(model,'R2')==0
    model.R2 = 2/model.lambda;
elseif model.R2~=inf
    model.R2 = min(2/model.lambda,model.R2);
end

if isfield(model,'step')==0
    model.step = 100*numel(Y);
end

if isfield(model,'n_cla')==0
    model.n_cla = max(Y);
end

if isfield(model,'iter')==0
    model.iter    = 0;
    model.aer     = [];

    model.epoch   = 0;
    model.time    = [];
    model.sum_tau = 0;

    % W - kernels representation part
    weights_k = zeros(n_kernel, 1);
    beta_k    = spalloc(model.n_cla, n, n*model.n_cla);
    % W - prior linear representation part
    weights_p = zeros(model.n_cla, n_prior);
    beta_p    = zeros(model.n_cla, n_prior);
end

if isfield(model, 'test')==1
    model.inititer = model.test(end);
else
    model.inititer = 0;
    model.test     = [];
end

if isfield(model,'p')==0
    model.q = 2*log(n_kernel);
    model.p = 1/(1-1/model.q);
else
    model.q = 1/(1-1/model.p);
end
qR=sqrt(model.R2)*model.q;

if isfield(model,'T')==0
    model.T = 5;
end

sqnorms_k = zeros(n_kernel, 1)+eps;
sqnorms_p = zeros(n_prior*model.n_cla, 1)+eps;
norm_theta = norm(sqrt([ sqnorms_k; sqnorms_p ])+eps,model.q);

for epoch=1:model.T
    model.epoch=model.epoch+1;
    idx_rand=randperm(n);
    
    model.errTot  = 0;
    model.lossTot = 0;

    n_update=0;
    for i=1:n
        model.iter=model.iter+1;

        idxs_subgrad=idx_rand(i);
        
        preds_k = beta_k * double(K(:, idxs_subgrad, :));
        preds_p = beta_p .* repmat(P(:, idxs_subgrad)', model.n_cla, 1); 
        val_f = preds_k(:, :)*weights_k + sum(preds_p.*weights_p, 2);

        yi = Y(idxs_subgrad);
                
        margin_true = val_f(yi);
        val_f(yi)   = -Inf;
        [margin_pred, yhat]=max(val_f);

        model.errTot  = model.errTot+(margin_true<=margin_pred);
        model.lossTot = model.lossTot+max(1-margin_true+margin_pred,0);
        
        if margin_true<1+margin_pred
            Kii = squeeze(double(K(idxs_subgrad,idxs_subgrad, :)));
            Pii = P(:, idxs_subgrad).^2;

            norm_subgrad = norm([ sqrt(2*Kii); P(:, idxs_subgrad); P(:, idxs_subgrad) ], model.q);    
            a=model.iter*model.lambda+model.sum_tau;
            tau=(-a+sqrt(a^2+model.q*(norm_subgrad+model.lambda/model.q*norm_theta)^2/model.R2))/2;
            model.sum_tau=model.sum_tau+tau;
            eta=model.q/(model.lambda*model.iter+model.sum_tau);

            % Gradient descent step
            fact = 1-eta*model.lambda/model.q;

	        % theta kernel part
			beta_k     = beta_k*fact;
            sqnorms_k  = sqnorms_k*(fact^2);        

            sqnorms_k = sqnorms_k+2*eta*(preds_k(yi, :)'-preds_k(yhat, :)')*fact+(2*eta^2*Kii(:));
                        
            beta_k(yi, idxs_subgrad)   = beta_k(yi, idxs_subgrad)+eta;
            beta_k(yhat, idxs_subgrad) = beta_k(yhat, idxs_subgrad)-eta;
            	    
            % theta prior part
            beta_p     = beta_p*fact; 
            sqnorms_p  = sqnorms_p*(fact^2);

            term2 = zeros(n_prior*model.n_cla, 1); 
            term2((yi-1)*n_prior+1:yi*n_prior, 1)     = preds_p(yi, :)';
            term2((yhat-1)*n_prior+1:yhat*n_prior, 1) = -preds_p(yhat, :)';
	    
            term3 = zeros(n_prior*model.n_cla, 1); 
            term3((yi-1)*n_prior+1:yi*n_prior, 1)     = Pii;
            term3((yhat-1)*n_prior+1:yhat*n_prior, 1) = Pii;

            sqnorms_p = sqnorms_p + 2*eta*term2*fact + eta^2*term3; 

            beta_p(yi, :)   = beta_p(yi, :) + eta*P(:, idxs_subgrad)'; 
            beta_p(yhat, :) = beta_p(yhat, :) - eta*P(:, idxs_subgrad)';

            n_update=n_update+1;

            norms      = sqrt([ sqnorms_k; sqnorms_p]);
            norm_theta = norm(norms+eps,model.q);
            weights    = (norms/norm_theta).^(model.q-2)/model.q;
            weights_k  = weights(1:n_kernel);
            weights_p  = reshape(weights(n_kernel+1:end), n_prior, model.n_cla)'; 
        else
            a=model.iter*model.lambda+model.sum_tau;
            tau=(-a+sqrt(a^2+model.q*(model.lambda/model.q*norm_theta)^2/model.R2))/2;
            model.sum_tau=model.sum_tau+tau;
            eta=model.q/(model.lambda*model.iter+model.sum_tau);

            % Gradient descent step
            fact=1-eta*model.lambda/model.q;

            sqnorms_k = sqnorms_k*(fact^2);
            sqnorms_p = sqnorms_p*(fact^2);            
            beta_k    = beta_k*fact;	        
            beta_p    = beta_p*fact;

            norm_theta = norm_theta*fact;
        end
        
        if qR<norm_theta
            % weights does not change by the scaling
            fact = qR/norm_theta;
            sqnorms_k  = sqnorms_k*(fact^2);
            sqnorms_p  = sqnorms_p*(fact^2);
            beta_k    = beta_k*fact;	        
            beta_p    = beta_p*fact;
     
            norm_theta = norm_theta*fact;
        end

        if mod(model.iter+model.inititer,model.step)==0
            model.test(end+1) = model.iter+model.inititer;
            model.time(end+1) = cputime-timerstart;
            if isfield(model,'eachRound')~=0
                model.beta_k=beta_k;
                model.beta_p=beta_p;
                model.sqnorms_k=sqnorms_k;
                model.sqnorms_p=sqnorms_p;
                model.weights_k=weights_k;
                model.weights_p=weights_p;
                if isfield(model, 'outputER')==0
                   model.outputER = [];
                   model.outputER = feval(model.eachRound, model);
                else
                   model.outputER(end+1) = feval(model.eachRound, model);
                end
            end
	        timerstart=cputime;
        end
    end

    if epoch==model.T
        model.test(end+1)=model.iter+model.inititer;  
        model.time(end+1)=cputime-timerstart;
        if isfield(model,'eachRound')~=0
           model.beta_k=beta_k;
           model.beta_p=beta_p;
           model.sqnorms_k=sqnorms_k;
           model.sqnorms_p=sqnorms_p;
           model.weights_k=weights_k;
           model.weights_p=weights_p;
            if isfield(model, 'outputER')==0
                model.outputER = [];
                model.outputER = feval(model.eachRound, model);
            else
                model.outputER(end+1) = feval(model.eachRound, model);
            end
        end
       	timerstart=cputime;
    end
end

model.beta_k=beta_k;
model.beta_p=beta_p;
model.sqnorms_k=sqnorms_k;
model.sqnorms_p=sqnorms_p;
model.weights_k=weights_k;
model.weights_p=weights_p;
