# The following is R code for the seizure data example of Madsen and Birkes (2013).


# Set up to do parallel processing.
library(doParallel) 
n_cores<-detectCores()-1
cl<-makeCluster(n_cores)
registerDoParallel(cl)
clusterCall(cl, function() library(mvtnorm))

g<-function(y,mu1,mu2,var1,var2,c) {
  Sigma<-matrix(c(1,c,c,1),2,2)  # This is one term in the sum in equation (8) of Madsen and Birkes.
  g<-(pmvnorm(upper=c(qnorm(pnbinom(y[[1]]-1,size=mu1^2/(var1-mu1),mu=mu1)),qnorm(pnbinom(y[[2]]-1,size=mu2^2/(var2-mu2),mu=mu2))),sigma=Sigma)[1]
       +pmvnorm(upper=c(qnorm(1-pnbinom(y[[1]],size=mu1^2/(var1-mu1),mu=mu1)),qnorm(1-pnbinom(y[[2]],size=mu2^2/(var2-mu2),mu=mu2))),sigma=Sigma)[1]
       -pmvnorm(lower=c(qnorm(pnbinom(y[[1]],size=mu1^2/(var1-mu1),mu=mu1)),-Inf),upper=c(Inf,qnorm(pnbinom(y[[2]]-1,size=mu2^2/(var2-mu2),mu=mu2))),sigma=Sigma)[1]
       -pmvnorm(lower=c(-Inf,qnorm(pnbinom(y[[2]],size=mu2^2/(var2-mu2),mu=mu2))),upper=c(qnorm(pnbinom(y[[1]]-1,size=mu1^2/(var1-mu1),mu=mu1)),Inf),sigma=Sigma)[1])
}

library(mvtnorm)
library(geepack)
f<-function(c,mu1,mu2,var1,var2,rho_sb) {
  K1<-ceiling(qnbinom(sqrt(1-.005), size=mu1^2/(var1-mu1),mu=mu1))
  K2<-ceiling(qnbinom(sqrt(1-.005), size=mu2^2/(var2-mu2),mu=mu2)) # Need an upper limit for sums to infinity
  y<-expand.grid(x=seq(0,K1),y=seq(0,K2))
  ylist<-split(y,as.numeric(row.names(y)))
  s<-sum(dnbinom(y[,'x'],size=mu1^2/(var1-mu1),mu=mu1)*dnbinom(y[,'y'],size=mu2^2/(var2-mu2),mu=mu2)*unlist(lapply(ylist,g,mu1=mu1,mu2=mu2,var1=var1,var2=var2,c=c)))
  return(3*s-rho_sb) # From equation (8) in Madsen and Birkes.
}

# This section estimates the parameters of the seizure data.
data(seizure)
seizure<-seizure[c(1:48,50:59),]#Omit obs. 49 (patient 207).

ns<-dim(seizure)[1] #Number of subjects
cs<-5 #Number of observations per subject
n<-cs*ns #Number of observations.

seiz.long <- reshape(seizure,
                  varying=list(c("base","y1", "y2", "y3", "y4")),
                  v.names="y", times=0:4, direction="long")
seiz.long <- seiz.long[order(seiz.long$id, seiz.long$time),]
seiz.long$t <- ifelse(seiz.long$time == 0, 8, 2)
seiz.long$x <- ifelse(seiz.long$time == 0, 0, 1)
seiz.long$xtrt <- seiz.long$x*seiz.long$trt
m1 <- geese(y ~ offset(log(t))+x+trt+x:trt,id = id,
            data=seiz.long, corstr="exch", family=poisson)
#summary(m1)

#      trt age time  y id t x xtrt
#1.0    0  31    0  11  1 8 0  0
#2.0    0  30    0  11  2 8 0  0

muY.tmp<-exp(log(seiz.long$t)+m1$beta[1]+as.matrix(seiz.long)%*%c(m1$beta[3],0,0,0,0,0,m1$beta[2],m1$beta[4]))
data.tmp<-reshape(cbind(seiz.long,muY.tmp),direction="wide")
data<-cbind(data.tmp$id,data.tmp$y.0,data.tmp$y.1,data.tmp$y.2,data.tmp$y.3,data.tmp$y.4,data.tmp$muY.tmp.0,data.tmp$muY.tmp.1,data.tmp$muY.tmp.2,data.tmp$muY.tmp.3,data.tmp$muY.tmp.4)

# Here are the target means and variances for the seizure data.
muY<-array(t(data[,7:11]),dim=c(n,1))
varY<-muY*m1$gamma

# Calculate the rank correlation for the seizure data.
seizure.rank=seizure[,c(6,1:4)]
for (j in 1:cs) {
	seizure.rank[,j]=rank(seizure.rank[,j])    # Default is midrank.
	seizure.rank[,j]=rank(seizure.rank[,j],ties.method='random')
}
# CorrY is the block-diagonal matrix of rank correlations.
CorrY<-kronecker(diag(1,ns),cor(seizure.rank[,1:5])) # Matrix of rho^{S(b)}
diag(CorrY)<-1

# Calculate the ties correction.
K<-ceiling(max(muY+5*sqrt(varY))) # Need an upper limit for infinite sums.
Ef.j.cube<-rowSums(matrix((dnbinom(kronecker(0:K,matrix(1,length(muY),1)),size=kronecker(matrix(1,K+1,1),muY^2/(varY-muY)),mu=kronecker(matrix(1,K+1,1),muY)))^3,length(muY),K+1))
a.tmp<-sqrt((1-Ef.j.cube)%*%(1-t(Ef.j.cube)))


# This section figures out all the unique combinations of pairs of means, variances, and rank correlation
# so we don't have to solve the equation (9) of Madsen & Birkes any more times than necessary.
mu1<-kronecker(matrix(1,1,n),muY)[upper.tri(kronecker(matrix(1,1,n),muY))]
mu2<-t(kronecker(matrix(1,1,n),muY))[upper.tri(t(kronecker(matrix(1,1,n),muY)))]
var1<-kronecker(matrix(1,1,n),varY)[upper.tri(kronecker(matrix(1,1,n),varY))]
var2<-t(kronecker(matrix(1,1,n),varY))[upper.tri(t(kronecker(matrix(1,1,n),varY)))]
c<-CorrY[upper.tri(CorrY)]
a<-a.tmp[upper.tri(a.tmp)]
c.star<-as.numeric(NA_character_)
mvc<-cbind(mu1,mu2,var1,var2,c,a,c.star)
unique.comb<-data.frame(unique(mvc))
n.unique<-dim(unique.comb)[1]

# For each "unique combination," solve equation (9) of Madsen & Birkes
ptm<-proc.time() 
c.star.list<-foreach(i = 1:n.unique) %dopar% {
  if (unique.comb[i,]$c==0) {
    0
  } else { 
    uniroot(f,c(max(-1,unique.comb[i,]$c-.2),min(1,unique.comb[i,]$c+.2)),tol=.001,mu1=unique.comb[i,]$mu1,mu2=unique.comb[i,]$mu2,var1=unique.comb[i,]$var1,var2=unique.comb[i,]$var2,rho_sb=unique.comb[i,]$c*unique.comb[i,]$a)$root
  }
}
unique.comb$c.star<- unlist(c.star.list)
et<-proc.time() - ptm
n.unique
#[1] 32
#> et
#user  system elapsed 
#0.06    0.02  141.71  

# This section takes the solutions of equation (9) and populates mvc, which has a row for each pair of points.
rlabels<-kronecker(matrix(1,1,n),1:n)[upper.tri(kronecker(matrix(1,1,n),1:n))]
clabels<-t(kronecker(matrix(1,1,n),1:n))[upper.tri(t(kronecker(matrix(1,1,n),1:n)))]
for (i in 1:n.unique) {
  mvc[which(abs(rowSums(kronecker(as.matrix(unique.comb[i,1:5]),matrix(1,nrow(mvc),1))-mvc[,1:5]))<=.Machine$double.eps),'c.star']<-unique.comb[i,'c.star']
}
 
# This part works.
ptm<-proc.time() 
nvectors<-10000 # Number of datasets to simulate. 
CorrZ<-matrix(0,n,n) 
CorrZ[cbind(rlabels,clabels)]<-mvc[,'c.star']  # Turn the c.star column of mvc into the matrix CorrZ. At first, it's upper triangular
CorrZ<-CorrZ+t(CorrZ)          # Get the lower triangle too.
diag(CorrZ)<-1
R<-chol(CorrZ) # To simulate correlated standard normals.
Y<-array(0,c(n,nvectors)) # Initialize Y to all zeros.
U<-Y
Kp<-ceiling(max(10*sqrt(varY)))
probs<-matrix(unlist(lapply(as.list(0:Kp),pnbinom,size=muY^2/(varY-muY),mu=muY)),length(muY),Kp+1) # probs[i,j] is P(Y_i=j-1)
for (v in 1:nvectors) {
	U[,v]<-pnorm(t(R)%*%rnorm(n))  #The U's are uniform(0,1).
	for (j in 1:Kp) {
    Y[which(probs[,j]<U[,v] & U[,v]<=probs[,j+1]),v]<-j
	}
}

# Check to see how long the above took.
et1<-proc.time() - ptm


## Obtain sample correlations, means, and variances, and plot vs. targets.
ObsCorr<-cor(t(Y),method="spearman")
xmax<-1.1*max(ObsCorr[lower.tri(ObsCorr)],CorrZ[lower.tri(CorrY)])
xmin<-min(ObsCorr[lower.tri(ObsCorr)],CorrZ[lower.tri(CorrZ)])-0.1*min(ObsCorr[lower.tri(ObsCorr)],CorrZ[lower.tri(CorrY)])

# Decimate the points a bit to get a smaller graphics file.
pts<-cbind(CorrY[lower.tri(ObsCorr)],ObsCorr[lower.tri(ObsCorr)])
pts.non0<-which(pts[,1]>0)
n.non0<-length(pts.non0)
pts.0<-which(pts[,1]<0.05)
n.0<-length(pts.0)
pts.dec<-pts[c(sample(pts.non0,size=0.5*n.non0),sample(pts.0,size=0.5*n.0)),]
par(oma=c(2,2,0,0))
plot(pts.dec[,1],pts.dec[,2],xlim=c(xmin,xmax),ylim=c(xmin,xmax),xlab="Target Rescaled Spearman Correlation",ylab="Sample Pearson Correlation of Midranks",cex.axis=1.5,cex.lab=1.5) #main='Figure 1(c)',
lines(c(xmin,xmax),c(xmin,xmax),lty=2)

plot(CorrY[lower.tri(ObsCorr)],ObsCorr[lower.tri(ObsCorr)],xlim=c(xmin,xmax),ylim=c(xmin,xmax),xlab="Target Rescaled Spearman Correlation",ylab="Sample Pearson Correlation of Midranks",cex.axis=1.5,cex.lab=1.5) #main='Figure 1(c)',
lines(c(xmin,xmax),c(xmin,xmax),lty=2)

ObsMean<-rowMeans(Y)
xmax<-1.1*max(ObsMean,muY)
xmin<-xmax/1.1-xmax
par(oma=c(2,2,0,0))
plot(muY,ObsMean,xlab="Target Means",ylab="Sample Means",cex.axis=1.5,cex.lab=1.5)
lines(c(xmin,xmax),c(xmin,xmax),lty=2)

ObsVar<-diag(cov(t(Y)))
xmax<-1.1*max(ObsVar,varY)
xmin<-xmax/1.1-xmax
par(oma=c(2,2,0,0))
plot(varY,ObsVar,xlab="Target Variances",ylab="Sample Variances",cex.axis=1.5,cex.lab=1.5)
lines(c(xmin,xmax),c(xmin,xmax),lty=2)