function [w, infos] = sgd2(problem, options)
% Inputs:
%       problem     function (cost/grad/hess)
%       options     options
% Output:
%       w           solution of w
%       infos       information
%

d = problem.dim();
n = problem.samples();

% extract options
if ~isfield(options, 'step_init')
    step_init = 0.1;
else
    step_init = options.step_init;
end
step = step_init;

if ~isfield(options, 'step_alg')
    step_alg = 'fix';
else
    if strcmp(options.step_alg, 'decay')
        step_alg = 'decay';
    elseif strcmp(options.step_alg, 'decay-2')
        step_alg = 'decay-2';
    elseif strcmp(options.step_alg, 'fix')
        step_alg = 'fix';
    else
        step_alg = 'decay';
    end
end

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

if ~isfield(options, 'tol_optgap')
    tol_optgap = 1.0e-12;
else
    tol_optgap = options.tol_optgap;
end

if ~isfield(options, 'batch_size')
    batch_size = 10;
else
    batch_size = options.batch_size;
end
num_of_bachces = ceil(n / batch_size);

if ~isfield(options, 'max_epoch')
    max_epoch = 100;
else
    max_epoch = options.max_epoch;
end

if ~isfield(options, 'w_init')
    w = randn(d,1);
else
    w = options.w_init;
end

if ~isfield(options, 'f_opt')
    f_opt = -Inf;
else
    f_opt = options.f_opt;
end

if ~isfield(options, 'permute_on')
    permute_on = 1;
else
    permute_on = options.permute_on;
end

if ~isfield(options, 'verbose')
    verbose = false;
else
    verbose = options.verbose;
end

if ~isfield(options, 'store_w')
    store_w = false;
else
    store_w = options.store_w;
end


% initialize
iter = 0;
epoch = 0;
grad_calc_count = 0;

% store first infos
clear infos;
infos.iter = epoch;
infos.time = [0];
infos.grad_calc_count = grad_calc_count;
f_val = problem.cost(w);
optgap = f_val - f_opt;
infos.optgap = optgap;
infos.optgap2=[optgap];
infos.cost = f_val;
infos.cost2=f_val;
if store_w
    infos.w = w;
end

% display infos
if verbose > 0
    fprintf('SGD: Epoch = %03d, cost = %.16e, optgap = %.4e\n', epoch, f_val, optgap);
end

% set start time
% start_time = tic();
k=options.k;
% main loop
while (optgap > tol_optgap) && (epoch < max_epoch)
        
    % permute samples
    if permute_on
        perm_idx = randperm(n);
    else
        perm_idx = 1:n;
    end
    
    for j=1:num_of_bachces
        % update step-size
        if strcmp(step_alg, 'decay')
            step = step_init / (1 + step_init * lambda * iter);
        elseif strcmp(step_alg, 'decay-2')
            step = step_init / (1 + epoch);
        end
        
        
        %% compute batch index and calculate gradient
        start_index = (j-1) * batch_size + 1;
        nn=min(j*batch_size,n);%to avoide the bad case for last loop
        indice_j = perm_idx(start_index:nn);
        grad =  problem.grad(w, indice_j);
        
        if batch_size>(n-start_index+1)
            break;
        end
        
        % update w
        w = w - step * grad;
        w = hardthreshold(w, k, d);
        iter = iter + 1;
        
        
        %% store info
        infos.time=[infos.time nn+n*(epoch)];
        
        f_val = problem.cost(w);
        optgap = f_val - f_opt;
        
        infos.optgap = [infos.optgap optgap];
        infos.cost = [infos.cost f_val];
        
        if store_w
            infos.w = [infos.w w];
        end
     if mod(j,100)== 0
        fprintf('SGD: Epoch = %03d, cost = %.24e, optgap = %.4e\n', epoch, f_val, optgap);
     end
        
    end
    infos.optgap2=[infos.optgap2,optgap];
    infos.cost2=[infos.cost2,f_val];
    epoch=epoch+1;
    
    %% measure elapsed time

    %% count gradient evaluations
    grad_calc_count = grad_calc_count + num_of_bachces * batch_size;

    if verbose > 0
        fprintf('SGD: Epoch = %03d, cost = %.16e, optgap = %.4e\n', epoch, f_val, optgap);
    end
    
end

if optgap < tol_optgap
    fprintf('Optimality gap tolerance reached: tol_optgap = %g\n', tol_optgap);
elseif epoch == max_epoch
    fprintf('Max epoch reached: max_epochr = %g\n', max_epoch);
end

end
