#
# Write a title of an analysis
#
banner <- function(..., width=50, ch="-") {
  cat("\n",rep(ch,width),"\n",sep="")
  cat(...,"\n",sep="")
  cat(rep(ch,width),"\n\n",sep="")
}
#
# Coefficient of relationship
#
kinship.rel <- function(ped) {
  kinval<-function(i,j, kin) { 
    ifelse(i>j, kin[i,j], kin[j,i])
  }
  num<-nrow(ped)
  nfound<-sum(is.na(ped$fa))
  kin<-diag(num) ; rownames(kin)<-colnames(kin)<-ped$id
  frat<-kin
  if (num>nfound) {
    fa.p<-match(paste(ped$famid,ped$fa), paste(ped$famid,ped$id))
    mo.p<-match(paste(ped$famid,ped$mo), paste(ped$famid,ped$id))
    for(i in (nfound+1):num) {
      for(j in 1:(i-1)) {
        kin[j,i]<-kin[i,j]<-0.5*(kinval(fa.p[i],j,kin)+ kinval(mo.p[i],j,kin))
        frat[j,i]<-frat[i,j]<-0.25*(kinval(fa.p[i], fa.p[j],kin)+
                                    kinval(mo.p[i], mo.p[j],kin)+
                                    kinval(fa.p[i], mo.p[j],kin)+
                                    kinval(mo.p[i], fa.p[j],kin))
      }
      kin[i,i]<-1.0+0.5*kinval(fa.p[i], mo.p[i],kin)
    }
  }
  list(famid=ped$famid[1],kinship=kin, fraternity=frat)
}
#
# read IBD matrices
#
kinship.emp <- function(ped, ibds) {
  num<-nrow(ped)
  if (is.null(ibds)) {
    kin=NULL
  }else{
    kin<-as.matrix(ped[,ibds[1:num]])
    rownames(kin)<-colnames(kin)<-ped$id
  }
  list(famid=ped$famid[1],kinship=kin)
}
#
# Variance components MVN maximizer
#
varcomp <- function(trait,ped,covariates=NULL,ibds=NULL, method.opt="Nelder-Mead") {
  banner("Variance components analysis of \"", trait,"\"")
  ntrait<-length(trait)
  vc.data<-create.vc.data(trait,covariates,ibds,ped) 
  ymean<-mean(ped[,trait], na.rm=T)
  yvar<-var(ped[,trait], na.rm=T)
  fixed<-c(ymean, rep(0,length(covariates)))
  nfixed<-ntrait*length(covariates)+1
  nobs<-sum(complete.cases(ped[,c(trait,covariates)]))
#
# Base AE model
#
  mod="ae"
  maxlik<-optim(c(fixed,0.1*yvar,0.9*yvar), optlik, gr = NULL, 
                method = method.opt, lower = -Inf, upper = Inf,
                control = list(fnscale=-2.0, reltol=1e-12), 
                hessian = FALSE, vc.data, nfixed, mod)
  va<-maxlik$par[nfixed+1]
  vq<-0.0
  ve<-maxlik$par[nfixed+2]
  base.ll <- maxlik$value
#
# AQE
#
  if (!is.null(ibds)) {
     mod="aqe"
     maxlik<-optim(c(fixed, 0.1*yvar, 0.0, 0.9*yvar), optlik, gr = NULL, 
                   method = method.opt, lower = -Inf, upper = Inf,
                   control = list(fnscale=-2.0, reltol=1e-12), 
                   hessian = FALSE, vc.data, nfixed, mod)
     va<-maxlik$par[nfixed+1]
     vq<-maxlik$par[nfixed+2]
     ve<-maxlik$par[nfixed+3]
  }
  test.ae.e <-nobs*(log(yvar)+1)+2*maxlik$value
  fixed<-maxlik$par[1:nfixed]
  vt<-va+vq+ve
  cat("Number of families         = ", 
      length(unique(ped$famid[complete.cases(ped[,trait])])),"\n",
      "Number of observations     = ", nobs,"\n",
      "Trait mean (intercept)     = ", fixed[1],"\n",sep="")
  if (nfixed>1) {
    for(i in 2:nfixed) {
      cat(formatC(covariates[i-1],width=27,flag="-"),"= ",fixed[i],"\n",sep="")
    }
  }
  cat("\n")
  cat("Additive genetic variance  = ", va ," (",round(100*va/vt,1),"%)\n",sep="")
  if (mod=="aqe") {
    cat("QTL genetic variance       = ", vq ," (",round(100*vq/vt,1),"%)\n",sep="")
  }
  cat("Environmental variance     = ", ve ," (",round(100*ve/vt,1),"%)\n",
      "Model loglikelihood        = ", maxlik$value,"\n\n",sep="")
  if (mod=="aqe") {
    lod=(maxlik$value-base.ll)/log(10)
    lrts=2*(maxlik$value-base.ll)
    cat("Lod score testing VQ=0     = ", lod, 
        " (df=1, P=",0.5*pchisq(lrts,1,lower=FALSE),")\n\n",sep="")
  }
  cat("Total function evaluations = ", maxlik$counts[1],"\n\n",sep="")
}
create.vc.data <- function(trait,covariates,ibds,ped) {
  use.vc<-function(x) {
    use<-complete.cases(data.frame(x$y,x$covariates))
    x$y<-as.matrix(x$y)[use,]
    x$covariates<-as.matrix(x$covariates)[use,]
    x$rel<-x$rel[use,use]
    if (!is.null(ibds)) x$ibd<-x$ibd[use,use]
    x
  }
  if (is.null(ibds)) {
    vc.data<-by(ped,ped$famid,
                function(x) list(y=x[,trait], 
                                 covariates=cbind(Intercept=rep(1,nrow(x)),
                                                  x[,covariates]),
                                 rel=kinship.rel(x)$kinship,
                                 ibd=matrix(0, nr=nrow(x), nc=nrow(x))))
  }else{
    vc.data<-by(ped,ped$famid,
                function(x) list(y=x[,trait], 
                                 covariates=cbind(Intercept=rep(1,nrow(x)),
                                                  x[,covariates]),
                                 rel=kinship.rel(x)$kinship,
                                 ibd=kinship.emp(x,ibds)$kinship))
  }
  lapply(vc.data,use.vc)
}
#
# function to evaluate the MVN likelihood
#
lik <- function(m, va, vq, ve, vc.data) {
   res<-sum(unlist(lapply(vc.data, lik.fam, m, va, vq, ve)))
   res
}
lik.fam <- function(vc.data, m, va, vq, ve) {
  ln<-function(x) ifelse(x>0, log(x), 0)
  yhat<-vc.data$covariates %*% matrix(m, nc=1)
  S<-va*vc.data$rel + vq*vc.data$ibd + diag(ve,nrow(vc.data$rel))
  res <- -0.5* (ln(det(S)) + t(vc.data$y-yhat) %*% solve(S) %*% (vc.data$y-yhat))
  res
}
#
# and the version to allow calling of optim() function maximizer
#
optlik <- function(par, vc.data, nfixed, mod) {
  if (mod=="ae") {
    lik(par[1:nfixed],par[nfixed+1], 0.0, par[nfixed+2],vc.data) 
  }else if (mod=="aqe") {
    lik(par[1:nfixed],par[nfixed+1], par[nfixed+2], par[nfixed+3],vc.data) 
  }
}
#
varcomp.GxE <- function(trait,ped,covariates,ibds=NULL, method.opt="Nelder-Mead") {
  banner("GxE Variance components analysis of \"", trait,"\"" )
  cat("Moderator: ", covariates,"\n",sep="")
  ntrait<-length(trait)
  vc.data<-create.vc.data(trait,covariates,ibds,ped) 
  ymean<-mean(ped[,trait], na.rm=T)
  yvar<-var(ped[,trait], na.rm=T)
  fixed<-c(ymean, rep(0,length(covariates)))
  nfixed<-ntrait*length(covariates)+1
  nobs<-sum(complete.cases(ped[,c(trait,covariates)]))
#
# AE base model
#
  mod="ae"
  maxlik<-optim(c(fixed,0.1*yvar,0.9*yvar,0.0), optlik2, gr = NULL, 
                method = method.opt, lower = -Inf, upper = Inf,
                control = list(fnscale=-2.0, reltol=1e-12), 
                hessian = FALSE, vc.data, nfixed, mod)
  va<-maxlik$par[nfixed+1]
  vq<-0.0
  ve<-maxlik$par[nfixed+2]
  b.ve <-maxlik$par[nfixed+3]
  base.ll <- maxlik$value
#
# AQE if required
#
  if (!is.null(ibds)) {
     mod="aqe"
     maxlik<-optim(c(fixed, 0.1*yvar, 0.0, 0.9*yvar, 0.0), optlik2, gr = NULL, 
                   method = method.opt, lower = -Inf, upper = Inf,
                   control = list(fnscale=-2.0, reltol=1e-12), 
                   hessian = FALSE, vc.data, nfixed, mod)
     va<-maxlik$par[nfixed+1]
     vq<-maxlik$par[nfixed+2]
     ve<-maxlik$par[nfixed+3]
     b.ve <-maxlik$par[nfixed+4]
  }
  fixed<-maxlik$par[1:nfixed]
  vt<-va+vq+ve
  cat("Number of families         = ", 
      length(unique(ped$famid[complete.cases(ped[,trait])])),"\n",
      "Number of observations     = ", nobs,"\n",
      "Trait mean (intercept)     = ", fixed[1],"\n",sep="")
  if (nfixed>1) {
    for(i in 2:(nfixed)) {
      cat(formatC(covariates[i-1],width=27,flag="-"),"= ",fixed[i],"\n",sep="")
    }
  }
  cat("\n")
  cat("Additive genetic variance  = ", va ," (",round(100*va/vt,1),"%)\n",sep="")
  if (mod=="aqe") {
    cat("QTL genetic variance       = ", vq ," (",round(100*vq/vt,1),"%)\n",sep="")
  }
  cat("Environmental variance     = ", ve ," (",round(100*ve/vt,1),"%)\n",
      "Model loglikelihood        = ", maxlik$value,"\n\n",sep="")
  if (mod=="aqe") {
    lod=(maxlik$value-base.ll)/log(10)
    lrts=2*(maxlik$value-base.ll)
    cat("Lod score testing VQ=0     = ", lod, 
        " (df=1, P=",0.5*pchisq(lrts,1,lower=FALSE),")\n\n",sep="")
  }
  cat("Moderator coefficient (E)  = ", b.ve,"\n\n",sep="")
  cat("Total function evaluations = ", maxlik$counts[1],"\n\n",sep="")
}
#
# function to evaluate the MVN likelihood
#
lik2 <- function(m, va, vq, ve, b.ev, vc.data) {
   res<-sum(unlist(lapply(vc.data, lik2.fam, m, va, vq, ve, b.ev)))
   res
}
lik2.fam <- function(vc.data, m, va, vq, ve, b.ev) {
  ln<-function(x) ifelse(x>0, log(x), 0)
  yhat<-vc.data$covariates %*% matrix(m, nc=1)
  mv<-vc.data$covariates[, ncol(vc.data$covariates)]
  ny<-nrow(yhat)
  S<-va*vc.data$rel +  
     vq*vc.data$ibd + 
     diag(ve,nrow(vc.data$rel)) + diag(b.ev*mv, ny)
  res <- -0.5* (ln(det(S)) + t(vc.data$y-yhat) %*% solve(S) %*% (vc.data$y-yhat))
  res
}
#
# and the version to allow calling of optim() function maximizer
#
optlik2 <- function(par, vc.data, nfixed, mod) {
  if (mod=="ae") {
    lik2(par[1:nfixed],par[nfixed+1], 0.0, par[nfixed+2], 
         par[nfixed+3], vc.data) 
  }else if (mod=="aqe") {
    lik2(par[1:nfixed],par[nfixed+1], par[nfixed+2],
        par[nfixed+3],par[nfixed+4], vc.data) 
  }
}
#
# Read in pedigree and ibd data from MERLIN 
#
ped <- read.table("merlin2R.rta", header=TRUE, na="x")

#
# Compare moderated and unmoderated linkage lods
#
# varcomp("dBP",ped,ibds=c("ibd1","ibd2","ibd3","ibd4","ibd5","ibd6"), 
#         covariates="env")
varcomp.GxE("dBP",ped,ibds=c("ibd1","ibd2","ibd3","ibd4","ibd5","ibd6"), 
            covariates="env")
quit(save="no")