/*
************************************************************************
*   (C) Copyright 1999, Peter Lenk. All Rights Reserved.
*	MixReg.GSS
*	HB Linear Regression Model.
*	Heterogeneity for subject level parameters have a mixture of normals
*	Allows for subject level design matrices.
*	
*		Y_i = X_i*beta_i + epsilon_i for i = 1,..., nsub
*		epsilon_i is N(0,sigma^2 I)
*
*		beta_i = sum_{m=1}^M psi_m N(beta_i|theta_m,Lambda_n)
*		0 < psi_1 < psi_2 .. < psi_M and sum_{m=1}^M psi_m = 1.
*		delta_i is N(0,Lambda)
*
*	PRIORS
*		sigma^2 is Inverted Gamma(r0/2,s0/2)
*		theta_m is normal(u0,v0).
*		Lambda_m is Inverted Wishart(f0, g0) 
*
*---->Note:
*		This program saves subject-level beta_i as columns, not rows.
*		beta is a nsub by rankx matrix
*
***********************************************************************
*/
new;
mmod		= 3;			@ mmod = number of mixture components @
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							@
flagtrue	= 1;				@ 1 -> knows true parameters from simulation				@

pcount  = 5;   		@ max number of times that assignments of subjects 			@
					@ to groups can conflict with ordering: n_1 <= ... <= n_k	@
					@ where n_j = number of subjects in each group.         	@
pflag = 0;			@ flag for runs in viloations								@
					@ pflag = 0 -> last iterations not violation				@
					@ pflag = 1 -> last iteration is a violation				@


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


mint 		= seqa(1,1,mmod);
/*
********************************************************************
*	Get data
********************************************************************
*/
@ 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 @
ynames		= xynames[rows(xynames)];	
xnames		= "Constant"|xynames[1:rows(xynames)-1];

@ Last row of xydata is Y. @
xdata		= xydata[.,1:cols(xydata)-1];		@ Get independent variables 	@
ydata		= xydata[.,cols(xydata)];		@ Get dependent variable	@


loadm	yrows	= yrows;	@ yrows gives the number of observations per subject @
nsub	= rows(yrows);		@ Number of subjects. @
ntot	= sumc(yrows);		@ Total number of observations. @
@ Create pointer based on yrows to access xdata and ydata @
b	= cumsumc(yrows);	@ cumsumc is the cumulative sum of yrows @
a	= 1|(1+b[1:nsub-1]);
iptxy	= a~b;
@ To use iptxy to get the design matrix and data for subject i: @
@ x_i = xdata[iptxy[i,1]:iptxy[i,2],.] @
@ y_i = ydata[iptxy[i,1]:iptxy[i,2]] @

xdim	= cols(xdata);
@ add intercepts to xdata  @
xdata 	= ones(ntot,1)~xdata;
rankx	= xdim+1;
@ Create pointer to get into stacked lambda matrix @
if mmod > 1;
	b = seqa(rankx,rankx,mmod);   @ {rankx, 2*rankx, ..., mmod*rankx } @
	a = 1|(1+b[1:mmod-1]);		@ {1, rankx+1, 2*rankx+1, ..., (mmod-1)*rankx+1} @
	iptgp = a~b;
else;
	iptgp = 1~rankx;
endif;


@ Compute some sufficient statistics @
@ Get point to access stacked matrices of xtx, xtxi, and xty @
b	= seqa(rankx,rankx,nsub);
a	= 1|(1+b[1:nsub-1]);
iptxt	= a~b;
xtx	= zeros(rankx*nsub,rankx);
xty	= zeros(rankx*nsub,1);
bhat	= zeros(rankx,nsub);		@ MLE of beta_i @
sse	= 0;
for i0 (1,nsub,1);	i = i0;
	xi	= xdata[iptxy[i,1]:iptxy[i,2],.];
	yi	= ydata[iptxy[i,1]:iptxy[i,2],.];
	xitxi	= xi'xi;
	xitxii	= invpd(xitxi);
	xityi	= xi'yi;
	xtx[iptxt[i,1]:iptxt[i,2],.] = xitxi;
	xty[iptxt[i,1]:iptxt[i,2],.] = xityi;
	bhat[.,i] = xitxii*xityi;
	resid	= yi - xi*bhat[.,i];
	sse	= sse + resid'resid;
endfor;
s2hat	= sse/ntot;		@ MLE of sigma2 @
sighat	= sqrt(s2hat);




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

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

@ {psi_m}, the mixture probabilities, have ordered Dirchlet proir. @
w0    = ones(mmod,1);  
xgam = seqa(1,1,mmod)';		@ used in generating ordered dirichlet @

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

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

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

/*
*******************************************************************
*	Initialize MCMC
******************************************************************
*/
psi		= mint/sumc(mint);			@ Class probabilities @
@ Assign membership based on random start                  @
zprob	= ones(nsub,1).*.psi';		@ Individual level prob.  .*. is Kronecker product. @
z 		= rndzmn(zprob);			@ Generate multinomials @
classn 	= sumc(z .== mint' );		@ number in each of the mmod classes @
rclass  = rankindx(classn,1);

beta = bhat;
theta	= zeros(rankx,mmod);
lambda	= zeros(rankx*mmod,rankx);
lambdai	= lambda;
bvars	= zeros(rankx*mmod,1);
for i0 (1,mmod,1); i = i0;
	if classn[i] > rankx+1;		@ We have some observations in this class @
		zi	= z .== i;
		bi	= selif(beta',zi)';	@ Get beta for subjects assigned to group i @
		mbi	= meanc(bi');
		resid	= bi - mbi;
		lbdai	= resid*resid'/classn[i];
	else;
		mbi	= zeros(rankx,1);
		lbdai	= eye(rankx);
	endif;
	theta[.,i] 							= mbi;
	lambda[iptgp[i,1]:iptgp[i,2],.] 	= lbdai;
	lbda								= invpd(lbdai);
	lambdai[iptgp[i,1]:iptgp[i,2],.] 	= lbda;
	bvars[iptgp[i,1]:iptgp[i,2]]		= diag(lbda);
endfor;
sigma2	= s2hat;
sigma	= sqrt(sigma2);

/*
********************************************************************************
* Arrays for saving iterations & computing posterior means & std
********************************************************************************
*/

betam	= zeros(rankx,nsub);	@ posterior mean of beta 	@
betas	= zeros(rankx,nsub);	@ posterior std of beta 	@
sigmag	= zeros(smcmc,1);		@ MCMC iterations for error std @
psig	= zeros(smcmc,mmod);
zprobm	= zeros(nsub,mmod);		@ Individual level class membership probabilities @
thdim	= rankx*mmod;
thetag	= zeros(smcmc,thdim);	
lambdag = zeros(smcmc,mmod*rankx);	@ MCMC iterations for heterogeniety variance @
lambdam	= zeros(rankx*mmod,rankx);
lambdas = lambdam;
loglikeg	= zeros(smcmc,1);
loglike		= 0;
/*
********************************************************************
*	Do MCMC
********************************************************************
*/
@ Do the initial transition period @
nmcmc	= 0;						@ Just a counter @
ipcount = 0;						@ Counter of order restriction violations @
for i1 (1,nblow,1);	imcmc = i1;
	nmcmc = nmcmc + 1;	
	call getmixreg;
endfor;

for i1 (1,smcmc,1);	imcmc = i1;		@ Save smcmc iterations 			@
	for i2 (1,skip,1); jmcmc = i2;		@ Save every skip iterations 	@
		nmcmc = nmcmc + 1;
		call getmixreg;
	endfor;
	sigmag[imcmc]		= sigma;
	thetag[imcmc,.]		= vec(theta)';
	psig[imcmc,.]		= psi';
	betam				= betam + beta;
	betas				= betas + beta^2;
	zprobm				= zprobm + zprob;
	lambdam				= lambdam + lambda;
	lambdas				= lambdas + lambda^2;
	lambdag[imcmc,.]	= bvars';
	loglikeg[imcmc,.]	= loglike;
endfor;


/*
******************************************************************
*	Compute Posterior Means and STD
******************************************************************
*/
@ Compute Posterior Means @
betam		= betam/smcmc;
zprobm		= zprobm/smcmc;
sigmam		= meanc(sigmag);
psim		= meanc(psig);
thetam		= meanc(thetag);
thetam		= reshape(thetam,mmod,rankx)';
lambdam		= lambdam/smcmc;
loglikem	= meanc(loglikeg);

@ Compute Posterior STD @
betas		= sqrt( abs(betas - smcmc*betam^2)/smcmc);
sigmas		= stdc(sigmag);
psis		= stdc(psig);
thetas		= stdc(thetag);
thetas		= reshape(thetas, mmod, rankx)';
lambdas		= sqrt( abs(lambdas - smcmc*lambdam^2)/smcmc);
loglikes	= stdc(loglikeg);

@ Predict yi @
yhat	= zeros(ntot,1);
br	= zeros(nsub,1);	@ multiple R for each subject @
estd	= br;			@ error std			@
for i0 (1, nsub, 1); i = i0;
	xi	= xdata[iptxy[i,1]:iptxy[i,2],.];
	yi	= ydata[iptxy[i,1]:iptxy[i,2]];
	yhati	= xi*betam[.,i];
	yhat[iptxy[i,1]:iptxy[i,2]] = yhati;
	cy	= corrx(yi~yhati);	@ correlation matrix of yi and yhati @
	br[i]	= cy[1,2];
	resid	= yi - yhati;
	estd[i] = sqrt(resid'resid/yrows[i]);
endfor;


/*
****************************************************************
*	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("Mixture Probabilities versus Iteration");
xy(t,psig);
title("Heterogeneity Means versus Iteration");
xy(t,thetag);
title("Heterogeneity VARS versus Iteration");
xy(t,lambdag);
_plctrl = -1;
if flagtrue == 1;
	load betat 		= betat;
	load sigmat		= sigmat;
	load psit		= psit;
	load thetat		= thetat;
	load lambdat	= lambdat;
	title("True & HB Slope vs Intercept");
	xy(betat[1,.]'~betam[1,.]',betat[2,.]'~betam[2,.]');
endif;
title("HB & ML Slope vs Intercept");
xy(betam[1,.]'~bhat[1,.]',betam[2,.]'~bhat[2,.]');
graphset;


end;

/*
****************************************************************
* GETMIXREG
*	Does one iteration of the HB regression model.
*	INPUT
*		Global Variables
*	OUTPUT
*		Global Variables
****************************************************************
*/
PROC (0) = getmixreg;
local 
sse, i0,i,thetak,lambdaik,xi,yi,xitxi,xityi,vibn,vibn12,ebin,yhati,resid,sn,
fk,k,betak,mbetak,c,b,fnk,gnki,gnk,gnk12,w,lambdak,resid2,dlambik, zmaxp,
rclass, pflag,  lamb0,lamb2,j;


zprob	= zeros(nsub,mmod);	@ used in computing the P(z[i] = k) @

/*
***********************************************************
* Generate beta_i
* If Y_i belongs to class k, then
* beta_i is N(mbin, vbn)
* vbn 	= ( X_i'X_i/sigma2 + Lambda_k^{-1} }^{-1}
* mbin	= vbn*( X_i'Y_i/sigma2 + Lambda_k^{-1}*Theta_k)
**********************************************************
*/
sse	= 0;
for i0 (1,nsub,1); i = i0;
	thetak		= theta[.,z[i]];
	lambdaik	= lambdai[iptgp[z[i],1]:iptgp[z[i],2],.];
	xi 			= xdata[iptxy[i,1]:iptxy[i,2],.];
	yi 			= ydata[iptxy[i,1]:iptxy[i,2],.];
	xitxi 		= xtx[iptxt[i,1]:iptxt[i,2],.];
	xityi 		= xty[iptxt[i,1]:iptxt[i,2],.];
	vibn		= xitxi/sigma2 + lambdaik;
	vibn12		= chol(vibn);
	ebin		= xityi/sigma2 + lambdaik*thetak;
	beta[.,i]	= cholsol(ebin + vibn12'rndn(rankx,1), vibn12);
	yhati		= xi*beta[.,i];
	resid		= yi - yhati;
	sse			= sse + resid'resid;
endfor;

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

@ Compute log likelihood @
loglike	= -(ntot*ln(sigma2) + sse/sigma2)/2;
/*
**************************************************************
*  Generate theta_k and lambda_k
**************************************************************
*/
for fk (1,mmod,1); k = fk;
	@ do we have observations in class k? @
	if classn[k] > 0.5;
		betak      = (selif(beta', z .== k))';          @ beta with z=k@
		mbetak     = meanc(betak');
	else;
		mbetak     = zeros(rankx,1);
	endif;
	lambdaik   = lambdai[iptgp[k,1]:iptgp[k,2],.];

	/*
	************************************************************
	*  Generate theta_k given z, beta etc.
	*  theta_k is N(u_nk,v_nk)
	*  v_n,k = (n_k lambda_k^{-1} + v_0,k^{-1})
	*  u_n,k = v_n,k(n_k lambda_k^{-1} meanc(beta_k) + v_0,k^{-1} u_0,k
	************************************************************
	*/
	c 		= chol(classn[k]*lambdaik + v0i);
	b 		= classn[k]*lambdaik*mbetak + v0iu0;
	thetak	= cholsol(b+c'rndn(rankx,1),c);
	theta[.,k] = thetak;
	/*
	************************************************************
	*  Generate Lambda_k from IW_p(fnk, Gnk)
	*    fnk = f0k + classn_k
	*    Gnk =
	* (G0k^{-1} +  sum I(z_i=k)(beta_i-theta_k)(beta_i-theta_k)')^{-1}
	************************************************************
	*/
	if classn[k] > 0.5;    @ observations in class k @
		fnk   = f0 + classn[k];
		resid = betak - theta[.,k];
		gnki  = g0i + resid*resid';