/*
************************************************************************
*   (C) Copyright 1999, Peter Lenk. All Rights Reserved.
*	HBREG1.GSS
*	HB Linear Regression Model.
*	Common Design Matrix X.
*	(HBREG2 allows different design matrics.)
*	
*		Y_i = X*beta_i + epsilon_i for i = 1,..., nsub
*		Assume common design matrix X for all subjects.
*		epsilon_i is N(0,sigma^2*I)
*
*		beta_i = Theta'z_i + delta_i
*		delta_i is N(0,Lambda)
*		B		= Z*Theta + Delta
*
*	PRIORS
*		sigma^2 is Inverted Gamma(r0/2,s0/2)
*		Theta is maxtrix normal(u0,v0).
*		That is, vec(Theta) is N(vec(u0),v0).
*		vec(theta) stacks the columns of theta.
*		Lambda is Inverted Wishart(f0, g0) 
***********************************************************************
*/
new;
inx		= "xdata";			@ Gauss file for common, independent variables 				@
iny		= "ydata";			@ Gauss file for dependent variables						@
inz		= "zdata";			@ Gauss file for covariates									@
outfile	= "hbreg1.res";	@ Specify output file for saving results 					@
							@ outfile is a string variable that contains a file name 	@
flagtrue	= 1;			@ 1 -> knows true parameters from simulation				@

/*
********************************************************************
*	Initialize parameters for MCMC 
********************************************************************
*/
smcmc		= 10000;		@ number of iterations to save for analysis 				@
skip		= 1;		@ Save every skip iterations								@
nblow		= 10000;		@ Initial transition iterations 							@
nmcmc		= nblow + skip*smcmc;	@ total number of iterations					@

/*
********************************************************************
*	Get data
********************************************************************
*/
open f1 	= ^inx;	@ Get Gauss file for X data 				@
xdata		= readr(f1,rowsf(f1));		
ci			= close(f1);
xnames		= setvars(inx);
xnames		= "Constant"|xnames;

open f1 	= ^iny;	@ Get Gauss file for Y data 				@
ydata		= readr(f1,rowsf(f1));		
ci			= close(f1);
ynames		= setvars(iny);

nsub	= rows(ydata);
mobs	= cols(ydata);
xdim	= cols(xdata);


open f1 	= ^inz;	@ Get Gauss file for Z data 				@
zdata		= readr(f1,rowsf(f1));		
ci			= close(f1);

zdim	= cols(zdata);
if rows(zdata) < 2;     @ No z-covariates @
    @ No z-covarites => zdata is the observation 0 @
    zdata = ones(nsub,1);
    znames = "CNST";
    zdim    = 0;
else;                   @ Got z-covariates @
    zdata	= ones(nsub,1)~zdata;
    znames		= setvars(inz);
    znames		= "CNST"|znames;
endif;



if not rows(xdata) == mobs;
	errorlog "XDATA has the wrong number of rows.";
endif;

@ add intercepts to xdata @
xdata 	= ones(mobs,1)~xdata;

rankx	= cols(xdata);
rankz	= cols(zdata);
thdim	= rankx*rankz;		@ dimension of vec(theta) @

@ Compute some sufficient statistics @
xtx 	= xdata'xdata;
xtxi	= invpd(xtx);
yx		= ydata*xdata;
ztz		= zdata'zdata;



@ MLE for beta @

bhat	= yx*xtxi;			@ nsub by rankx matrix: rows are bhat_i'	@
yhat	= bhat*xdata';		@ estimate values of ydata					@
resid	= ydata - yhat;
sse		= resid^2;
sse		= sumc(sumc(sse));
s2hat	= sse/(nsub*mobs);	@ MLE of error variance sigma2				@
sighat	= sqrt(s2hat);


/*
********************************************************************
*	Initialize Priors 
********************************************************************
*/

@ Prior for sigma2 is IG(r0/2, s0/2) @
r0 	= 2; s0 = 2;
rn 	= r0 + nsub*mobs;


@ Prior for theta is N(u0,v0) @

u0 		= zeros(rankz,rankx);
v0   	= 100*eye(thdim);			@ thdim = rankx*rankz    @
v0i  	= invpd(v0);    			@ used in updating theta @
v0iu0 	= v0i*vecr(u0);          	@ used in updating theta @

@ Lambda^{-1} is W_rankx(f0,g0 ) @
@ f0 = prior df, g0 = prior scale matrix @
f0 		= rankx+5;  f0n = f0 + nsub;
g0i 	= f0*eye(rankx);   @ g0^{-1} @

/*
*******************************************************************
*	Initialize MCMC
******************************************************************
*/

beta 	= zeros(nsub,rankx);
sigma 	= 1;
sigma2	= s2hat;
theta	= zeros(rankz,rankx);
thetav	= vecr(theta);
lambda	= eye(rankx);
lambdai	= invpd(lambda);

@ Define data structures for saving iterates & computing posterior means & std @
betam	= zeros(nsub,rankx);	@ posterior mean of beta 	@
betas	= zeros(nsub,rankx);	@ posterior std of beta 	@
sigmag	= zeros(smcmc,1);
thetag	= zeros(smcmc,thdim);
c		= rankx*(rankx+1)/2;
lambdag = zeros(smcmc,c);		@ save  lambda @
rsquarehb = zeros(nsub,1);      @ Bayesian R-Square @


/*
********************************************************************
*	Do MCMC
********************************************************************
*/
@ Do the initial transition period @
for i1 (1,nblow,1);	imcmc = i1;
	call gethbreg;
endfor;

for i1 (1,smcmc,1);	imcmc = i1;		@ Save smcmc iterations 			@
	for i2 (1,skip,1); jmcmc = i2;		@ Save every skip iterations 	@
		call gethbreg;
	endfor;
	betam			= betam + beta;
	betas			= betas + beta^2;
	sigmag[imcmc]	= sigma;
	thetag[imcmc,.]	= vecr(theta)';
	lambdag[imcmc,.]= vech(lambda)';
    yhathb          = beta*(xdata'); @ predicted value for HB @
    for j0 (1,nsub,1); j = j0;
        cr = corrx(ydata[j,.]'~yhathb[j,.]');
        rsquarehb[j] = rsquarehb[j] + cr[1,2]^2;
    endfor;

endfor;


/*
******************************************************************
*	Compute Posterior Means and STD
******************************************************************
*/

betam		= betam/smcmc;
sigmam		= meanc(sigmag);
thetam		= reshape(meanc(thetag),rankz,rankx);
lambdam		= xpnd(meanc(lambdag));
rsquarehb   = rsquarehb/smcmc;


betas		= sqrt( abs(betas - smcmc*betam^2)/smcmc);
sigmas		= stdc(sigmag);
thetas		= reshape(stdc(thetag),rankz,rankx);
lambdas		= xpnd( stdc(lambdag));

yhhb		= betam*(xdata');			@ predicted value for HB @
yhml		= bhat*(xdata');			@ predicted value for ML @
cy			= corrx(vec(ydata)~vec(yhhb));
rhb			= cy[1,2];
cy			= corrx(vec(ydata)~vec(yhml));
rml			= cy[1,2];

@ save individual level betas to file @
s = seqa(1,1,nsub);
bout = s~(rsquarehb*100)~betam;
bnames = 0 $+ "X" $+ ftocv(seqa(1,1,xdim),2,0);
bnames = "Subject"|"R2*100"|"CNST"|bnames;

eout = export(bout,"beta.xls",bnames');


/*
****************************************************************
*	Do some output
****************************************************************
*/
call outputanal;


@ Plot saved iterations against iterations number @
t 	= seqa(nblow+skip,skip,smcmc);		@ saved iteration number @
title("Error STD versus Iteration");
xy(t,sigmag);
title("Theta versus Iteration");
xy(t,thetag);
title("Lambda versus Iteration");
xy(t,lambdag);
title("MLE & HB for " $+ xnames[2]$+ " versus " $+ xnames[1]);
_plctrl = -1;
xy(bhat[.,1]~betam[.,1],bhat[.,2]~betam[.,2]);
graphset;			@ Return to default settings @

end;

/*
****************************************************************
* GETHBREG
*	Does one iteration of the HB regression model.
*	INPUT
*		Global Variables
*	OUTPUT
*		Global Variables
****************************************************************
*/
PROC (0) = gethbreg;
local vibn, vbn, vbn12, ebn, yhat, sse, sn, resid, gni, gn, gn12, w;

	/*
	***********************************************************
	* Generate beta
	* beta_i is N(mbin, vbn)
	* vbn 	= ( X'X/sigma2 + Lambda^{-1} }^{-1}
	* mbin	= vbn*( X'Y_i/sigma2 + Lambda^{-1}*Theta*Z_i)
	**********************************************************
	*/
	vibn	= xtx/sigma2 + lambdai;
	vbn		= invpd(vibn);
	vbn12	= chol(vbn);
	ebn		= (yx/sigma2 + zdata*theta*lambdai)*vbn;
	beta	= ebn + rndn(nsub,rankx)*vbn12;

	/*
	***************************************************************
	* Generate sigma2
	* sigma2 is IG(rn/2, sn/2)
	* rn = r0 + nsub*mobs
	* sn = s0 + sum_{i=1}^{nsub} (Y_i - X*beta_i)'(Y_i - X*beta_i)
	****************************************************************
	*/
	yhat 	= beta*(xdata');
	sse 	= (ydata - yhat)^2;
	sse 	= sumc(sumc(sse));
	sn		= s0 + sse;
	sigma2	= sn/(2*rndgam(1,1,rn/2)); 
	sigma	= sqrt(sigma2);

	/*
	***********************************************************
	* Generate Theta and Lambda from multivariate model:
	*	B = Z*Theta + N(0,Lambda)
	************************************************************
	*/
	{theta, lambda, lambdai} = 
	getmulreg(beta,zdata,ztz,theta,lambda,lambdai,v0i,v0iu0,f0n,g0i);

endp;


/*
****************************************************************
* GETMULREG
*	Generate multivariate regression parameters.
*		Yd		= Xd*parmat + epsilon
*	
*	INPUT
*		yd		= dependent variables
*		xd		= independet variables
*		xdtxd		= xd'xd
*