vbhetlm1<-function(y,X,Z,a=1+1e-5,b=1+1e-5,tol=1e-10) {
# Details:	Maximizing the lower bound for the heteroscedastic linear regression model.
# 		The prior mean vectors are assumed to be 0 and the prior covariance matrices
# 		are assumed to be the identity matrices, up to unknow scale parameters.
# References: "Variational approximation for heteroscedastic models and matching 
#		pursuit algorithms" by David Nott et al.
# Author:	Minh-Ngoc Tran (ngoctm@nus.edu.sg) with a helper function from David Nott
# Inputs:
# y		vector of responses
# X		design matrix for the mean model (you need to include the intercept, 
#   		if desired).
# Z		design matrix for the variance model (you need to include the intercept, 
#   		if desired).
# a,b		hyperparameters of the inverse gamma priors 
# tol		stopping tolerance for variational lower bound
# 
# Outputs:
# mub		mean in variational normal posterior for mean coefficients
# muv		mean in variational normal posterior for variance coefficients
# sb		covariance matrix in variational normal posterior for mean
#  		coefficients
# sv		covariance matrix in variational normal posterior for variance
#  		coefficients
# lb		vector of lower bound values over iterations
#
######################################################################

n<-length(y)
p<-dim(X)[2]
q<-dim(Z)[2]

vlb<-function(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv) {
  t1<-solve(psb,sb)
  t2<-solve(psv,sv)
  lb<-(p+q)/2-n/2*log(2*pi)+1/2*log(det(t1))+1/2*log(det(t2))
  lb<-lb-1/2*sum(diag(t1))-1/2*sum(diag(t2))
  v1<-mub-pmub
  lb<-lb-1/2*sum(v1*solve(psb,v1))
  v2<-muv-pmuv
  lb<-lb-1/2*sum(v2*solve(psv,v2))
  lv<-Z%*%muv
  lb<-lb-1/2*sum(lv)
  lm<-X%*%mub
  ln<-(y-lm)^2
  ln<-ln+apply((X%*%sb)*X,1,sum)
  ld<-exp(lv-1/2*apply((Z%*%sv)*Z,1,sum))
  lb<-lb-1/2*sum(ln/ld)
  return(lb)
}

lmfit = lm(y ~ X-1)
r = residuals(lmfit)
r<-log(r*r)
init.lm<-lm(r ~ Z-1)
muv<-coef(init.lm)
sv<-vcov(init.lm)
dd<-1/exp(Z%*%muv-1/2*apply((Z%*%sv)*Z,1,sum))
Xtilde<-matrix(rep(sqrt(dd),times=p),nrow=n)*X
ytilde<-sqrt(dd)*y
weighted.lm<-lm(ytilde ~ Xtilde-1)
mub<-weighted.lm$coef
sb<-vcov(weighted.lm)/summary(weighted.lm)$sigma^2
sig2b = (b+.5*sum(mub*mub)+.5*sum(diag(sb)))/(a+1+p/2)
psb = sig2b*diag(p)
sig2v = (b+.5*sum(muv*muv)+.5*sum(diag(sv)))/(a+1+q/2)
psv = sig2v*diag(q)
pmub = rep(0,p)
pmuv = rep(0,q)

lb<-vlb(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv)
lball<-lb
lbold<-lb-abs(lb)

while((lb-lbold)>tol) {
  lbold<-lb
  w<-(y-X%*%mub)
  w<-w*w+apply((X%*%sb)*X,1,sum)
  working.glm<-bgamglm(w,Z,pmuv,psv,muv,tol=tol)
  muv<-working.glm$coef
  sv = working.glm$variance

  sig2v = (b+.5*sum(muv*muv)+.5*sum(diag(sv)))/(a+1+q/2)
  psv = sig2v*diag(q)

  dd<-1/exp(Z%*%muv-1/2*apply((Z%*%sv)*Z,1,sum))
  D = Diagonal(x=as.numeric(dd))
  sb = solve(t(X)%*%D%*%X+solve(psb))
  mub = sb%*%(solve(psb,pmub)+t(X)%*%D%*%y)
  
  sig2b = (b+.5*sum(mub*mub)+.5*sum(diag(sb)))/(a+1+p/2)
  psb = sig2b*diag(p)

  lb<-vlb(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv)
  lball<-c(lball,lb)
}
list(mub=mub,muv=muv,sb=sb,sv=sv,sig2b=sig2b,sig2v=sig2v,lb=lball)

} # end of function vbhetlm1


bgamglm<-function(w,Z,pm,pv,p0,tol) {
# Helper function for the function vbhetlm1.
# This function is to compute posterior mode for gamma generalized
# linear model with log link, known shape parameters
# and normal prior on coefficients
#
# w - vector of responses
# Z - design matrix
# pm, pv - prior mean and covariance matrix
# p0 - starting guess for mode
# tolerance for stopping

xold = p0
V = .5*w*exp(-Z%*%xold)
fvalold = .5*sum(Z%*%xold)+sum(V)+0.5*sum((xold-pm)*solve(pv,(xold-pm)))
Df = .5*colSums(Z)-t(Z)%*%V+solve(pv,xold-pm)
W = Diagonal(x = as.numeric(V))
D2f = t(Z)%*%W%*%Z+solve(pv)
xnew = xold - solve(D2f,Df)
V = .5*w*exp(-Z%*%xnew)
fvalnew = .5*sum(Z%*%xnew)+sum(V)+0.5*sum((xnew-pm)*solve(pv,(xnew-pm)))
fvalall = fvalold

while((fvalold-fvalnew)>tol) {
  xold = xnew
  fvalold = fvalnew
  fvalall = c(fvalall,fvalnew)
  Df = .5*colSums(Z)-t(Z)%*%V+solve(pv,xold-pm)
  W = Diagonal(x = as.numeric(V))
  D2f = t(Z)%*%W%*%Z+solve(pv)
  xnew = xold - solve(D2f,Df)
  V = .5*w*exp(-Z%*%xnew)
  fvalnew = .5*sum(Z%*%xnew)+sum(V)+0.5*sum((xnew-pm)*solve(pv,(xnew-pm)))
}
list(coef=xnew,variance=solve(D2f),fval=fvalall)
}

