% Paul Schrimpf
% 2011-27-10
% 628 hw4 

clear; close all;
%% load data
data=csvread('hw4-nlsy.csv');
data=data(2:end,:); % first row is all zeros because it was the column
                    % headings
y = log(data(:,9)); % log(average earnings)
col = data(:,8); % whether ever attended college
tuit = data(:,5); % local tuition
lc = data(:,6); % whether there's a local college
afqt = data(:,7); % covariates = mother's education, standardized afqt score

%% OLS and IV estimates
n = numel(y);
X = [col afqt ones(n,1)];
bols = (X'*X) \ X'*y;
ixx = (X'*X) \ eye(size(X,2));
% heteroskedasticity robust variance
xe = X.*((y-X*bols)*ones(1,size(X,2)));
Vols = ixx*xe'*xe*ixx;

% iv
Z = [tuit afqt ones(n,1)];
pz = Z*inv(Z'*Z)*Z';
biv = (X'*pz*X) \ (X'*pz*y);
% heteroskedasticity robust variance
xze = (pz*X).*((y-X*biv)*ones(1,size(X,2)));
ixzx = inv(X'*pz*X);
Viv = ixzx*xze'*xze*ixzx;

%% estimating the propensity score
k = 3; 
bfn = @(z,k) cos(pi*k*z) % basis functions using cosines, need 0<=z<=1

%bfn = @(z,d) 1.*(d==0) + z.*(d==1) + z.^2.*(d==2) + ...
%      (d>=3).*(z>=(d-3)/(k-2)).*(z-(d-3)/(k-2)).^3 % cubic splines

w = [tuit afqt];
% put between zero and one
wmin=min(w,[],1);
w = (w-ones(n,1)*wmin);
wmax=max(w,[],1); 
w = w./(ones(n,1)*wmax);
% construct basis functions up to order k evaluated on observations
[k1 k2] = ndgrid(0:k,0:k)
include=find(k1+k2<=k);
sw = zeros(n,numel(include));
np = 50;
[wg1 wg2] = meshgrid(0:1/(np-1):1); % for plotting
wg1 = wg1(:);
wg2 = wg2(:);
sgrid = zeros(numel(wg1),numel(include));
for j=1:numel(include)
  sw(:,j) = bfn(w(:,1),k1(j)).*bfn(w(:,2),k2(j));
  sgrid(:,j) = bfn(wg1,k1(j)).*bfn(wg2,k2(j));
end
alpha = (sw'*sw)\(sw'*col);
% plot estimates
wg1 = wg1*wmax(1)+wmin(1); % convert back to original support
wg2 = wg2*wmax(2)+wmin(2); % convert back to original support
figure;
mesh(reshape(wg1,np,np),reshape(wg2,np,np),reshape(sgrid*alpha,np,np));
xlabel('Tuition');
ylabel('AFQT');
zlabel('P(college)');
% plot support of propensity score
% note: this assumes that the support of (tuit, afqt) is rectangular. That is
% not exactly true in this data, but it is somewhat close.
figure;
Agrid = reshape(wg2,np,np); % should be constant on rows
Pgrid = reshape(sgrid*alpha,np,np);
plot(Agrid(:,1),min(Pgrid,[],2),Agrid(:,1),max(Pgrid,[],2))
legend('min p|x','max p|x');
xlabel('AFQT');
ylabel('P(College)')

%% Estimating the Marginal treatment effect
phat = sw*alpha;
k = 3; 
bfn = @(z,k) z.^k; % basis functions using polynomials
dbfn = @(z,k) k*z.^(k-1);

%bfn = @(z,d) 1.*(d==0) + z.*(d==1) + z.^2.*(d==2) + ...
%      (d>=3).*(z>=(d-3)/(k-2)).*(z-(d-3)/(k-2)).^3 
%dbfn = @(z,d) 0.*(d==0) + 1.*(d==1) + 2*z.*(d==2) + ...
%      (d>=3).*(z>=(d-3)/(k-2)).*3.*(z-(d-3)/(k-2)).^2 % cubic splines

%bfn = @(z,k) cos(pi*k*z); % basis functions using cosines, need 0<=z<=1
%dbfn = @(z,k) -pi*k*sin(pi*k*z); % derivative of basis function

px = [phat afqt];
minpx = min(px,[],1);
px = px - ones(n,1)*minpx;
maxpx = max(px,[],1);
px = px./(ones(n,1)*maxpx);
% construct basis functions up to order k evaluated on observations
[k1 k2] = ndgrid(0:k,0:k)
include=find(k1+k2<=k);
spx = zeros(n,numel(include));
dspx = spx;
np = 50;
[pxg1 pxg2] = meshgrid(0:1/(np-1):1); % for plotting
pxg1 = pxg1(:);
pxg2 = pxg2(:);
sgp = zeros(numel(pxg1),numel(include));
for j=1:numel(include)
  spx(:,j) = bfn(px(:,1),k1(j)).*bfn(px(:,2),k2(j));
  dspx(:,j) = dbfn(px(:,1),k1(j)).*bfn(px(:,2),k2(j));
  sgp(:,j) = bfn(pxg1,k1(j)).*bfn(pxg2,k2(j));
  dsgp(:,j) = dbfn(pxg1,k1(j)).*bfn(pxg2,k2(j));
end
alpha = (spx'*spx)\(spx'*y);
% plot estimates
pxg1 = pxg1*maxpx(1)+minpx(1); % convert back to original support
pxg2 = pxg2*maxpx(2)+minpx(2); % convert back to original support
P1 = reshape(pxg1,np,np);
P2 = reshape(pxg2,np,np);
S = reshape(sgp*alpha,np,np);
D = reshape(dsgp*alpha,np,np);
for j=1:np % don't report estimates out of the support of phat
  notinsupp = (P1(j,:)<min(Pgrid(j,:)) | P1(j,:)>max(Pgrid(j,:)));
  S(j,notinsupp) = NaN;
  D(j,notinsupp) = NaN;
end
figure;
mesh(P1,P2,S);
xlabel('u');
ylabel('AFQT');
zlabel('E[y|u,x]');
figure;
mesh(P1,P2,D);
xlabel('u');
ylabel('AFQT');
zlabel('MTE');
% marginal treatment effect at average afqt
figure;
plot(P1(10,:),D(10,:))
xlabel('u')
ylabel('MTE')

%% computing confidence bands

%% estimating a PRTE
