function [beta,Ghat,bbeta] = selectIV_rw(y,treat,W1,Z,boots)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%		Code for estimating the selection-IV estimator in Laffers & Schmidpeter (2020)		%
%		based on the model of Breunig,Mammen & Simoni (2016) and  Breunig et al. (2015) 	%
%       Input:  y - outcome & selection variable subject to missingness (NaN if missing)	%
%               (y is of dim. Nx2)                                                          % 
%               treat - (potentially) endogenous treatment									%
%               W1 - covariates																%
%               boots - number of bootstrap replications									%
%																							%
%		Output:  beta - estimates from IV													%
%				 Ghat - estimated selection probability										%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



%% Set up data
d = (isnan(y(:,2))) + (isinf(y(:,2)));
d = logical(1-d);
hous = y(:,2);
y = y(:,1);

%% Estimate selection
[Ghat,~,~] = selfunction(y,d,treat,W1,Z);
y(logical(1-d)) = NaN;
yo = hous ./ Ghat;
%% Estimate Effects
[beta] = linear(yo,d,treat,W1,Z,Ghat);


if boots > 0  

	disp('')
	disp('Bootstrapping');
	
	yo(isnan(yo)) = 0; % replace with zeros
	bbeta = zeros(boots,size(beta,1));
	
	%% Centered Residuals
	resv = yo - [ones(size(treat,1),1),treat,W1]*beta;
	resv = resv - mean(resv);
    
    parfor l = 1:boots % BS Loop
         
		d = logical(1-isnan(y(:,1)));
		p = binornd(1,.5,size(y,1),1); % Rademacher
		vnu = -1*p + 1*(1-p);
		resb = resv .* vnu;
		bYb = [ones(size(W1,1),1),treat,W1]*beta + resb;          
		[bb] = linear(bYb,d,treat,W1,Z,Ghat);	
		bbeta(l,:) = bb';
    end
    
	bbeta = nanstd(bbeta)';
    bbeta = round(bbeta,2);
    
else
    
	bbeta = NaN;
       
end

%before output add the wages to Ghat so we can get easily hat(S)
Ghat = [round(Ghat,2),y];
beta = round(beta,2);

end



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%% HELPER FUNCTIONS %%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%% Selection
function [Ghat,y,Gsel] = selfunction(y,d,treat,W1,Z)

y(isnan(y)) = 0;
y(isinf(y)) = 0;

% normalize as it seems to give better GMM performance
y = normalize(y);
treat = normalize(treat);
W1 = normalize(W1);
Z = normalize(Z);

% Optimization 
Ve = [ones(size(W1,1),1),y,W1,treat];
Vz = [ones(size(W1,1),1),W1,Z,treat]; 
mmg = @(x)moments(x,d,Ve,Vz);
smmg = @(x)smoments(x,d,Ve,Vz);
beta0 = zeros(size(Ve,2),1);
beta0(2) = norminv(mean(d));

% Get better starting values then do full optimization
options_d = optimset('LargeScale','off','GradObj','off','Hessian','off','HessUpdate','bfgs','FinDiffType','central',...
                         'Display','off','TolFun',1e-6,'TolX',1e-16,'MaxFunEvals',2000000,'MaxIter',20000);
Gsel = fminunc(mmg,beta0,options_d);

options_d = optimset('LargeScale','on','GradObj','on','Hessian','off','DerivativeCheck','off',...
                            'Display','off','TolFun',1e-6,'TolX',1e-16,'MaxFunEvals',35000,'MaxIter',150);
Gsel =fminunc(smmg,Gsel,options_d);

Ghat = logs(Ve*Gsel);
Gsel = [Gsel];

end


%%%%%%%% Normalization to 0-1 Function
function z = normalize(x)
z = zeros(size(x,1),size(x,2));
    
for ii = 1:size(x,2)
	z(:,ii) = (x(:,ii)-min(x(:,ii)))/(max(x(:,ii)-min(x(:,ii))));
end
end


%%%%%%%% Logistic function
function z = logs(x)

z = exp(x)./(1+exp(x));

end


%%%%%%%% GMM Moment Functionfunction
function mmval = moments(theta,d,V,Z)
    
mmh = d./logs(V*theta);
mm  = sum(Z.*repmat((mmh - 1),1,size(Z,2)));
mm(isnan(mm)) = -log(realmin);
mm(isinf(mm)) = -log(realmin);   
mmval = mm*mm';
   
end


%%%%%%%% SGMM Moment Functionfunction
function [mmval,g] = smoments(theta,d,V,Z)
mmval = moments(theta,d,V,Z);

eps = 0.00001;
epsu = .1;
epsl = 1e-5;
g = zeros(size(theta,1),1);
for i=1:length(theta)
        
	tiny = theta(i)*eps;
	tiny = min(tiny,epsu);
	tiny = max(tiny,epsl);
	theta1 = theta;
	theta1(i)= theta(i) - tiny;
	theta2 = theta;
	theta2(i)= theta(i) + tiny;
	g(i)=(moments(theta2,d,V,Z)-moments(theta1,d,V,Z))'/(2*tiny)';
	
end
    
end


%%%%%%%% Linear IV-Estimator
function [beta] = linear(y,d,treat,W1,Z,g)

y(isnan(y)) = 0;
y(isinf(y)) = 0;
dY = d.*y;
g = d./g;
d = logical(ones(size(d,1),1));
dY = dY(d);   
treat = treat(d,:);
W1 = W1(d,:);
Z = Z(d,:);
Z = [ones(size(treat,1),1),Z,W1];
treat = [ones(size(treat,1),1),treat,W1];
    
Vzx = Z*pinv((Z.*g)'*Z)*Z'*treat; %reweighting only Z'Z part JoE p. 275
ivlow = (Vzx'*Vzx);
ivupp = (Vzx'*dY);
beta = ivlow\ivupp;
end




    
    
    
    





