/*
************************************************************************
*   (C) Copyright 1999, Peter Lenk. All Rights Reserved.
*	HBMREG2.GSS
*
*	HBMREG1*.GSS
*		Has a matrix X_{ij} for each (ij) observation and a vector beta_i.
*	HBMREG2*.gss
*		has a vector x_{ij} for each (ij) obsevations and a matrix B_i.
*
*	HBMREG3.gss
*		This model allows different variables in the design matrics.
*		Example: 
*			Y_j = country j's inflation rate
*			X_j = constant ~ country j's nominal interest rate.
*		In the "standard" multivariate model, each Y_j uses the same X's.	
*
*		Y_j 	= X_j*beta_j + epsilon_j
*		j = 1, ..., npop is the model for population j.
*
*		Y_j is a nobs by 1 vector.
*		X_j is a nobs by rankx design matrix
*		beta_j is rankx by 1 vector of regression coefficients.
*
*		beta_j = Theta'z_j + delta_j for j = 1, ..., npop
*
*		B = Z*Theta + Delta
*		B = {beta_1', beta_2', ..., beta_npop'}
*		Z = {z_1', z_2', ..., z_npop' }
*		Z is npop x q
*
*		B* = vec(B') stacks the beta_j.
*		Theta* = vec(Theta)
*		B* = (Z.*.I_p)Theta* + Delta*
*
*		INPUT
*		xydata = { X1 X2 ... Xnpop Y1 Y2 ... Ynpop }
*				Xj's include a vector of ones for the intercepts.
*		zdata	= { z1', z2', ..., znpop' }
*				first column is a vector of ones for the intercepts.
*		parall = { npop, nobs, rankx, rankz }
*
***********************************************************************
*/
new;

outfile		= "results1.dat";	@ Specify output file for saving results 					@
								@ outfile is a string variable that contains a file name 	@
inxy 		= "xydata";			@ Name of Gauss file with X,Y data							@
inz	 		= "zdata";			@ Name of Gauss file with Z data							@
flagtrue	= 1;				@ 1 -> knows true parameters from simulation				@
/*
********************************************************************
*	Initialize parameters for MCMC 
********************************************************************
*/
smcmc		= 100;		@ number of iterations to save for analysis 				@
skip		= 1;		@ Save every skip iterations								@
nblow		= 100;		@ Initial transition iterations 							@
nmcmc		= nblow + skip*smcmc;	@ total number of iterations					@

/*
********************************************************************
*	Get data
********************************************************************
*/
load parall = parall;		@ Parameters @
npop	= parall[1];		@ number of populations @
nobs	= parall[2];		@ number of observations per population @
rankx	= parall[3];		@ rank of each X design matrix @
rankz	= parall[4];		@ rank of Z design matrix @

@ Input Gauss files @
open f1 	= ^inxy;	@ Get Gauss file for X, Y data 				@
				@ Opens Gauss file & assigns file handle f1 @
xydata		= readr(f1,rowsf(f1));		
			@ readr reads in Gauss file with file handle f1.			@
			@ rowsf(f1) returns the number of rows in the Gauss file. 	@
			@ readr reads rowsf(f1) rows, which is the entir dataset.	@
ci			= close(f1);
xynames		= setvars(inxy);			@ Get the variable names that accompnay X, Y data @
nxy			= cols(xydata);

ynames		= xynames[nxy-npop+1:nxy];	
xnames		= xynames[1:nxy-npop];
xdata		= xydata[.,1:nxy-npop];
ydata		= xydata[.,nxy-npop+1:nxy];
bnames		= 0 $+ "Int" | "X" $+ ftocv(seqa(1,rankx,1),2,0);




open f1		= ^inz;
zdata		= readr(f1,rowsf(f1));	
ci			= close(f1);
znames		= setvars(inz);



dimomega	= rankx*rankz;				@ dimension of vec(omega) @



@ Compute some sufficient statistics @
xtx	= xdata'xdata;
xty	= xdata'ydata;
ztz	= zdata'zdata;



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

@ Prior for Sigma is IW(sf0,sg0)  @
sf0	= npop+4; sfn = sf0+nobs;
sg0	= eye(npop);
sg0i = invpd(sg0);


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

u0 		= zeros(dimomega,1);
v0   	= 100*eye(dimomega);		@ dimomega = rankx*rankz    @
v0i  	= invpd(v0);    		@ used in updating omega @
v0iu0 	= v0i*vec(u0);          @ used in updating omega @

@ Lambda^{-1} is W_rankx(f0,g0 ) @
@ f0 = prior df, g0 = prior scale matrix @
g0i 	= eye(rankx);   @ g0^{-1} @
f0 		= rankx+2;  f0n = f0 + npop;
g0i		= zeros(rankx,rankx);
f0		= 0;	f0n = f0 + npop;



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

beta 		= zeros(npop,rankx);
sigma		= eye(npop);
sigmai		= invpd(sigma);

omega		= zeros(rankz,rankx);
lambda		= eye(rankx);
lambdai		= invpd(lambda);

@ Define data structures for saving iterates & computing posterior means & std @
betam	= zeros(npop,rankx);	@ posterior mean of beta 	@
betas	= zeros(npop,rankx);	@ posterior std of beta 	@
c		= npop*(npop+1)/2;
sigmag	= zeros(smcmc,c);
omegag	= zeros(smcmc,dimomega);
c		= rankx*(rankx+1)/2;
lambdag = zeros(smcmc,c);	@ save unique elements of lambda @


/*
********************************************************************
*	Do MCMC
********************************************************************
*/
@ Do the initial transition period @
for i1 (1,nblow,1);	imcmc = i1;
	/*
	************************************************************
	* Generate beta and Sigma from multivariate model
	************************************************************
	*/
	{beta, sigma, sigmai} =
	gethbmreg3(ydata,xdata,zdata,xtx,xty,beta,sigma,sigmai,omega,lambda,lambdai, sfn, sg0i);

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

endfor;

for i1 (1,smcmc,1);	imcmc = i1;		@ Save smcmc iterations 			@
	for i2 (1,skip,1); jmcmc = i2;		@ Save every skip iterations 	@

		/*
		************************************************************
		* Generate beta and Sigma from multivariate model
		************************************************************
		*/
		{beta, sigma, sigmai} =
		gethbmreg3(ydata,xdata,zdata,xtx,xty,beta,sigma,sigmai,omega,lambda,lambdai, sfn, sg0i);

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

	endfor;
	sigmag[imcmc,.]	= vech(sigma)';
	omegag[imcmc,.]	= vecr(omega)';
	betam			= betam + beta;
	betas			= betas + beta^2;
	lambdag[imcmc,.]	= vech(lambda)'; @ vech gets unique elements of symmetric matrix @
endfor;


/*
******************************************************************
*	Compute Posterior Means and STD
******************************************************************
*/
@ Compute Posterior Means @
betam		= betam/smcmc;
sigmam		= xpnd(meanc(sigmag));
omegam		= meanc(omegag);
omegam		= reshape(omegam,rankz,rankx);
lambdam		= xpnd(meanc(lambdag));	@ xpnd is opposite of vech @

@ Compute Posterior STD @
betas		= sqrt( abs(betas - smcmc*betam^2)/smcmc);
sigmas		= xpnd(stdc(sigmag));
omegas		= stdc(omegag);
omegas		= reshape(omegas, rankz, rankx);
lambdas		= xpnd(stdc(lambdag));

@ Predict yi @
yhat	= zeros(nobs,npop);

rmse	= zeros(npop,1);
multr	= zeros(npop,1);
rsquare	= zeros(npop,1);

bot = 1; top = rankx;
for i0 (1, npop, 1); i = i0;
	xi			= xdata[.,bot:top];
	yi			= ydata[.,i];
	betai		= betam[i,.]';
	yhati		= xi*betai;
	yhat[.,i] 	= yhati;
	resid		= yi - yhati;
	rmse[i]		= sqrt(resid'resid/nobs);
	c			= corrx(yi~yhati);
	multr[i]	= c[1,2];
	rsquare[i]	= c[1,2]^3;

	bot = top + 1; top = top + rankx;
endfor;




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


@ Plot saved iterations against iterations number @
t 	= seqa(nblow+skip,skip,smcmc);		@ saved iteration number @
title("Error Variance versus Iteration");
xy(t,sigmag);
title("Theta versus Iteration");
xy(t,omegag);
title("Lambda versus Iteration");
xy(t,lambdag);

graphset;			@ Return to default settings @


end;

/*
****************************************************************
* GETHBMREG3
*	Does one iteration of the HB regression model.
*	INPUT
*		ydata
*		xdata
*		zdata
*		xtx
*		xty
*		beta
*		sigma
*		sigmai
*		omega
*		lambda
*		lambdai
*		sfn
*		sg0i
*	OUTPUT
*		beta
*		sigma
*		sigmai

{beta, sigma, sigmai} =
gethbmreg3(ydata,xdata,zdata,xtx,xty,beta,sigma,sigmai,omega,lambda,lambdai, sfn, sg0i);

*
****************************************************************
*/
PROC (3) = gethbmreg3(ydata,xdata,zdata,xtx,xty,beta,sigma,sigmai,omega,lambda,lambdai, sfn, sg0i);

local omstar, c, d, vibn, vibn12, ebn, bstar, resid, bot, top,
 fj, j, xi, sgni, sgn, npop, rankx; 
	rankx 	= cols(beta);
	npop	= rows(beta);
	omstar	= vecr(omega);
	/*
	**********************************************************
	* Generate bstar  = vec(beta') 
	**********************************************************
	*/

	c		= sigmai.*.ones(rankx,rankx);
	d		= eye(npop).*.lambdai;
	vibn	= c.*xtx + d;
	vibn12	= chol(vibn);
	c		= sigmai.*.ones(rankx,1);

	ebn		= sumc( (c.*xty)' ) + (zdata.*.lambdai )*omstar;
	bstar	= cholsol(ebn+vibn12'rndn(npop*rankx,1), vibn12);
	beta	= reshape(bstar,npop,rankx);



	resid = zeros(nobs,npop);
	bot	= 1; top = rankx;
	for fj (1,npop,1); j = fj;