%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% This function maximizes the unbalanced panel model numerically in minit
% tries out of maxit.
%
% Functions included that are called
% function [par,resnorm,exitflag,output, C,BB,SIimmi] = SPRE_MLE_unbal(JQT,
% ET, IT,IN, WN, WpWN, WWN, Sm,So, uo, n,m,N,T, z0, options)
% function stop = outfun(z, optimValues,state)
%
% Input:
% y...     nx1 vector, dependent variable
% X...     nxK matrix, K explanatory variables including the constant
% WN...    NxN sparse spatial weighting matrix
% So...    nxNT selectormatrix of observed units
% Sm...    NT-nxNT selectormatrix of missing units
% N...     number of cross-section units
% T...     number of time periods
% z0...    3x1 vector of starting values for phi, rho1 and rho2
% minit... number of ml-maximizations at random starting values
% maxit... maximial number of ml-maximizations at random starting values
% tol....  precision, stopping creterion for maximization
%
% Output:
% estp...  K+4 vector of estimated paramters with maximal liklihood out of
%          minit maximizations
% Logl...  Log liklihood
% timer... minit x 1 vector with epalsed maximization time
% iterv... minit x 1 vector with number of iterations
% V1md ... KxK VC matrix of estimated parameters beta
% V2md ... 4x4 VC matrix of estimated parameterts  (s2m,s2n,r4ho1, rho2)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Maximum Likelihood maximization function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [estp,Logl,timer,iterv, V1md, V2md] =GSPRE_MLE_unbal(y,X,So,Sm,WN,N,T,z0, minit,maxit,tol)

% Initialization
C=sparse(zeros(N,N));
BB=sparse(zeros(N,N));
n=size(So,1);
sy=So'*y;
sX=So'*X;
JQT=ones(T,T)/T;
ET=sparse(speye(T)-JQT);
z00=z0;
CsX=zeros(N,size(X,2));

% Paramterspace for rho1 and rho2
opt.tol = 1e-3; opt.disp = 0;
lambda = eigs(WN,1,'SR',opt);
rmin = 1/lambda+0.0001;
rmax = 1-0.0001;
clear lambda;

% Likelihood
    function [par,resnorm,exitflag,output, C,BB,SIimmi] = SPRE_MLE_unbal(JQT, ET, IT,IN, WN, WpWN, WWN, Sm,So, uo, n,m,N,T, z0, options)

        function F = SPRE_MLE_F(z)

            %% B'B,det(B)
            BB=sparse((IN-z(3)*WpWN+(z(3)^2)*WWN));
            ldetBB=4*sum(log(diag(chol(IN-z(3)*WN))));

            %% A'A/Tphi+B'B, inv(inv(A'A/Tphi)+inv(B'B))
            C=sparse(chol(...
                sparse( (IN-z(2)*WpWN+(z(2)^2)*WWN)/(T*z(1)) ) +...
                BB...
                ))\IN;
            C=(sparse(-BB*(C*C')*BB))  ;

            %% det((inv(A'A/Tphi)+inv(B'B))
            ldetAB=- 2*sum(log(diag(chol(BB+C))))   ;

            R=sparse(zeros(N*T-n,N*T-n));
            SSuo=0;
            CSm=zeros(N,N);
            Cuo=zeros(N,1);
            um=zeros(1,N*T-n);
            for i=1:T
                Rj=sparse(zeros(N*T-n,N*T-n));
                SS1oj=0;
                umj=zeros(1,N*T-n);
                CSm=sparse(C*(Sm(:,N*(i-1)+1:N*i))'/T);
                Cuo=C*uo(N*(i-1)+1:N*i,1)/T;
                for j=1:T
                    Rj=Rj+Sm(:,N*(j-1)+1:N*j)*CSm;
                    SS1oj=SS1oj+uo(N*(j-1)+1:N*j,1)'*Cuo;
                    umj=umj+uo(N*(j-1)+1:N*j,1)'*CSm;
                end
                clear CSm Cuo;
                R=sparse(R+(Rj+Sm(:,N*(i-1)+1:N*i)*BB*(Sm(:,N*(i-1)+1:N*i))'));
                SSuo=SSuo+ (SS1oj+uo(N*(i-1)+1:N*i,1)'*BB*uo(N*(i-1)+1:N*i,1))/(n);
                um=um+(umj+uo(N*(i-1)+1:N*i,1)'*BB*(Sm(:,N*(i-1)+1:N*i))');
            end
            clear Rj

            R=chol(R);
            ldOMmm=2*sum(log(diag(R)));
            R=R\(eye(N*T-n,N*T-n));
            SIimmi=sparse(R*R');
            clear R

            SSo1=SSuo-um*SIimmi*um'/n;
            ldOM= -ldetAB+(T-1)*ldetBB;

            F=-(-(n/2)*log(pi)-n*log(SSo1)/2-(ldOMmm-ldOM)/2);

        end

        [par,resnorm,exitflag,output] = fmincon(@SPRE_MLE_F,z0,[],[],[],[],[],[],@SPRE_MLE_C,options);

        function [c2, ceq2] = SPRE_MLE_C(z)
            % Nonlinear inequality constraints
            c2 = [-z(1)+0.01; -z(2)+rmin; z(2)-rmax; -z(3)+rmin; z(3)-rmax];
            % Nonlinear equality constraints
            ceq2 = [];
        end
    end


% Stopping ML-iterations
    function stop = outfun(z, optimValues,state)
        stop = false;
        % Check if fval is less than .000001
        if abs(optimValues.fval) > 1000000000
            stop = true;
        end
    end


% JQT, IT,ET
JQT=ones(T,T)/T;
IN=speye(N);
IT=speye(T);

% W'+W, W'W
WpWN=sparse(WN+WN');
WWN=sparse(WN'*WN);

% Initial residuals
beta=inv(X'*X)*X'*y;
uo=So'*(y-X*beta);
clear beta;

% Loop for ML-estimation starting
ii=1;
jj=1;
maxvalue=ones(minit,1)*10000;
estpar=ones(minit,3)*10000;
timer=ones(minit,1);
iterv=ones(minit,1);
runv=ones(minit,1);

while ii<=minit
    jj=jj+1;
    tic
    rtol=10000;
    rtol1=10000;
    m=1;
    iter=0;

    % Starting guess for the remaining parameters
    rand('state',sum(100*clock));
    z0=z00+0.2*(rand(3,1)-0.5*ones(3,1))';   %starting values

    while rtol > tol
        if 10^(-2*(m-1)-1)>10^(-6)
            otol=10^(-2*(m-1)-1);
        else
            otol=tol;
        end

        % Options and tolerance
        options=optimset('Display','off',  'Diagnostics','off', 'MaxFunEvals',10000, ...
            'LargeScale','off', 'MaxIter',200,'OutputFcn', @outfun,...
            'TolFun',otol, 'TolX', otol);


        [theta,resnorm,exitflag,output, C,BB,SIimmi]=SPRE_MLE_unbal(JQT, ET,...
            IT,IN, WN, WpWN, WWN,Sm,So,uo,n,m,N,T, z0, options);

        clear XXT XYT;
        XXT=zeros(2,2);
        XYT=zeros(2,1);
        for k=1:T
            XXT1=zeros(2,2);
            XYT1=zeros(2,1);
            CsX=C*sX(N*(k-1)+1:N*k,:)/T;
            Csy=C*sy(N*(k-1)+1:N*k,1)/T;
            for l=1:T
                XXT1=XXT1+sX(N*(l-1)+1:N*l,:)'*CsX;
                XYT1=XYT1+sX(N*(l-1)+1:N*l,:)'* Csy;
            end
            XXT=XXT + (XXT1+sX(N*(k-1)+1:N*k,:)'*BB*sX(N*(k-1)+1:N*k,:))/n;
            XYT=XYT + (XYT1+sX(N*(k-1)+1:N*k,:)'*BB*sy(N*(k-1)+1:N*k,1))/n;
        end
        clear XXT1 XYT1 CsX Csy;

        clear XXT2 XYT2;
        XXT2=zeros(2,N*T-n);
        XYT2=zeros(1,N*T-n);
        for k=1:T
            XXT22=zeros(2,N*T-n);
            XYT22=zeros(1,N*T-n);
            CSm=C*(Sm(:,N*(k-1)+1:N*k))'/T;
            for l=1:T
                XXT22=XXT22+sX(N*(l-1)+1:N*l,:)'*CSm;
                XYT22=XYT22+sy(N*(l-1)+1:N*l,:)'*CSm;
            end
            XXT2=XXT2 + (XXT22+sX(N*(k-1)+1:N*k,:)'*BB*(Sm(:,N*(k-1)+1:N*k))');
            XYT2=XYT2 + (XYT22+sy(N*(k-1)+1:N*k,:)'*BB*(Sm(:,N*(k-1)+1:N*k))');
        end
        clear CSm

        uo=So'*(y-X*inv(XXT-XXT2*SIimmi*XXT2'/n)*(XYT-XXT2*SIimmi*XYT2'/n));
        rtol2=resnorm;
        rtol=abs((rtol1-rtol2)/rtol1);
        rtol1=rtol2;
        z0=theta;
        iter=iter+output.iterations;
        step=([m,resnorm,rtol,otol]);
        m=m+1;

    end
    if (exitflag >0)
        maxvalue(ii,1)=resnorm;
        estpar(ii,:)=theta;
        timer(ii,1)=toc;
        iterv(ii,1)=iter;
        runv(ii,1)=ii;
        ii=ii+1;
    end
    if jj>maxit, break, end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

format short
runs=[runv,timer,maxvalue,estpar];
[Logl,zmin]=min(maxvalue);
fprintf('***************************************************************\n');
disp('      runs      sec     Logl       phid     rho1d     rho2d ')
disp(runs)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Estimate slope paramters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
bmd=inv(XXT-XXT2*SIimmi*XXT2'/n)*(XYT-XXT2*SIimmi*XYT2'/n);
b1md=bmd(1,1);
b2md=bmd(2,1);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Estimate s2n
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
SSuo=0;
for i=1:T
    SS1oj=0;
    Cuo=C*uo(N*(i-1)+1:N*i,1)/T;
    for j=1:T
        SS1oj=SS1oj+uo(N*(j-1)+1:N*j,1)'*Cuo;
    end
    SSuo=SSuo+ (SS1oj+uo(N*(i-1)+1:N*i,1)'*BB*uo(N*(i-1)+1:N*i,1))/(n);
end
clear Cuo;

um=zeros(1,N*T-n);
for i=1:T
    umj=zeros(1,N*T-n);
    CfullSm=C*full(Sm(:,N*(i-1)+1:N*i))'/T;
    for j=1:T
        umj=umj+uo(N*(j-1)+1:N*j,1)'*CfullSm;
    end
    um=um+(umj+uo(N*(i-1)+1:N*i,1)'*BB*full(Sm(:,N*(i-1)+1:N*i))');
end
clear CfullSm;
s2nmd=SSuo-um*SIimmi*um'/n;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Collect estimated paramters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
s2mmd=estpar(zmin,1)* s2nmd;
phimd=estpar(zmin,1);
rho1md=estpar(zmin,2);
rho2md=estpar(zmin,3);
estp=[b1md b2md s2mmd s2nmd rho1md rho2md Logl n T]';

clear y X sX sy u uo XXT1 XYT1 XXT22 XYT22 XYT XYT2 CXt Cu Cyt um umj IT ET;
clear JQT SSuo um uo SS1oj
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% VC-matrix of estimated  b
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
invXtSXt=inv(XXT-XXT2*SIimmi*XXT2'/n);
V1md=s2nmd*invXtSXt/n;
clear XXT XXT2 invXtSXt;
pack;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% VC-matrix of estimated s2m snm rho1 rho2
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Omega_oo_inverse
%initialize
n=size(So,1);
m=size(Sm,1);
SIoo=sparse(zeros(n,n));
SIoo1=sparse(zeros(n,n));
SIom=sparse(zeros(n,m));
SIom1=sparse(zeros(n,m));
Sop=So';
Smp=Sm';

for t1=1:T
    for t2=1:T
        SIoo=SIoo+So(:,N*(t1-1)+1:N*t1)*C*Sop(N*(t2-1)+1:N*t2,:)/T;
        SIom=SIom+So(:,N*(t1-1)+1:N*t1)*C*Smp(N*(t2-1)+1:N*t2,:)/T;
    end
    SIoo1=SIoo1+So(:,N*(t1-1)+1:N*t1)*BB*Sop(N*(t1-1)+1:N*t1,:);
    SIom1=SIom1+So(:,N*(t1-1)+1:N*t1)*BB*Smp(N*(t1-1)+1:N*t1,:);
end

SIoo=(SIoo+SIoo1-(SIom+SIom1)*SIimmi*(SIom+SIom1)')/s2nmd;
clear SIoo1 SIom SIom1 SIimmi Smp Sm C
pack

% Derivatives for s2m and rho1
A=chol(IN-rho1md*WpWN+(rho1md^2)*WWN)\IN;
S1oo=sparse(zeros(n,n));
S3oo=sparse(zeros(n,n));
for t1=1:T
    for t2=1:T
        S1oo=S1oo+So(:,N*(t1-1)+1:N*t1)*A*A'*Sop(N*(t2-1)+1:N*t2,:);
        S3oo=S3oo+So(:,N*(t1-1)+1:N*t1)*...
            (s2nmd*phimd*A*A'*(WpWN-2*rho1md*WWN)*A*A')*...
            Sop(N*(t2-1)+1:N*t2,:);
    end
end
clear A
S1=sparse(SIoo*S1oo);
S3=sparse(SIoo*S3oo);
clear S3oo S1oo A
pack;

% Derivatives for s2n and rho2
BB=chol(BB)\IN;
S2oo=sparse(zeros(n,n));
S4oo=sparse(zeros(n,n));

for t1=1:T
    S2oo=S2oo+So(:,N*(t1-1)+1:N*t1)*BB*BB'*Sop(N*(t1-1)+1:N*t1,:);
    S4oo=S4oo+So(:,N*(t1-1)+1:N*t1)*...
        (s2nmd*BB*BB'*(WpWN-2*rho2md*WWN)*BB*BB')*Sop(N*(t1-1)+1:N*t1,:);
end

S2=sparse(SIoo*S2oo);
S4=sparse(SIoo*S4oo);

clear WpWN WWN WN BB IN IT JQT So Sm X y Sx Sy uo SIoo S2oo S4oo Sop;
pack

% Putting all togethter: 0.5*trace(invSIoo*Sj*invSIoo*Sk), j,k=1,...4

V2md=zeros(4,4);

V2md(1,1)=0.5*reshape(S1',1,(n)^2)*reshape(S1,1,(n)^2)';
V2md(1,2)=0.5*reshape(S1',1,(n)^2)*reshape(S2,1,(n)^2)';
V2md(1,3)=0.5*reshape(S1',1,(n)^2)*reshape(S3,1,(n)^2)';
V2md(1,4)=0.5*reshape(S1',1,(n)^2)*reshape(S4,1,(n)^2)';

V2md(2,1)=V2md(1,2);
V2md(2,2)=0.5*reshape(S2',1,(n)^2)*reshape(S2,1,(n)^2)';
V2md(2,3)=0.5*reshape(S2',1,(n)^2)*reshape(S3,1,(n)^2)';
V2md(2,4)=0.5*reshape(S2',1,(n)^2)*reshape(S4,1,(n)^2)';

V2md(3,1)=V2md(1,3);
V2md(3,2)=V2md(2,3);
V2md(3,3)=0.5*reshape(S3',1,(n)^2)*reshape(S3,1,(n)^2)';
V2md(3,4)=0.5*reshape(S3',1,(n)^2)*reshape(S4,1,(n)^2)';

V2md(4,1)=V2md(1,4);
V2md(4,2)=V2md(2,4);
V2md(4,3)=V2md(3,4);
V2md(4,4)=0.5*reshape(S4',1,(n)^2)*reshape(S4,1,(n)^2)';

V2md=inv(V2md);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

end