function model = slvcrf_mmlearning(featureSet, gtLabelSet, trainList, category, pars)
% max margin training of sparse latent variable conditional random fields with graph cut
% written by Jimei Yang @UCM, Oct 2011

% model initilization
num_images = length(trainList);
try
    load(pars.m0file);
catch
    tic;fprintf('initializing...');
    X = cell(1,num_images);
    Y = cell(num_images,1);

    for ii = 1 : num_images
        ti = trainList(ii);
        load(gtLabelSet.path{ti});
        if ~isempty(find(gtLabel.gt==1))
            load(featureSet.path{ti});
            X{ii} = feaSet.feaArr;
            Y{ii} = gtLabel.gt;
        end
    end
    X = cell2mat(X);
    Y = cell2mat(Y);
    model = slvcrf_initialize(X, Y, pars);
    save(pars.m0file, 'model');
    fprintf('done in %.2f seconds.\n',toc);
end

% start training slv-crf
% parameter parsing
gamma = pars.gamma;
lambda = pars.lambda;
nu = pars.nu;
num_bases = pars.numBases;
num_iters = pars.maxIters;
rho = pars.rho;
n0 = num_images/10;

bp_pars.maximize = 0;
bp_pars.max_iters = 200;
bp_pars.tol = 1e-3;

% crf iterations
for i = 1 : num_iters
    tic;
    % Take a random permutation of the samples
    ind_rnd = randperm(num_images);
    for batch=1 : num_images
        
        % load the data
        bIdx = trainList(ind_rnd(batch));
        load(featureSet.path{bIdx});
        data = feaSet.feaArr;
        num_rows = length(unique(feaSet.y));
        num_cols = length(unique(feaSet.x));
        
        load(gtLabelSet.path{bIdx});
        label = gtLabel.gt;
        D = model.D;
        w = model.w;
        
        % build the graph
        graph = StarEdge_MakeEdgeNums_Lattice2(num_rows, num_cols);
        [aff, dist] = BuildGraph(data, num_rows, num_cols);
        graph.A = aff; graph.D = dist; 
        graph.rows = num_rows; graph.cols = num_cols;
        
        % sparse coding
%         W = exp(-dist).*aff;
%         A = D'*D; h = D'*data;
%         [S,f] = graph_sparse_coding(A, h, W, lambda, gamma, 1e-2, 100);
        S = L1QP_FeatureSign_Set(data, D, gamma, lambda);
        
        % energy minimization of E(xn,y) - delta(y,yn) by graphcut
        [label_hat, energy] = crf_infer_gc(model, S, graph, label);
%         % maximization of <w,psi(xn,y)> + delta(y,yn) by lbp
%         [label_hat, belief] = crf_infer_lbp(model, S, graph, label, bp_pars);
        
        % compute the slv-crf features
        [psi_hat, phi_hat] = slvcrf_feature(model, S, label_hat, graph);
        [psi, phi] = slvcrf_feature(model, S, label, graph);

        % compute the gradient of crf parameters
        grad_w = (psi_hat - psi) + nu*w;
        grad_S = phi_hat - phi;

        % compute the gradient of dictionary
        % find the active set and compute beta
        B = zeros(num_bases, num_rows*num_cols);
        for j = 1 : num_rows*num_cols
            active_set = find(S(:, j) ~= 0);
            DtD = D(:,active_set)'*D(:,active_set) + gamma*eye(length(active_set));
            DtD_inv = inv(DtD);
            beta = DtD_inv * grad_S(active_set, j);
            B(active_set, j) = beta;
        end
        grad_D = -D*B*S' + (data - D*S)*B'; % dD = -D*B*S' + (X - D*S)*B';
                
        % update the learning rate
        rho_i = min(rho, rho*n0/batch);
        
        % update crf parameters
        w(1:end-2) = w(1:end-2) - rho_i*grad_w(1:end-2);
%         if w(end)<0; w(end) = 0; end
%         if w(end-1)<0; w(end-1) = 0; end
        model.w = w;
        
        D = D - rho_i*grad_D;
        D = D ./ repmat(sum(D.^2).^.5,size(D,1),1);
        model.D = D;
        
        % Show progress in epoch
        fprintf('Iter = %d, Batch = %d, Elapsed Time = %f\n', i, batch, toc);
    end
    save(fullfile('tmp',sprintf('model_%d_%d.mat',category,i)),'model');
end
    

