function [w, U, V] = fm_train(y, X, y_test, X_test, lambda_w, lambda_U, lambda_V, d, epsilon, do_pcond, sub_rate, verbose)
% Train a factorization machine using the proposed method in the paper below.
%   Wei-Sheng Chin, Bo-Wen Yuan, Meng-Yuan Yang, and Chih-Jen Lin, An Efficient Alternating Newton Method for Learning Factorization Machines, Technical Report, 2016.
% function [w, U, V] = fm_train(y, X, y_test, X_test, lambda_w, lambda_U, lambda_V, d, epsilon, do_pcond, sub_rate, verbose)
% Inputs:
%   y: training labels, an l-dimensional binary vector. Each element should be either +1 or -1.
%   X: training instances. X is an l-by-n matrix if you have l training instances in an n-dimensional feature space.
%   y_test: test labels. An l'-dimensional label vector. Use empty matrix if you don't have a test set.
%   X_test: test instances. An l'-by-n matrix. Use empty matrix if you don't have a test set.
%   lambda_w: the regularization coefficient of linear term.
%   lambda_U, lambda_V: the regularization coefficients of the two interaction matrices.
%   d: dimension of the latent space.
%   epsilon: stopping tolerance in (0,1). Use a larger value if the training time is too long.
%   do_pcond: a flag. Use 1/0 to enable/disable the diagonal preconditioner.
%   sub_rate: sampling rate in (0,1] to select instances for the sub-sampled Hessian matrix.
%   verbose: a flag controling the level of output information. Set 1/0 for more/less information.
% Outputs:
%   w: linear coefficients. An n-dimensional vector.
%   U, V: the interaction (d-by-n) matrices.
    tic;
    max_iter = 1000;
    [l, n] = size(X);
    w = zeros(1,n);
    rand('seed', 0);
    U = 2*(0.1/sqrt(d))*(rand(d,n)-0.5);
    V = 2*(0.1/sqrt(d))*(rand(d,n)-0.5);
    y_tilde = X*w'+0.5*(sum((U*X').*(V*X'),1))';
    [loss, sigma] = calc_loss(y.*y_tilde);
    f = 0.5*(lambda_w*sum(w.*w)+lambda_U*sum(sum(U.*U))+lambda_V*sum(sum(V.*V)))+loss;
    G_norm_0 = 0;
    if (verbose > 0)
        fprintf('iter        time              obj          |grad|        loss   loss_test         |gradw| (#nt,#cg)           |gradU| (#nt,#cg)           |gradV| (#nt,#cg)\n');
    else
        fprintf('iter        time              obj          |grad|        loss   loss_test\n');
    end
    for k = 1:max_iter
        adjusted = false;
        [w, y_tilde, sigma, f, loss, nt_iters_w, G_norm_w, cg_iters_w, adjusted] = update_block(y, X, w, 2*ones(1,l), y_tilde, sigma, f, loss, lambda_w, do_pcond, sub_rate, adjusted);
        [U, y_tilde, sigma, f, loss, nt_iters_U, G_norm_U, cg_iters_U, adjusted] = update_block(y, X, U, V*X', y_tilde, sigma, f, loss, lambda_U, do_pcond, sub_rate, adjusted);
        [V, y_tilde, sigma, f, loss, nt_iters_V, G_norm_V, cg_iters_V, adjusted] = update_block(y, X, V, U*X', y_tilde, sigma, f, loss, lambda_V, do_pcond, sub_rate, adjusted);
        G_norm = norm([G_norm_w, G_norm_U, G_norm_V]);
        if (k == 1)
            G_norm_0 = G_norm;
        end
        if (G_norm <= epsilon*G_norm_0 || ~adjusted)
            break;
        end
        if (~isempty(y_test) && ~isempty(X_test))
            loss_test = calc_loss(y_test.*(X_test*w'+0.5*(sum((U*X_test').*(V*X_test'),1))'));
        else
            loss_test = nan;
        end
        if (verbose > 0)
            fprintf('%4d  %11.3f  %14.6f  %14.6f  %10.6f  %10.6f  %14.6f (%3d,%3d)    %14.6f (%3d,%3d)    %14.6f (%3d,%3d)\n', k, toc, f, G_norm, loss/length(y), loss_test/length(y_test), G_norm_w, nt_iters_w, cg_iters_w, G_norm_U, nt_iters_U, cg_iters_U, G_norm_V, nt_iters_V, cg_iters_V);
        else
            fprintf('%4d  %11.3f  %14.6f  %14.6f  %10.6f  %10.6f\n', k, toc, f, G_norm, loss/length(y), loss_test/length(y_test));
        end
        if (k == max_iter)
            fprintf('Warning: reach max training iteration. Terminate training process.\n');
        end
    end
end

% See Algorithm 3 in the paper.
function [U, y_tilde, sigma, f, loss, nt_iters, G_norm, total_cg_iters, adjusted] = update_block(y, X, U, Q, y_tilde, sigma, f, loss, lambda, do_pcond, sub_rate, adjusted)
    epsilon = 0.8;
    nu = 0.1;
    max_nt_iter = 100;
    min_step_size = 1e-20;
    l = size(X,1);
    G0_norm = 0;
    total_cg_iters = 0;
    nt_iters = 0;
    for k = 1:max_nt_iter
        G = lambda*U+0.5*Q*sparse([1:l], [1:l], -y.*sigma)*X;
        G_norm = sqrt(sum(sum(G.*G)));
        if (k == 1)
            G0_norm = G_norm;
        end
        if (G_norm <= epsilon*G0_norm)
            return;
        end
        nt_iters = k;
        if (k == max_nt_iter)
            fprintf('Warning: reach newton iteration bound before gradient norm is shrinked enough.\n');
        end
        D = sparse([1:l], [1:l], sigma.*(1-sigma));
        [S, cg_iters] = pcg(X, Q, G, D, lambda, do_pcond, sub_rate);
        total_cg_iters = total_cg_iters+cg_iters;
        Delta = 0.5*(sum(Q'.*(X*S'),2));
        US = sum(sum(U.*S)); SS = sum(sum(S.*S)); GS = sum(sum(G.*S));
        theta = 1;
        while (true)
            if (theta < min_step_size)
                fprintf('Warning: step size is too small in line search. Switch to the next block of variables.\n');
                return;
            end
            y_tilde_new = y_tilde+theta*Delta;
            [loss_new, sigma_new] = calc_loss(y.*y_tilde_new);
            f_diff = 0.5*lambda*(2*theta*US+theta*theta*SS)+loss_new-loss;
            if (f_diff <= nu*theta*GS)
                U = U+theta*S;
                loss = loss_new;
                f = f+f_diff;
                y_tilde = y_tilde_new;
                sigma = sigma_new;
                adjusted = true;
                break;
            end
            theta = theta*0.5;
        end
    end
end

% See Algorithm 4 in the paper.
function [S, cg_iters] = pcg(X, Q, G, D, lambda, do_pcond, sub_rate)
    zeta = 0.3;
    cg_max_iter = 100;
    if (sub_rate < 1)
        l = size(X,1);
        whole = randperm(l);
        selected = sort(whole(1:max(1, floor(sub_rate*l))));
        X = X(selected,:);
        Q = Q(:,selected);
        D = D(selected,selected);
    end
    l = size(X,1);
    S = zeros(size(G));
    Minv = ones(size(G));
    if (do_pcond)
        M = sqrt(lambda+(1/sub_rate)*0.25*(Q.*Q)*(D*(X.*X)));
        Minv(M > 0) = 1./M(M > 0);
    end
    R = -Minv.*G;
    D_hat = R;
    G0G0 = sum(sum(R.*R));
    gamma = G0G0;
    cg_iters = 0;
    while (gamma > zeta*zeta*G0G0)
        cg_iters = cg_iters+1;
        Dh = Minv.*D_hat;
        z_hat = 0.5*sum(Q'.*(X*Dh'),2);
        Dh = Minv.*(lambda*Dh+0.5*(1/sub_rate)*Q*sparse([1:l], [1:l], D*z_hat)*X);
        alpha = gamma/sum(sum(D_hat.*Dh));
        S = S+alpha*D_hat;
        R = R-alpha*Dh;
        gamma_new = sum(sum(R.*R));
        beta = gamma_new/gamma;
        D_hat = R+beta*D_hat;
        gamma = gamma_new;
        if (cg_iters >= cg_max_iter)
            fprintf('Warning: reach max CG iteration. CG process is terminated.\n');
            break;
        end
    end
    S = Minv.*S;
end

function [loss, sigma] = calc_loss(yytilde)
    mask_pos = yytilde > 0;
    mask_neg = yytilde <= 0;
    expyy = exp(yytilde.*(-mask_pos+mask_neg));
    sigma = (mask_neg+expyy.*mask_pos)./(1+expyy);
    loss = sum(-yytilde.*mask_neg+log1p(expyy));
end
