#This program is for moderation analysis using normal distribution based ML;
#Results of LS analysis is also in the default output;
#The details are in the paper: Moderation Analysis Using a Random Coefficient Regression Model 
#by Ke-Hai Yuan, Ying Cheng, Scott Maxwell, 2011.
#Six standard errors are calculated for each estimator, 
#Four are included in the output;
#####################################################################
#  SE based on the sandwich-type covariance matrix for LS analysis  #
#####################################################################
SW_LS=function(hatee,CC_in,C_mat) { 
C_mat=as.matrix(C_mat);
n=nrow(C_mat);
p_c=ncol(C_mat);
B0_mat=matrix(0,p_c,p_c);
B2_mat=matrix(0,p_c,p_c);
B3_mat=matrix(0,p_c,p_c);
B4_mat=matrix(0,p_c,p_c);
for (i in 1:n){
	c_i=C_mat[i,];
	c_it=t(c_i);
	cc_i=c_i%*%c_it;	
	t_ci=c_it%*%CC_in%*%c_i; #for B2_mat;
	vc_i=1/(1-t_ci);
	vc_i2=vc_i*vc_i; #for B3_mat;

	r_ci=min(4,(n*t_ci/p_c)); #for B4_mat;
	vc_ir=vc_i^r_ci;

	delta_i2=hatee[i];
###############################;
	B0_i=delta_i2*cc_i;
	B2_i=c(vc_i)*B0_i;
	B3_i=c(vc_i2)*B0_i; 
	B4_i=c(vc_ir)*B0_i; 
	
	B0_mat=B0_mat+B0_i;	
	B2_mat=B2_mat+B2_i;	
	B3_mat=B3_mat+B3_i;	
	B4_mat=B4_mat+B4_i;
};

Cov_LSSW0=CC_in%*%B0_mat%*%CC_in;
Var_LSSW0=diag(Cov_LSSW0);
SE_LSSW0=sqrt(Var_LSSW0);

SE_LSSW1=sqrt(n/(n-p_c))*SE_LSSW0;

Cov_LSSW2=CC_in%*%B2_mat%*%CC_in;
Var_LSSW2=diag(Cov_LSSW2);
SE_LSSW2=sqrt(Var_LSSW2);

Cov_LSSW3=CC_in%*%B3_mat%*%CC_in;
Var_LSSW3=diag(Cov_LSSW3);
SE_LSSW3=sqrt(Var_LSSW3);

Cov_LSSW4=CC_in%*%B4_mat%*%CC_in;
Var_LSSW4=diag(Cov_LSSW4);
SE_LSSW4=sqrt(Var_LSSW4);
SE_LSSW=list(SE_LSSW0,SE_LSSW1,SE_LSSW2,SE_LSSW3,SE_LSSW4);

return(SE_LSSW);
}
#####################################################################
# SE based on the sandwich-type covariance matrix for NML analysis  #
#####################################################################
SW_ML=function(hate,tau_vec,C_mat,H_mat){
C_mat=as.matrix(C_mat);
C_matt=t(C_mat);
CC_mat=C_matt%*%C_mat;
CC_in=solve(CC_mat);
H_mat=as.matrix(H_mat);
HH_mat=t(H_mat)%*%H_mat;
HH_in=solve(HH_mat);

n=nrow(C_mat);
p_c=ncol(C_mat);
p_h=ncol(H_mat);
p=p_c+p_h;

A_mat11=matrix(0,p_c,p_c);
A_mat12=matrix(0,p_c,p_h);
A_mat22=matrix(0,p_h,p_h);
B0_mat11=matrix(0,p_c,p_c);
B0_mat12=matrix(0,p_c,p_h);
B0_mat22=matrix(0,p_h,p_h);
B2_mat11=matrix(0,p_c,p_c);
B2_mat12=matrix(0,p_c,p_h);
B2_mat22=matrix(0,p_h,p_h);
B3_mat11=matrix(0,p_c,p_c);
B3_mat12=matrix(0,p_c,p_h);
B3_mat22=matrix(0,p_h,p_h);
B4_mat11=matrix(0,p_c,p_c);
B4_mat12=matrix(0,p_c,p_h);
B4_mat22=matrix(0,p_h,p_h);
for (i in 1:n){
	c_i=C_mat[i,];
	c_it=t(c_i);
	cc_i=c_i%*%c_it;
	h_i=H_mat[i,];
	h_it=t(h_i);
	ch_i=c_i%*%h_it;
	hh_i=h_i%*%h_it;

	t_ci=c_it%*%CC_in%*%c_i; #for B2_mat;
	t_hi=h_it%*%HH_in%*%h_i;
	vc_i=1/(1-t_ci);
	vh_i=1/(1-t_hi);
	vch_i=vc_i*vh_i;
	svch_i=sqrt(vch_i);

	vc_i2=vc_i*vc_i; #for B3_mat;
	vh_i2=vh_i*vh_i;

	r_ci=min(4,(n*t_ci/p_c)); #for B4_mat;
	r_hi=min(4,(n*t_hi/p_h));
	vc_ir=vc_i^r_ci;
	vh_ir=vh_i^r_hi;
	svch_ir=sqrt(vc_ir*vh_ir);

	delta_i=hate[i];
	delta_i2=delta_i*delta_i;
	tau_i2=tau_vec[i];
	tau_i4=tau_i2*tau_i2;
	tau_i6=tau_i4*tau_i2;
	tau_i8=tau_i6*tau_i2;
	dt_i2=delta_i2-tau_i2;
###########################################;
	A_i11=cc_i/tau_i2;
	A_i12=delta_i*ch_i/tau_i4;
	A_i22=(2*delta_i2-tau_i2)*hh_i/(2*tau_i6);
	A_mat11=A_mat11+A_i11;
	A_mat12=A_mat12+A_i12;
	A_mat22=A_mat22+A_i22;


	B_i11=delta_i2*cc_i/tau_i4;
	B_i12=delta_i*dt_i2*ch_i/(2*tau_i6);
	B_i22=dt_i2*dt_i2*hh_i/(4*tau_i8);
	B0_mat11=B0_mat11+B_i11;
	B0_mat12=B0_mat12+B_i12;
	B0_mat22=B0_mat22+B_i22;		
	
	B2_i11=c(vc_i)*B_i11; 
	B2_i12=c(svch_i)*B_i12;
	B2_i22=c(vh_i)*B_i22; 
	B2_mat11=B2_mat11+B2_i11;
	B2_mat12=B2_mat12+B2_i12;
	B2_mat22=B2_mat22+B2_i22;		

	B3_i11=c(vc_i2)*B_i11; 
	B3_i12=c(vch_i)*B_i12;
	B3_i22=c(vh_i2)*B_i22;
	B3_mat11=B3_mat11+B3_i11;
	B3_mat12=B3_mat12+B3_i12;
	B3_mat22=B3_mat22+B3_i22;		

	B4_i11=c(vc_ir)*B_i11; 
	B4_i12=c(svch_ir)*B_i12;
	B4_i22=c(vh_ir)*B_i22;
	B4_mat11=B4_mat11+B4_i11;
	B4_mat12=B4_mat12+B4_i12;
	B4_mat22=B4_mat22+B4_i22;		
};
A_mat21=t(A_mat12);
A_mat=rbind(cbind(A_mat11,A_mat12),cbind(A_mat21,A_mat22) );

B0_mat21=t(B0_mat12);
B0_mat=rbind(cbind(B0_mat11,B0_mat12),cbind(B0_mat21,B0_mat22) );

B2_mat21=t(B2_mat12);
B2_mat=rbind(cbind(B2_mat11,B2_mat12),cbind(B2_mat21,B2_mat22) );

B3_mat21=t(B3_mat12);
B3_mat=rbind(cbind(B3_mat11,B3_mat12),cbind(B3_mat21,B3_mat22) );

B4_mat21=t(B4_mat12);
B4_mat=rbind(cbind(B4_mat11,B4_mat12),cbind(B4_mat21,B4_mat22) );


A_in=solve(A_mat);

Cov_SW0=A_in%*%B0_mat%*%A_in;
Var_SW0=diag(Cov_SW0);
SE_SW0=sqrt(Var_SW0);

SE_SW1=sqrt(n/(n-p))*SE_SW0;

Cov_SW2=A_in%*%B2_mat%*%A_in;
Var_SW2=diag(Cov_SW2);
SE_SW2=sqrt(Var_SW2);

Cov_SW3=A_in%*%B3_mat%*%A_in;
Var_SW3=diag(Cov_SW3);
SE_SW3=sqrt(Var_SW3);

Cov_SW4=A_in%*%B4_mat%*%A_in;
Var_SW4=diag(Cov_SW4);
SE_SW4=sqrt(Var_SW4);
SE_MLSW=list(SE_SW0,SE_SW1,SE_SW2,SE_SW3,SE_SW4);

return(SE_MLSW);
}
#####################################################################
#       LS regression with SE from sandwich-type covariances        #
#####################################################################
LSreg=function(C_mat,y) { 
	C_mat=as.matrix(C_mat);
	n=nrow(C_mat);
	p_c=ncol(C_mat);
	C_matt=t(C_mat);
	CC_mat=C_matt%*%C_mat;
	CC_in=solve(CC_mat);
	Cy=C_matt%*%y;
	hgamma=CC_in%*%Cy;
	hy=C_mat%*%hgamma;
	hate=y-hy;
	hatee=hate^2;
	RSS=sum(hatee);
	hsig2=RSS/(n-p_c);
	Cov_LS=hsig2*CC_in;
	SE_LS=sqrt(diag(Cov_LS));
	hsig2_LS=RSS/n;
	BIC_LS=n+n*log(hsig2_LS)+(p_c+1)*log(n);
	cat("==========================================================\n",append = T);
	cat("==========================================================\n");
	cat("               Least Squares               \n",append = T);

	cat("    \n",append = T);

	cat("BIC_LS=", BIC_LS, "\n",append = T);
	cat("    \n",append = T);

	z_LS=hgamma/SE_LS;
	SE_SW=SW_LS(hatee,CC_in,C_mat);
	SE_SW0=SE_SW[[1]];
	SE_SW3=SE_SW[[4]];
	SE_SW4=SE_SW[[5]];	
	z_SW0=hgamma/SE_SW0;
  z_SW3=hgamma/SE_SW3;
	z_SW4=hgamma/SE_SW4;
	
	tt_LS=cbind(hgamma,SE_LS,z_LS,SE_SW0,z_SW0,SE_SW3,z_SW3,SE_SW4,z_SW4);
	tt_LS=round(tt_LS,3);

	est_lab=c("hat.gamma", "SE_LS", "z_LS", "SE_SW0", "z_SW0", "SE_SW3", "z_SW3", "SE_SW4", "z_SW4");
	colnames(tt_LS)=est_lab;
	cat("-------regression coefficients-------\n",append = T);

	print(tt_LS, quote=F,right=TRUE);

	cat(sep="",  "\n");
	cat("hat.sigma^2=", hsig2, "\n", append = T);

	LS_stat=list(hgamma, hsig2);
return(LS_stat);
}
#####################################################################
#            Obtain the C matrix and label the columns              #
#####################################################################

C_mat=function(L1,L2) { 

p_1=ncol(L1);
C_mat=NULL;
C_name=NULL;
L1_name=colnames(L1);
for (j in 1:p_1){
	L2j=L2[[j]];
	L2j_name=colnames(L2j);
	m_j=ncol(L2j);
	for (k in 1:m_j){
		xjuk=L1[,j]*L2j[,k];
		C_mat=cbind(C_mat,xjuk);
		namejk=paste(L1_name[j],L2j_name[k],sep='');
		C_name=c(C_name,namejk);
		};
};

colnames(C_mat)=C_name;

return(C_mat)
}

#####################################################################
#      evaluateing the R-square of level-1 coefficient beta         #
#          predicted by level-2 explanatory variables u             #
#####################################################################
Rsquare=function(x_name,u,gamma_ML,sig_ML) { 
	s_u=cov(u);
	m=ncol(u);
	u_name=colnames(u);
	xu_name=NULL;#collecting labeles for level-1 x and level-2 u;
	for (j in 1:m){
		xuj=paste(x_name,u_name[j],sep='');
		xu_name=c(xu_name,xuj);
	};
	
	gamma_xu=gamma_ML[c(xu_name),];
	sig_xx=sig_ML[c(paste(x_name,x_name,sep='')),];
	sig_model=t(gamma_xu)%*%s_u%*%gamma_xu;
	R_square=sig_model/(sig_model+sig_xx);
cat("    \n",append = T);
cat("R square for the coefficient of", x_name, "predicted by", colnames(u), "=", R_square, "\n");
cat("    \n",append = T);
};
#####################################################################
#       NML regression with SE from sandwich-type covariances       #
#####################################################################
NMLreg=function(L1,L2,H_mat,y) { 
C_mat=C_mat(L1,L2);
C_mat=as.matrix(C_mat);
H_mat=as.matrix(H_mat);
n=nrow(C_mat);
p_c=ncol(C_mat);
p_h=ncol(H_mat);

epsilon=.0001;
#starting value from LS;
  LS_stat=LSreg(C_mat,y);
  gamma_0=gamma_LS=LS_stat[[1]];
  hsig2_LS=LS_stat[[2]];
  V_one=matrix(1,n,1);
  V_onet=t(V_one);
  sigmat_0=hsig2_LS*V_onet%*%H_mat/(abs(sum(H_mat))+1);
  sigma_0=t(sigmat_0);
  delta=.01;

  cat("==========================================================\n")
  cat("==========================================================\n")

  cat("          Normal Distribution Maximum Likelihood          \n",append = T)
  
for (k in 1:300){
  #updating gamma;
	tau_vec=H_mat%*%sigma_0;	
	W=c(1/tau_vec);
	CW=C_mat*W; #columnwise multiplication;
	CWt=t(CW);
	CWC=CWt%*%C_mat;
	CWy=CWt%*%y; 
	CWC_in=solve(CWC);
	gamma_1=CWC_in%*%CWy; 

  #updating sigma;
	hy=C_mat%*%gamma_1;
	hate=y-hy;
	hatee=hate*hate; 
	tau_vec2=tau_vec^2;
	W2=c(.5/tau_vec2);
	HW=H_mat*W2; #columnwise multiplication;
	HWt=t(HW);
	HWH=HWt%*%H_mat;
	HWee=HWt%*%hatee;
	HWH_in=solve(HWH);
	sigma_1=HWH_in%*%HWee;

  #checking convergence;
	dt=max(abs(sigma_1-sigma_0),abs(gamma_1-gamma_0));
	sigma_0=sigma_1;
	gamma_0=gamma_1;
	if (dt<epsilon) break
}
if (k==300) {cat("The IRLS does not converge with 300 iterations in obtaining the NMLEs \n",append = T) 
    stop()};
cat("    \n",append = T);

cat("IRLS algorithm takes", k, "iterations to converge", "\n",append = T);

sum1=sum(log(tau_vec));
sum2=sum(hatee/tau_vec);
BIC_ML=sum1+sum2+(p_c+p_h)*log(n);
cat("    \n",append = T);
cat("BIC_ML=", BIC_ML, "\n",append = T);
cat("    \n",append = T);
  SE_MLgamma=sqrt(diag(CWC_in));
  SE_MLgamma=as.matrix(SE_MLgamma);
  z_MLgamma=gamma_1/SE_MLgamma;
  SE_SW=SW_ML(hate,tau_vec,C_mat,H_mat);#five sets of SEs are in this function;

  SE_SW0=SE_SW[[1]];
  SE_SW1=SE_SW[[2]];
  SE_SW4=SE_SW[[5]];    
  
  SE_SWgamma0=SE_SW0[1:p_c];
  z_SWgamma0=gamma_1/SE_SWgamma0;
  
  SE_SWgamma1=SE_SW1[1:p_c];
  z_SWgamma1=gamma_1/SE_SWgamma1;
  
  SE_SWgamma4=SE_SW4[1:p_c];
  z_SWgamma4=gamma_1/SE_SWgamma4;
  
  
  tt_MLgamma=cbind(gamma_1,SE_MLgamma,z_MLgamma,SE_SWgamma0,z_SWgamma0,SE_SWgamma1,z_SWgamma1,SE_SWgamma4,z_SWgamma4);
  tt_MLgamma=round(tt_MLgamma,3);
  est_lab=c("hat.gamma", "SE_ML", "z_ML", "SE_SW0", "z_SW0", "SE_SW1", "z_SW1", "SE_SW4", "z_SW4");
  cat("------------regression coefficients------------\n",append = T);
	colnames(tt_MLgamma)=est_lab;
	print(tt_MLgamma, quote=F, right=TRUE);
	cat("    \n", append = T);

  cat("------------variance parameters------------\n",append = T);
  est_lab=c("hat.sigma", "SE_ML", "z_ML", "SE_SW0", "z_SW0", "SE_SW1", "z_SW1", "SE_SW4", "z_SW4");
  SE_MLsig=sqrt(diag(HWH_in));  
  SE_MLsig=as.matrix(SE_MLsig);

  SE_SWsig0=SE_SW0[(p_c+1):(p_c+p_h)];
  SE_SWsig1=SE_SW1[(p_c+1):(p_c+p_h)];
  SE_SWsig4=SE_SW4[(p_c+1):(p_c+p_h)];    

  z_MLsig=sigma_1/SE_MLsig;
  z_SWsig0=sigma_1/SE_SWsig0;
  z_SWsig1=sigma_1/SE_SWsig1;
  z_SWsig4=sigma_1/SE_SWsig4;  
  
  tt_MLsig=cbind(sigma_1,SE_MLsig,z_MLsig,SE_SWsig0,z_SWsig0,SE_SWsig1,z_SWsig1,SE_SWsig4,z_SWsig4);
  tt_MLsig=round(tt_MLsig,3);
  colnames(tt_MLsig)=est_lab;
  print(tt_MLsig, quote=F,right=TRUE);

  cat("    \n", append = T);
  cat("Note: x0e is the label for sigma_0^2+sigma_e^2\n");
  cat("-------------------------------------------\n",append = T);

#The following is to generate R-square using labels;
p_1=ncol(L1);
L1_name=colnames(L1);
if (L1_name[1]=="x0"){ #11;
	for (j in 2:p_1){    #22;
		xj=L1_name[j];
		L2j=L2[[j]];
		p2j=ncol(L2j);
		L2j_name=colnames(L2j);
		if (L2j_name[1]=="u0"){ #33;
			u=L2j[,2:p2j];
            }else{
			u=L2j
			};#33;
		Rs_j=Rsquare(xj,u,gamma_1,sigma_1);
		};#22;#11;
}else{#44;
	for (j in 1:p_1){#55;
		xj=L1_name[j];
		L2j=L2[[j]];
		p2j=ncol(L2j);
		L2j_name=colnames(L2j);
		if (L2j_name[1]=="u0"){#66;
			u=L2j[,2:p2j];
            }else{
			u=L2j
			}#66;
		Rs_j=Rsquare(xj,u,gamma_1,sigma_1);
		};#55;#44;
};
gamma_sig=list(gamma_1,sigma_1,gamma_LS,hsig2_LS,SE_SW)
return(gamma_sig)
}