function N = samplesize4(maxn, delta, conf)
% 
% DESCRIPTION:
% This formula calculates the sample size for sizing a
% SMART trial to select the strategy with the highest
% mean outcome with some specified probability.
% We assume a SMART design with two decision points:
%   1) two initial treatments (A1 in {0,1})
%   2) two treatments (A2 in {0,1}) for
%      non-responders (R = 1) and maintenance treatment
%      for responders (R = 0).
% We make the following assumptions:
% - The marginal variances of the final outcome given the
%   strategy are all equal.  This means that,
%   Var[Y|A1=a1, A2=a2] is the same for all (a1, a2) in 
%   {(1,1), (1,0), (0,1), (0,0)}
% - The final outcome Y is normally distributed.
% - The sample sizes will be large enough so that the
%   estimator for the mean is approximately normally 
%   distributed.
% - The correlation between
%   * the final outcome Y given treatment strategy (1,1)
%     and Y given treatment strategy (1,0) is the same as
%   * the correlation between Y given treatment strategy
%     (0,1) and Y given treatment strategy (0,0);
%   we denote this identical correlation by r. 
%  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  
% INPUT
% maxn: the maximum sample size the user has available, 
%       might be constrained by cost
% delta: standardized effect size the user desires to be
%        able to detect
% conf: the probability that the strategy estimated to have
%       the largest mean does indeed have the largest mean
% 
% OUTPUT
% N: sample size
% 
% EXAMPLE
% samplesize4(3000, .2, .90);
% 
% 
% 

x0=[1,maxn]';
options = optimset('Display','off','TolFun', 1, 'TolX',.01);
 
% find u (i.e. sample size) s.t. p0 = conf
[x,fval,exitflag] = fzero(@(u) prob(u, delta, conf), x0, options)
 
N=ceil(x) %N is the sample size
 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 
function [p0] = prob(n, delta, conf)

% function used in samplesize4
 
niter=20000;
% we assume the variance is constant across the treatment 
% strategies;
sigma2=1/4;
% s is the variance of the *estimator* of the strategy 
% means:
s=sigma2.*4; 
% only making one treatment strategy have an effect, this 
% would be the hardest to detect:
one1=[delta.*sqrt(sigma2).*ones(niter,1),zeros(niter,3)]; 
 
for i=1:100

 % full range of possibilities for the correlation:
    r=(i-1)./100; 
    % Correlation comes from people the treatment share,
    % directly related to the response rate.
    % Structure of Sigma comes from sharing responders and 
    % treatments:
    Sigma=[s,r*s,0,0;r*s,s,0,0;0,0,s,r*s;0,0,r*s,s]./n;
        try % to make sure the covariance matrix is valid:
           R=chol(Sigma); 
        catch
           i 
           break;
        end;
    % here is the normality assumption:
    R=mvnrnd(one1,Sigma); 
  
    % how often out of 20000 is the first larger than the 
    % remaining three, this indicates probability (somewhat
    % analagous to power for hypothesis test) of finding 
    % the best strategy:
    test=(R(:,1)> max(R(:,2:4),[],2));
    order(i)=mean(test); 
    
end;
 
% take the one with the worst power, use that to find the 
% sample size:
[p0,j]=min(order'); 
p0=p0-conf;