function [z, INFO] = TOS(prox1,prox2,subgrad,obj,y0,maxit,gamma_0,strategy)
%AdapTOS Adaptive Three Operator Splitting
%   contact: alpy@mit.edu OR gua@mit.edu

switch lower(strategy)
    case 'smooth', flag_strategy = 0;
    case 'nonsmooth', flag_strategy = 1;
    case 'adaptive', flag_strategy = 2;
    otherwise, error('Unknown step-size strategy.');
end

% INFO collects information about convergence of the method
INFO.err        = nan(maxit,1);
INFO.err_av     = nan(maxit,1);
INFO.feas       = nan(maxit,1);
INFO.feas_av    = nan(maxit,1);
INFO.normy      = nan(maxit,1);
INFO.normgrad   = nan(maxit,1);
INFO.stepsize   = nan(maxit,1);

% Initialization
sum_ut2 = 0; % sum of norm(u_t)^2 

y = y0;   % y_0
z_av = 0; % ergodic average of z_t (\bar{z}_t in the paper)
x_av = 0; % ergodic average of x_t (\bar{x}_t in the paper)
% gamma = 0;
if flag_strategy == 0, gamma = gamma_0; end

for t = 1:maxit
    
    % Decide step-size
    if flag_strategy == 1, gamma = gamma_0/sqrt(t);
    elseif flag_strategy == 2, gamma = gamma_0/sqrt(1+sum_ut2); 
    end
    
    % Main steps
    z = prox2(y,gamma);
    g = subgrad(z);
    x = prox1(2*z - y - gamma*g,gamma);
    y = y + x - z;
    
    % keep track of gradient norm sum squared
    norm_g = norm(g,'fro');
    sum_ut2 = sum_ut2 + norm_g^2;
    
    % Ergodic averages
    z_av = z_av*(1-1/t) + (1/t)*z;
    x_av = x_av*(1-1/t) + (1/t)*x;

    obj_val = obj(z);
    INFO.err(t) = obj_val;
    INFO.err_av(t) = obj(z_av);
    INFO.feas(t) = norm(x-z,'fro');
    INFO.feas_av(t) = norm(x_av-z_av,'fro');
    INFO.normy(t) = norm(y,'fro');
    INFO.normgrad(t) = norm_g;
    INFO.stepsize(t) = gamma;
    
    if mod(t,1000) == 0 % or use 2^floor(log2(t)) == t
        fprintf('Iteration: %d | Objective: %f \n', t, obj_val);
    end
    
end

INFO.iteration = t;

end
