###################################################################
#
#  Unit 7 - Example i.
#
#  Hierarchical models: Costs for U.S. Airlines
#
#  90 Observations on 6 firms for 1970-1984
#  
#  Variables:
#     Airline
#     Year
#     Output, in revenue passenger miles, index number
#     Total cost, in $1000
#     Fuel price
#     Load factor, the average capacity utilization of the fleet
#
###################################################################
#
# Author : Hedibert Freitas Lopes
#          Graduate School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoGSB.edu
#
###################################################################
q25 = function(x){quantile(x,0.25)}
q75 = function(x){quantile(x,0.75)}
Gibbs = function(sigma2,y,X,iC0th0,iC0,v1,v0s02){
  XtX = t(X)%*%X
  Xty = t(X)%*%y
  C1     = solve(iC0+XtX/sigma2)
  theta  = C1%*%(iC0th0+Xty/sigma2) + t(chol(C1))%*%rnorm(nrow(C1))
  sigma2 = 1/rgamma(1,v1/2,v0s02+sum((y-X%*%theta)^2)/2)  
  return(list(beta=theta,sig2=sigma2))
}


m       = 6
n       = 15
k       = 3
years   = 1970:1984
dataset = matrix(scan("airlines.txt"),n*m,6,byrow=T)
firm    = dataset[,1]
y1      = log(dataset[,3])
X1      = log(dataset[,4:6])

par(mfrow=c(2,2))
plot(firm,y1,xlab="Firms",ylab="",main="Log output")
plot(firm,X1[,1],xlab="Firms",ylab="",main="Log total cost")
plot(firm,X1[,2],xlab="Firms",ylab="",main="Log fuel price")
plot(firm,X1[,3],xlab="Firms",ylab="",main="Log load factor")

par(mfrow=c(2,2))
plot(years,y1[firm==1],xlab="years",ylab="",ylim=range(y1),type="l")
title("Log output")
for (i in 2:m)
  lines(years,y1[firm==i],col=i)
plot(years,X1[firm==i,1],xlab="years",ylab="",ylim=range(X1[,1]),type="l")
title("Log total cost")
for (i in 2:m)
  lines(years,X1[firm==i,1],col=i)
plot(years,X1[firm==i,2],xlab="years",ylab="",ylim=range(X1[,2]),type="l")
title("Log fuel price")
for (i in 2:m)
  lines(years,X1[firm==i,2],col=i)
plot(years,X1[firm==i,3],xlab="years",ylab="",ylim=range(X1[,3]),type="l")
title("Log load factor")
for (i in 2:m)
  lines(years,X1[firm==i,3],col=i)

# individual regression (Frequentist/Bayes with vague prior)
# ----------------------------------------------------------
table = NULL
resid   = NULL
for (i in 1:m){
  I        = firm==i
  y        = y1[I]    
  X        = cbind(1,X1[I,])     
  XtX      = t(X)%*%X
  iXtX     = solve(XtX)
  Xty      = t(X)%*%y
  b        = iXtX%*%Xty
  r        = y-X%*%b
  resid    = cbind(resid,r)
  s2       = sum(r^2)/(n-k-1)
  se       = sqrt(s2*diag(iXtX))
  table    = cbind(table,b)
  table    = cbind(table,se)
}
b.indreg      = table[1,seq(1,2*m,2)]
se.indreg     = table[1,seq(2,2*m,2)]
coef.indreg   = t(table[,seq(1,2*m,2)])[,2:(k+1)]
secoef.indreg = t(table[,seq(2,2*m,2)])[,2:(k+1)]
s2.indreg     = s2
cbind(round(t(table)[c(1,3,5,7,9,11),],2),round(t(table)[c(2,4,6,8,10,12),],2))

# Pooled regression (Frequentist/Bayes with vague prior)
# ------------------------------------------------------
X2       = cbind(1,X1)
XtX      = t(X2)%*%X2
iXtX     = solve(XtX)
Xty      = t(X2)%*%y1
b        = iXtX%*%Xty
r        = y1-X2%*%b
s2       = sum(r^2)/(length(y1)-k-1)
se       = sqrt(s2*diag(iXtX))
b.pooled = b[1]
se.pooled = se[1]
coef.pooled = b[2:(1+k)]
secoef.pooled = se[2:(1+k)]
s2.pooled = s2

plot(1:m,table[1,seq(1,12,2)],xlab="Firm",ylab="",main="Intercept")
abline(h=b[1],col=2)
plot(1:m,table[2,seq(1,12,2)],xlab="Firm",ylab="",main="Output")
abline(h=b[2],col=2)
plot(1:m,table[3,seq(1,12,2)],xlab="Firm",ylab="",main="Total Cost")
abline(h=b[3],col=2)
plot(1:m,table[4,seq(1,12,2)],xlab="Firm",ylab="",main="Fuel price")
abline(h=b[4],col=2)

# Individual effects model, nonhierarchical prior
# -----------------------------------------------
y = y1
X = matrix(0,n*m,m)
for (i in 1:m){
  XX = NULL
  for (j in 1:m)
    if (i==j){
      XX = cbind(XX,rep(1,n))
    }else{ 
      XX = cbind(XX,rep(0,n))
    }
  X[((i-1)*n+1):(i*n),] = XX
}
X = cbind(X,X1)
XtX  = t(X)%*%X
iXtX = solve(XtX)
Xty  = t(X)%*%y
b    = iXtX%*%Xty
r    = y-X%*%b
s2   = sum(r^2)/(length(y)-k-1)
se   = sqrt(s2*diag(iXtX))

b.nonhier = b[1:m]
se.nonhier = se[1:m]
coef.nonhier = b[(m+1):(m+k)]
secoef.nonhier = b[(m+1):(m+k)]
s2.nonhier = s2

cbind(b,se)
----------------------------
        parameter   standard
        estimates      error
----------------------------
alpha1  8.8019549 0.20826802
alpha2  8.7654389 0.21888229
alpha3  8.5947641 0.23069662
alpha4  8.9884205 0.24725654
alpha5  8.8268837 0.25908670
alpha6  8.8939219 0.26647114
COST    0.9187036 0.02976197
FUEL    0.4158084 0.01503234
FACTOR -0.5527304 0.10992766
----------------------------

par(mfrow=c(1,1))
plot(1:m,b.indreg,xlab="Firm",ylab="intercepts",pch=2,lwd=3,ylim=c(7,12))
for (i in 1:m){
  segments(i,b.indreg[i]-2*se.indreg[i],i,b.indreg[i]+2*se.indreg[i])
  points(i+0.1,b.nonhier[i],col=2,pch=3,lwd=3)
  segments(i+0.1,b.nonhier[i]-2*se.nonhier[i],i+0.1,b.nonhier[i]+2*se.nonhier[i],col=2)
}
abline(h=b.pooled,col=4,lwd=1,lty=2)
abline(h=b.pooled-2*se.pooled,col=4,lwd=1,lty=2)
abline(h=b.pooled+2*se.pooled,col=4,lwd=1,lty=2)
text(1.73,12,"Individual regressions",col=1)
text(2,11.8,"Individual effects, nonhierarchical prior",col=2)
text(1.67,11.6,"Pooled regression",col=4)

# Individual effects model, hierarchical prior
# --------------------------------------------
# Prior of mu_alpha
  mu.a = mean(b[1:m])
  sig2.a  = var(b[1:m])
# Prior of V_alpha
  ma = sig2.a
  da = 10
  va = ma^2/da^2+2
  vaVa2 = (va-1)*ma
  va = 2*va
  vaVa2 = 2*vaVa2
# Prior of beta_tilde
  bt = b[(m+1):(m+k)]
  Vb = 10*diag(se[(m+1):(m+k)]^2)
  iVb = solve(Vb)
  iVbbt = iVb%*%bt
# Prior of sigma2
  ma = var(r)[1,1]
  da = 10
  v0 = ma^2/da^2+2
  v0s02 = (va-1)*ma
  v0 = 2*v0
  v0s02 = 2*v0s02
  v1 = v0+n

# Initial values
alpha    = b[1:m]
mu.alpha = mean(alpha)
V.alpha  = var(alpha)/m
beta     = b[(m+1):(m+k)]
sig2     = var(r)[1,1]

# MCMC
M0  = 0
M   = 1000
LAG = 10
par = NULL
for (iter in 1:(M0+M*LAG)){
  # sampling alphas
  for (i in 1:m){
    z    = y[firm==i]-X1[firm==i,]%*%beta
    var  = 1/(1/V.alpha+n/sig2)
    mean = var*(mu.alpha/V.alpha+sum(z)/sig2) 
    alpha[i] = rnorm(1,mean,sqrt(var))
  }
  # sampling beta and sig2
  y1  = y-matrix(matrix(rep(alpha,n),n,m,byrow=T),n*m)
  run = Gibbs(sig2,y1,X1,iVbbt,iVb,v1,v0s02)
  beta = run$beta
  sig2 = run$sig2
  # sampling mu.alpha and V.alpha
  run = Gibbs(V.alpha,alpha,rep(1,m),mu.a/sig2.a,1/sig2.a,va+m,vaVa2)
  mu.alpha = run$beta
  V.alpha  = run$sig2
  par = rbind(par,c(alpha,mu.alpha,V.alpha,beta,sig2))
}
ind = seq((M0+1),(M0+LAG*M),by=LAG)
par1 = par
par = par1[ind,]

par = par1

par(mfrow=c(3,2))
for (i in 1:m){
  ts.plot(par[,i],xlab="iteration",ylab="")
  title(paste("alpha",i,sep=""))
}
par(mfrow=c(1,2))
ts.plot(par[,m+1],xlab="iteration",ylab="",main="mu.alpha")
ts.plot(par[,m+2],xlab="iteration",ylab="",main="V.alpha")

par(mfrow=c(2,2))
for (i in 1:k){
  ts.plot(par[,(m+2+i)],xlab="iteration",ylab="")
  title(paste("beta",i,sep=""))
}
ts.plot(par[,m+k+3],xlab="iteration",ylab="",main="sig2")

rbind(rbind(cbind(apply(par[,1:m],2,mean),sqrt(apply(par[,1:m],2,var))),
cbind(apply(par[,(m+3):(m+2+k)],2,mean),sqrt(apply(par[,(m+3):(m+2+k)],2,var)))),
rbind(apply(par[,(m+1):(m+2)],2,mean),sqrt(apply(par[,(m+1):(m+2)],2,var))))

b.hier = apply(par[,1:m],2,mean)
se.hier = sqrt(apply(par[,1:m],2,var))
coef.hier = apply(par[,(m+3):(m+k+2)],2,mean)
secoef.hier = sqrt(apply(par[,(m+3):(m+k+2)],2,var))
s2.hier = par[,(m+k+3)]

cbind(b.nonhier,se.nonhier,b.hier,se.hier)
-------------------------------------------------
                          PRIOR
          NON-HIERARCHICAL       HIERARCHICAL
-------------------------------------------------
        parameter   standard  parameter  standard
        estimates      error  estimates     error

-------------------------------------------------
alpha1  8.801955  0.2082680    8.803550 0.1316991
alpha2  8.765439  0.2188823    8.768506 0.1347153
alpha3  8.594764  0.2306966    8.603203 0.1316831
alpha4  8.988420  0.2472565    8.935631 0.1353796
alpha5  8.826884  0.2590867    8.782959 0.1394919
alpha6  8.893922  0.2664711    8.841829 0.1447149
m.alpha                        8.792618 0.1157588
V.alpha                        0.137541 0.0120244
-------------------------------------------------
m.alpha = mean(par[,m+1])
V.alpha = mean(sqrt(par[,m+2]))

par(mfrow=c(1,1))
plot(1:m,b.indreg,xlab="Firm",ylab="intercepts",pch=2,lwd=3,ylim=c(7,12),xlim=c(1,m+1.5),axes=F)
axis(2)
axis(1,at=1:7,label=c(1:6,"Pooled"))
for (i in 1:m){
  segments(i,b.indreg[i]-2*se.indreg[i],i,b.indreg[i]+2*se.indreg[i])
  points(i+0.1,b.nonhier[i],col=2,pch=3,lwd=3)
  segments(i+0.1,b.nonhier[i]-2*se.nonhier[i],i+0.1,b.nonhier[i]+2*se.nonhier[i],col=2)
  points(i+0.2,b.hier[i],col=4,pch=4,lwd=3)
  segments(i+0.2,b.hier[i]-2*se.hier[i],i+0.2,b.hier[i]+2*se.hier[i],col=4)

}
segments(7,b.pooled-2*se.pooled,7,b.pooled+2*se.pooled,col=3)
points(7,b.pooled,col=3,pch=5,lwd=3)
text(2.73,12,"Individual regressions",col=1,cex=1.5)
text(3.2,11.8,"Individual effects, nonhierarchical prior",col=2,cex=1.5)
text(3.1,11.6,"Individual effects, hierarchical prior",col=4,cex=1.5)


coefs = t(cbind(t(coef.indreg),coef.pooled,coef.nonhier,coef.hier))
cbind(coefs,c(s2.indreg,s2.pooled,s2.nonhier,se.hier))

-----------------------------------------------------------
                    Total      Fuel        Load    
                     Cost     Price      Factor      sigma2
-----------------------------------------------------------
Airline 1       1.1664029 0.3916901 -1.46136744 0.001413859
Airline 2       1.4648872 0.3103503 -1.52160585 0.015579075
Airline 3       0.7196370 0.4534382 -0.42409611 0.003543916
Airline 4       0.9371388 0.4590140 -0.37646810 0.131699070
Airline 5       1.0618380 0.2959101 -0.61319870 0.134715261
Airline 6       0.9675387 0.3001937  0.08667271 0.131683146
Pooled          0.8828541 0.4546868 -0.89146395 0.135379576
Nonhierarchical 0.9187036 0.4158084 -0.55273039 0.139491879
Hierarchical    0.8993396 0.4164568 -0.53900453 0.144714911
-----------------------------------------------------------

ses   = t(cbind(t(secoef.indreg),secoef.pooled,secoef.nonhier,secoef.hier))
L     = coefs-2*ses
U     = coefs+2*ses

par(mfrow=c(2,2))
names = c("Total Cost","Fuel Price","Load Factor")
for (j in 1:3){
  plot(1:9,coefs[,j],xlab="",ylab="",main="",axes=F,ylim=c(min(L[,j]),max(U[,j])))
  for (i in 1:9)
    segments(i,L[i,j],i,U[i,j])
  axis(2)
  axis(1,at=1:9,labels=c("Airline 1","Airline 2","Airline 3","Airline 4","Airline 5","Airline 6","Pooled","Nonhier","Hier"))
  title(names[j])
}

# Random coefficient model 
# ------------------------
b0     = apply(cbind(b.indreg,coef.indreg),2,mean)
V0     = diag(10000,k+1)
iV0    = solve(V0)
b00    = b0 
iVbb00 = iVb%*%b00
Vb     = diag(10000,k+1)
iVb    = solve(Vb)
v      = m+1
par1   = v+n 
V00    = diag(10000,k+1)
iV00   = solve(V00)
ma     = mean(s2.hier)
da     = 10
v0     = ma^2/da^2+2
v0s02  = (va-1)*ma
v0     = 2*v0
v0s02  = 2*v0s02

# Initial values
beta  = cbind(b.indreg,coef.indreg)
sig2  = ma

# Other quantities
X2 = array(0,c(n,k+1,m))
y2 = matrix(0,n,m)
for (i in 1:m){
  X2[,,i] = cbind(1,X1[((i-1)*n+1):(i*n),])
  y2[,i] = y[((i-1)*n+1):(i*n)]
}
nmv02 = (n*m+v0)/2

# MCMC
M0    = 1000
M     = 2000
LAG   = 10
MM    = M0+M*LAG
betas = array(0,c(m,k+1,MM))
sig2s = rep(0,MM)
b0s   = matrix(0,k+1,MM)
V0s   = array(0,c(k+1,k+1,MM))
for (iter in 1:MM){
  # sampling beta[i]
  for (i in 1:m){
    var  = solve(iV0+t(X2[,,i])%*%X2[,,i]/sig2)
    mean = var%*%(iV0%*%b0+t(X2[,,i])%*%y2[,i]/sig2)
    beta[i,] = mean + t(chol(var))%*%rnorm(k+1)
  }
  # sampling sig2
  mse = 0.0
  for (i in 1:m){
    res = y2[,i]-X2[,,i]%*%beta[i,]
    mse = mse + sum(res^2)
  }
  sig2 = 1/rgamma(1,nmv02,(v0s02+mse)/2)
  # sampling b0
  var  = solve(iVb+n*iV0)
  mean = var%*%(iVbb00+iV0%*%apply(beta,2,sum))
  b0   = mean + t(chol(var))%*%rnorm(k+1)

  # sampling V0 
  par2 = iV00 
  for (i in 1:m)
    par2 = par2 + (beta[i,]-b0)%*%t(beta[i,]-b0)
  par2 = solve(par2)
  draw = matrix(rnorm((k+1)*par1),par1,k+1)%*%chol(par2)
  iV0  = draw[1,]%*%t(draw[1,])
  for (i in 2:par1)
    iV0 = iV0 + draw[i,]%*%t(draw[i,])
  V0 = solve(iV0)

  # storing draws
  betas[,,iter] = beta
  sig2s[iter]   = sig2
  b0s[,iter]    = b0
  V0s[,,iter]   = V0
}
ind    = seq((M0+1),MM,by=LAG)
betas1 = betas[,,ind]
sig2s1 = sig2s[ind]
b0s1   = b0s[,ind]
V0s1   = V0s[,,ind]

names = c("Intercept","Total Cost","Fuel price","Load Factor")
par(mfrow=c(3,2))
plot(density(sig2s1),xlab="",ylab="",main="sigma2")
for (i in 1:(k+1))
  plot(density(b0s1[i,]),xlab="",ylab="",main=names[i])

par(mfrow=c(k+1,m))
for (i in 1:(k+1))
  for (j in 1:m)
    plot(density(betas1[j,i,]),xlab="",ylab="",main=paste("beta",j,i,sep=""),xlim=range(betas1[,i,]))

par(mfrow=c(k+1,k+1))
for (i in 1:(k+1))
  for (j in 1:(k+1)){
    if (i<=j){
      plot(density(V0s[i,j,]),xlab="",ylab="",main=paste("V0(",i,",",j,")",sep=""))
    }else{
      plot(0:5,xlab="",ylab="",col=0,axes=F)
    }
}

xx = NULL
for (i in 1:m)
 xx = cbind(xx,apply(betas[i,,],1,mean))
t(xx)

-----------------------------------------------------------
                    Total      Fuel        Load    
                     Cost     Price      Factor      sigma2
-----------------------------------------------------------
Airline 1       1.1664029  0.3916901 -1.46136744 0.001413859
Airline 2       1.4648872  0.3103503 -1.52160585 0.015579075
Airline 3       0.7196370  0.4534382 -0.42409611 0.003543916
Airline 4       0.9371388  0.4590140 -0.37646810 0.131699070
Airline 5       1.0618380  0.2959101 -0.61319870 0.134715261
Airline 6       0.9675387  0.3001937  0.08667271 0.131683146
Pooled          0.8828541  0.4546868 -0.89146395 0.135379576
Nonhierarchical 0.9187036  0.4158084 -0.55273039 0.139491879
Hierarchical    0.8993396  0.4164568 -0.53900453 0.144714911
-----------------------------------------------------------

















par2 = matrix(c(2,3,3,5),2,2)
draw = matrix(rnorm(2*20),20,2)%*%chol(par2)
iV0  = draw[1,]%*%t(draw[1,])
for (i in 2:20)
  iV0 = iV0 + draw[i,]%*%t(draw[i,])
V0 = solve(iV0)