# field trials example 
# source("test.s")

 # read in the data

y_matrix(scan("yield.dat"),ncol=14,byrow=T)
id_matrix(scan("variety.dat"),ncol=14,byrow=T)

nr_dim(y)[1]  # rows
nc_dim(y)[2]  # cols
n_nr*nc
ntr_length(unique(id))
nrep_as.vector(table(id))
 # varmap holds the locations of each variety
 # first 5 are id=1, next 5 are id=2,...
varmap_order(id)

# look at data
par(mfrow=c(2,2))
image(y)

 # make the matrix W

W_diag(rep(0,n)) 
for(i in 1:n){
  i0_(i-1)%/%nr + 1
  j0_(i-1)%%nr + 1
  for(j in 1:n){
    i1_(j-1)%/%nr + 1
    j1_(j-1)%%nr + 1
    if((abs(i0-i1)+abs(j0-j1))==1) W[i,j]_-1
  }
}
diag(W)_-apply(W,1,sum)

 # now set some MCMC parameters

iter_600
burn_100
ay_1; by_.005; aid_1; bid_.005; ax_1; bx_.005


 # arrays to hold the MCMC realizations
 # y = t + x + e

tout_matrix(NA,ncol=ntr,nrow=iter)
xout_matrix(NA,ncol=n,nrow=iter)
lamy_rep(NA,iter)
lamt_rep(NA,iter)
lamx_rep(NA,iter)

 # give initial values to parameters

xout[1,]_y
xout[1,]_lowess(1:n,as.vector(y),f=.1)$y  # or smooth version of y
image(matrix(xout[1,],ncol=14),zlim=range(y))
res_y-xout[1,]
lamy[1]_1
lamx[1]_.2
lamt[1]_.2
tout[1,]_as.vector(rep(1,5)%*%matrix(res[varmap],nrow=5)/5)

for(k in 2:iter){
  teff_matrix((tout[k-1,])[id],ncol=nc)
  rss_sum((y-teff-xout[k-1,])^2)
  lamy[k]_rgamma(1,ay+n/2)/(by+.5*rss)
 
  tss_sum(tout[k-1,]^2)
  lamt[k]_rgamma(1,aid+n/2)/(bid+.5*rss)

  xss_t(xout[k-1,])%*%W%*%xout[k-1,]
  lamx[k]_rgamma(1,ax+n/2)/(bx+.5*xss)

  res_y-xout[k-1,]
  tmean_as.vector(rep(1,5)%*%matrix(res[varmap],nrow=5))*lamy[k]
  tprec_nrep*lamy[k] + lamt[k]
  tmean_tmean/tprec
  tout[k,]_rnorm(ntr,tmean,1/sqrt(tprec))

  teff_matrix((tout[k,])[id],ncol=nc)
  prec_diag(rep(lamy[k],n))+lamx[k]*W
  vmat_solve(prec)
  vmat_.5*(vmat + t(vmat))
  mean_as.vector(lamy[k]*vmat%*%as.vector(y-teff))
  xout[k,]_rmultnorm(1,mean,vmat)
}

 # now get posterior means...

xhat_apply(xout[-(1:burn),],2,mean)
that_apply(tout[-(1:burn),],2,mean)
lamyhat_mean(lamy[-(1:burn)])
lamxhat_mean(lamx[-(1:burn)])

 # plot the solution

par(mfrow=c(2,2),oma=c(0,0,0,0), mar=c(5,5,2,2))
image(y)
mtext("data",side=3,line=1,cex=.8)
image(matrix(xhat,ncol=14),zlim=range(y))
mtext("posterior mean of fertility",side=3,line=1,cex=.8)

tbounds_apply(tout[-(1:burn),],2,quantile,c(.1,.9))
matplot(matrix(1:ntr,ncol=ntr,nrow=2,byrow=T),tbounds,lty=1,type="l",col=1)
points(1:ntr,that)
mtext("posterior mean and 80% CI for var effects",side=3,line=1,cex=.8)

boxplot(lamy[-(1:burn)],lamt[-(1:burn)],lamx[-(1:burn)],
        names=c("lam_y","lam_t","lam_x"))
mtext("posterior dist of precisions",side=3,line=1,cex=.8)