function [b, varcov_vars, varcov_hes, varcov_sw, varcov_sp, T_hat] = xtpsse(Y,X,S,spatial,cutoff,out,outtest)
%Estimates Hausman, Hall, Griliches (1984) Poisson Conditional Fixed Effects Model, 
%computes sandwich and spatial SE, and tests for time-invariant spatial
%dependence according to Bertanha and Moser (2014)
%
%
%Marinho Bertanha
%this version: 07/03/2014
%I'd appreciate your feedback: bertanha@stanford.edu




N=size(X,1);



if N~=sum(max(S,[],3))
    
    dropi = (max(S,[],3)==0);
    Y(dropi,:,:)=[];
    X(dropi,:,:)=[];
    S(dropi,:,:)=[];
    spatial(dropi,:)=[];
    if out==1
        fprintf('\n %-1d cross-section unit(s) dropped because never observed \n',sum(dropi));
    end
end

T=size(X,3);


if T~=sum(max(S,[],1))
    dropt = (reshape(max(S,[],1),T,1)==0);
    Y(:,:,dropt)=[];
    X(:,:,dropt)=[];
    S(:,:,dropt)=[];
    if out==1
        fprintf('\n %-1d time period(s) dropped because never observed \n',sum(dropt));
    end
end

if min(sum(S,3))==1
    dropi1=(sum(S,3)==1);
    Y(dropi1,:,:)=[];
    X(dropi1,:,:)=[];
    S(dropi1,:,:)=[];
    spatial(dropi1,:)=[];
    if out==1
        fprintf('\n %-1d cross-section unit(s) dropped because observed in only one time period \n',sum(dropi1));
    end
end

if sum(max(Y.*S,[],3)==0)>0
    dropi2=(max(Y.*S,[],3)==0);
    Y(dropi2,:,:)=[];
    X(dropi2,:,:)=[];
    S(dropi2,:,:)=[];
    spatial(dropi2,:)=[];
    if out==1
        fprintf('\n %-1d cross-section unit(s) dropped because of 0 outcome in all observed time periods \n',sum(dropi2));
    end
end

[N,K,T]=size(X);

if out==1
    if min(min(S))==1
        fprintf('\nBalanced Panel\n')
    else
        fprintf('\nUnbalanced Panel\n')
    end
    fprintf('Cross-section units: %-1d\n',N);
    fprintf('Time Periods: %-1d\n',T);
    fprintf('Total obs: %-1d\n\n\n',sum(sum(S)));

end

if nargin<7 %if outtest is not entered, outtest=1 and K2=K, the default;
    outtest=1;
    K2=K; %number of elements of the score to be used in the test-statistic
    
else %if outtest is entered
    if outtest>0
        K2=outtest;
        outtest=1;
    end

    
end



%*****************************************************************OLS
%We run a OLS of logY on X to get pre-estimates. If Y=0, logY=0 and
%D=1, which is not necessary if minY>0

%pool variables
Yp=reshape(Y,N*T,1);
Sp=reshape(S,N*T,1);
Xp=zeros(N*T,size(X,2));
for t=1:T
    Xp(((t-1)*N+1):t*N,:)=X(:,:,t);
end

%keep (i,t)s that are observed
Yp=Yp(Sp==1);
Xp=Xp(Sp==1,:);


logY=zeros(size(Yp));
logY(Yp>0)=log(Yp(Yp>0));

if min(min(Yp))==0
    D=zeros(size(Yp));
    D(Yp==0)=1;
    Xols=[Xp ones(size(Yp)) D]; %add a dummy for the log issue and a constant term

    
else
    Xols=[Xp ones(size(Yp))]; %add a constant term
    
end




beta_ols=(Xols'*Xols)\(Xols'*logY); %will work as initial guess for MLE



%**************************************************computes PCFE
%Maximizing Likelihood for Poisson Model
%It is strictly concave, so we use Newton-Raphson

if out==1
    disp('Maximizing...');
end

critnum=100;
crit=10^(-6); %minimum convergence criterion
iternum=0;
iter=1000; %max. number of iterations

beta0=beta_ols(1:K);
ll0 = LLPFE(Y,X,S,beta0);
[momg, D_hat]=SPFE(Y,X,S,beta0);
if out==1
    fprintf('Iteration        LogL           eps\n');
end

while (critnum>crit && iternum<iter)
    D_hat_inv=eye(K)/D_hat;
    beta1=beta0+D_hat_inv*(mean(momg)');
    ll1 = LLPFE(Y,X,S,beta1);
    [momg, D_hat]=SPFE(Y,X,S,beta1);
    critnum=max(abs(ll1-ll0));
    iternum=iternum+1;
    if out==1
        fprintf('     %-4d        %-10.4g     %-9.4g\n',[iternum ll1 critnum]);
    end
    beta0=beta1;
    ll0=ll1;
    
end
b=beta1;
D_hat_inv=eye(K)/D_hat;


%***********************fits spatial coord into regular lattice coordinates

dist=zeros(N*(N-1)/2,1);
l=1;
for i=1:(N-1)
    for j=(i+1):N
        dist(l,1)=sum((spatial(i,:)-spatial(j,:)).^2)^(.5);
        l=l+1;
        
    end
end
dist(dist==0)=[];
dmin=0.95*min(dist); %max diagonal size for square cells in the lattice
dstar=(sqrt(2)/2)*dmin; %sides of each square cell


coord=zeros(N,2);
spmin=min(spatial);

for i=1:N
    coord(i,:)=floor((spatial(i,:)-spmin)./(dstar*ones(1,2)))+[1,1];
end




%converts 'cutoff' into lattice units
coord_cutoff=round(cutoff./(dstar*ones(1,2)));


%*****************************************computes covariance matrix C_hat



if out==1
    fprintf('\n\n\nComputing Spatial Covariances...\n\n\n');
end

C_hat = zeros(K,K); %estimator for C

for i=1:N
    for j=1:N
        temp=abs(coord(i,:)-coord(j,:));
        if temp<=coord_cutoff
            if min(coord_cutoff)>0
                pr=prod(1-temp./coord_cutoff);
            else
                pr=1;
            end
            C_hat=C_hat+(1/N)*pr*(momg(i,:)')*momg(j,:);
        end


    end

end


varcov_sp=(1/N)*D_hat_inv*C_hat*D_hat_inv';
se_beta_sp=sqrt(diag(varcov_sp));

vars=zeros(K,K);
for i=1:N;
    vars=vars+(1/N)*(momg(i,:)')*(momg(i,:));
end

varcov_sw=(1/N)*D_hat_inv*vars*D_hat_inv';
se_beta_sw=sqrt(diag(varcov_sw));

varcov_vars=(1/N)*(eye(K)/vars);

varcov_hes=(1/N)*D_hat_inv;




if out==1
    fprintf('Robust Standard Errors (Sandwich formula)\n');
    fprintf('------------------------------------------------------\n');
    fprintf('betas            s.e.          t-stat          p-value\n');
    fprintf('%-9.4g        %-9.4g     %-9.4g       %-2.4f\n\n\n',[b  se_beta_sw b./se_beta_sw 2*(1-normcdf(abs(b./se_beta_sw)))]');
    fprintf('Spatial Standard Errors\n');
    fprintf('------------------------------------------------------\n');
    fprintf('betas            s.e.          t-stat          p-value\n');
    fprintf('%-9.4g        %-9.4g     %-9.4g       %-2.4f\n\n\n',[b  se_beta_sp b./se_beta_sp 2*(1-normcdf(abs(b./se_beta_sp)))]');   
    
end

%********************************************computes test-statistic
if outtest>0
    if out==1
        fprintf('\n\n\nComputing Test-statistic...\n\n\n');
    end

    Kstar=K2*(K2+1)/2;
    
    %number of neighbors
    Nl=prod(2*coord_cutoff+1)-1;
    
    %computing spatial lags
    spmomg=zeros(N,K2);
    
    for i=1:N
        for j=1:N
            temp=abs(coord(i,:)-coord(j,:));
            if (i~=j) && (temp(1)<=coord_cutoff(1)) && (temp(2)<=coord_cutoff(2))
                spmomg(i,:)=spmomg(i,:)+momg(j,1:K2);
            end
        end
        spmomg(i,:)=(1/Nl)*spmomg(i,:);
    end
    
    
    
    Z=zeros(Kstar,K2,N);
    A=zeros(Kstar);

    
    for i=1:N
        l=1;
        for k=1:K2
            Z(l:(l+k-1),k,i)=spmomg(i,1:k)';
            l=l+k;
            
        end
        A=A+Z(:,:,i)*Z(:,:,i)';
    end
    
    Gamma=(-1/N)*A;
    Omega = zeros(Kstar,Kstar);

   
    for i=1:N
        for j=1:N
            temp=abs(coord(i,:)-coord(j,:));
            if temp<=2*coord_cutoff
                if min(coord_cutoff)>0
                    pr=prod(1-temp./(2*coord_cutoff));
                else
                    pr=1;
                end

                
                Omega=Omega+(1/N)*pr*(Z(:,:,i)*momg(i,1:K2)')*(momg(j,1:K2)*Z(:,:,j)');
            end


        end

    end
 
        
    
    %if we cannot invert Gamma or Omega, we use less elements of the score function
    if min(rank(Gamma),rank(Omega))<Kstar
        if out==1
            fprintf('\nCannot invert Gamma or Omega matrix using all elements of the score \n');   
        end
        
        while min(rank(Gamma),rank(Omega))<Kstar 

            K2=min(K2-1,floor((sqrt(1+8*N)-1)/2));

            if K2==0
                error('Gamma or Omega is zero when using just one element of the score. Try to increase the cutoff.');
            end
            Kstar=K2*(K2+1)/2;
            Z=zeros(Kstar,K2,N);
            A=zeros(Kstar);

            for i=1:N
            l=1;
                for k=1:K2
                    Z(l:(l+k-1),k,i)=spmomg(i,1:k)';
                    l=l+k;

                end
            A=A+Z(:,:,i)*Z(:,:,i)';
            end

            Gamma=(-1/N)*A;
            Omega = zeros(Kstar,Kstar);
            for i=1:N
                for j=1:N
                    temp=abs(coord(i,:)-coord(j,:));
                    if temp<=2*coord_cutoff
                        if min(coord_cutoff)>0
                            pr=prod(1-temp./(2*coord_cutoff));
                        else
                            pr=1;
                        end
                        Omega=Omega+(1/N)*pr*(Z(:,:,i)*momg(i,1:K2)')*(momg(j,1:K2)*Z(:,:,j)');
                    end
                end

            end
        end
    end
    W=(eye(Kstar)/Gamma)*Omega*(eye(Kstar)/Gamma);

    if (K2<K)&&(out==1)        
        fprintf('\nTest-statistic computed using %-1d out of %-1d element(s) of the score \n\n',[K2,K]);
    end
    

    
    B=zeros(Kstar,1);
    for i=1:N
        B=B+Z(:,:,i)*momg(i,1:K2)';
    end
    
    Theta=A\B;
    
    T_hat=N*((Theta')/W)*Theta;
    
    
    if out==1
        fprintf('Null hypothesis of time-invariant spatial dependence\n');
        fprintf('------------------------------------------------------\n');
        fprintf('test-statistic       p-value (Chi2 - %-1d df) \n',Kstar);
        fprintf('%-9.4g            %-2.4f\n\n\n',[T_hat (1-chi2cdf(T_hat,Kstar))]');
    end

end

end

%************************************************************************
function [ out ] = LLPFE(Y,X,S,b)
%Computes Log-Likelihood function for the PCFE evaluated at 'b'
%for balanced and unbalanced panels

%Marinho Bertanha
%this version: 02/04/2014 
%previous versions: 01/31/2014, Sept,2011
%I'd appreciate your feedback: bertanha@stanford.edu


%Y Nx1xT
%X NxKxT
%S Nx1xT
%b Kx1


N=size(X,1);
%inverting (i,t) orders: easier for the loops below
Y2=permute(Y,[3 2 1]); %Y2 is Tx1xN
X2=permute(X,[3 2 1]); %X2 is TxKxN
S2=permute(S,[3 2 1]); %S2 is Tx1xN

out=0;

if min(min(S2))==1 %in case of a balanced panel
    for i=1:N

        out=out+sum(Y2(:,1,i).*log(exp(X2(:,:,i)*b)/sum(exp(X2(:,:,i)*b))));

    end
else    %in case of an unbalanced panel
    for i=1:N
        Y3=Y2(S2(:,1,i)==1,1,i);
        X3=X2(S2(:,1,i)==1,:,i);
        out=out+sum(Y3.*log(exp(X3*b)/sum(exp(X3*b))));

    end    
    
end

end



%************************************************************************
function [ sco,D_hat] = SPFE(Y,X,S,b)
%Computes the Score Function out of the Log-Likelihood function for the Poisson Conditional Fixed Effect Model
%the output is NxK matrix, each line with s(Yi,Xi,b)' (1xK)
% and also the hessian:
%D_hat (KxK) is the sample average of (-1)*hessian evaluated at 'b'

%Marinho Bertanha
%this version: 02/04/2014
%previous versions: 01/31/2014; Jul,2012
%I'd appreciate your feedback: bertanha@stanford.edu




%input
%Y Nx1xT
%X NxKxT
%S Nx1xT
%b Kx1


[N,K,T]=size(X);

%inverting (i,t) orders: easier for the loops below
Y2=permute(Y,[3 2 1]); %Y2 is Tx1xN
X2=permute(X,[3 2 1]); %X2 is TxKxN
S2=permute(S,[3 2 1]); %S2 is Tx1xN




sco=zeros(N,K);
D_hat=zeros(K,K);


if min(min(S2))==1      %in case of a balanced panel

    for i=1:N

        %for a detailed explanation of this refer to notes, 01/25/2012, pages
        %4-5
        num1=kron(exp(X2(:,:,i)*b),X2(:,:,i));
        DbPi=(num1(1:T+1:T^2,:)*sum(exp(X2(:,:,i)*b))-exp(X2(:,:,i)*b)*sum(num1(1:T+1:T^2,:)))/sum(exp(X2(:,:,i)*b))^2;
        Wi=diag(exp(X2(:,:,i)*b)/sum(exp(X2(:,:,i)*b))); 
        Wi=eye(T)/Wi;
        Ui=Y2(:,1,i)-(exp(X2(:,:,i)*b)/sum(exp(X2(:,:,i)*b)))*sum(Y2(:,1,i));
        sco(i,:)=((DbPi')*Wi*Ui)';
        D_hat=D_hat+(1/N)*sum(Y2(:,1,i))*(DbPi')*Wi*DbPi;



    end
    
else      %in case of an unbalanced panel
    for i=1:N

        Y3=Y2(S2(:,1,i)==1,1,i);
        X3=X2(S2(:,1,i)==1,:,i);
        T3=sum(S2(:,1,i)==1);

        num1=kron(exp(X3*b),X3);
        DbPi=(num1(1:T3+1:T3^2,:)*sum(exp(X3*b))-exp(X3*b)*sum(num1(1:T3+1:T3^2,:)))/sum(exp(X3*b))^2;
        Wi=diag(exp(X3*b)/sum(exp(X3*b))); 
        Wi=eye(T3)/Wi;
        Ui=Y3-(exp(X3*b)/sum(exp(X3*b)))*sum(Y3);
        sco(i,:)=((DbPi')*Wi*Ui)';
        D_hat=D_hat+(1/N)*sum(Y3)*(DbPi')*Wi*DbPi;



    end
end


end








