library(MCMCglmm)
library(coda)
library(reshape)

source("owls_R_funs.R")
load("../DATA/Owls.rda")

## MCMCglmm approach
Owls <- rename(Owls,c(SiblingNegotiation="NCalls"))
Owls$ArrivalTime <- scale(Owls$ArrivalTime,center=TRUE,scale=FALSE)

offvec <- c(1,1,2,rep(1,5)) ## indicates whether fixed effect is the offset
                            ## (log brood size) or not
offvec2 <- c(1,1,2,rep(1,5),rep(3,54)) ## indicates whether parameter is fixed effect (1), offset (2), random effect (3)

fixef2 <- NCalls~trait-1+
  at.level(trait,1):logBroodSize+

  at.level(trait,1):((FoodTreatment+ArrivalTime)*SexParent)

prior_overdisp  <- list(R=list(V=diag(c(1,1)),nu=0.002,fix=2),
                        G=list(list(V=diag(c(1,1e-6)),nu=0.002,fix=2)))

prior_overdisp_broodoff <- within(prior_overdisp,
                                  { B <- list(mu=c(0,1)[offvec],
                                              V=diag(c(1e8,1e-6)[offvec]))})


st1 <- system.time(ofit_MCMCglmm <-
                   MCMCglmm(fixef2,
                            rcov=~idh(trait):units,
                            random=~idh(trait):Nest,
                            prior=prior_overdisp_broodoff,
                            pr=TRUE,  ## save random effects
                            data=Owls,
                            family="zipoisson",
                            verbose=FALSE))

st1 <- unname(st1["elapsed"])

## should provide estimates for variance terms as well ...
ofit_MCMCglmm_mcmc <- as.mcmc(ofit_MCMCglmm$Sol)[,offvec2==1]
ofit_ranef <- as.mcmc(ofit_MCMCglmm$Sol)[,offvec2==3]
ofit_vcov <- ofit_MCMCglmm$VCV[,c(1,3)]

p_order <- c(1,3:7,2) ## fixed effects first, then logit.pz
fixef <- unname(colMeans(ofit_MCMCglmm_mcmc)[p_order])

gdiag <- try(geweke.diag(ofit_MCMCglmm_mcmc))
if (inherits(gdiag,"try-error")) gdiag <- NA
results1 <- list(obj=NULL,
                coef=c(fixef,colMeans(ofit_vcov)[1]),
                rcoef=colMeans(ofit_ranef),
                vcoef=colMeans(ofit_vcov),
                 ## include fixed effects, logit p, and sigma (regular scale)
                confint.mcmc=rbind(HPDinterval(ofit_MCMCglmm_mcmc)[p_order,],
                  HPDinterval(ofit_vcov)[1,]),
                confint.quantile=rbind(t(apply(ofit_MCMCglmm_mcmc,2,
                  quantile,c(0.025,0.975)))[p_order,],
                  quantile(ofit_vcov[,1],c(0.025,0.975))),
                rconfint.mcmc=HPDinterval(ofit_ranef),
                vconfint.mcmc=HPDinterval(ofit_vcov),
                time=c(fit=st1,mcmc=st1),
                terminfo=NULL,
                convinfo=list(effsize=effectiveSize(ofit_MCMCglmm_mcmc),
                  geweke=gdiag))

fits <- list(MCMCglmm=ofit_MCMCglmm)
save("fits",file="owl_R_fits.RData")

###########################
## GLMMADMB
###########################
library(glmmADMB)
st2 <- system.time(ofit_glmmadmb <-
                   glmmadmb(NCalls~(FoodTreatment+ArrivalTime)*SexParent+
                            offset(logBroodSize),
                            random=~1|Nest,
                            data=Owls,
                            zeroInflation=TRUE,
                            family="poisson"))
## "est covariance matrix may not be positive definite" warning

st2 <- unname(st2["elapsed"])

fits <- c(fits,list(glmmadmb=ofit_glmmadmb))
save("fits",file="owl_R_fits.RData")

rvar <- ofit_glmmadmb$S$Nest

results2 <- list(obj=-logLik(ofit_glmmadmb),
                coef=c(coef(ofit_glmmadmb),qlogis(ofit_glmmadmb$pz),rvar),
                rcoef=ranef(ofit_glmmadmb),
                stdEr=stdEr(ofit_glmmadmb),
                confint.quad=rbind(confint(ofit_glmmadmb),
                  with(ofit_glmmadmb,qlogis(pz + c(-1.96,1.96)*sd_pz)),
                  with(ofit_glmmadmb,rvar + c(-1.96,1.96)*sd_S$Nest)),
                 time=c(fit=st2))

library(lme4)

## test compatibility:
## st3 <- system.time(tmpcl2 <- zipme.f(10,Owls))
## st4 <- system.time(tmpcl3 <- zipme(cformula=NCalls~(FoodTreatment+ArrivalTime)*SexParent+
##                                    offset(logBroodSize)+(1|Nest),
##                                    zformula=z ~ SexParent*FoodTreatment + SexParent*ArrivalTime + (1|Nest),
##                                    data=Owls,maxitr=20,tol=1e-6))

st3 <- system.time(ofit_zipme <- zipme(cformula=NCalls~(FoodTreatment+ArrivalTime)*SexParent+
                                   offset(logBroodSize)+(1|Nest),
                                   zformula=z ~ 1,
                                   data=Owls,maxitr=20,tol=1e-6))

st3 <- unname(st3["elapsed"])

tmpci <- function(m) {
  se <- sqrt(diag(vcov(m)))
  fixef(m) +  se %o% c(-1.96,1.96)
}
results3 <- list(obj=NA, ## not sure how to compute this
                 coef=c(fixef(ofit_zipme$cfit),coef(ofit_zipme$zfit),
                   c(VarCorr(ofit_zipme$cfit)$Nest)),
                 rcoef=ranef(ofit_zipme$cfit),
                 confint.quad=rbind(tmpci(ofit_zipme$cfit),
                   confint.default(ofit_zipme$zfit),
                   c(NA,NA)),
                 time=c(fit=st3))

fits <- c(fits,list(zipme=ofit_zipme))
save("fits",file="owl_R_fits.RData")

results <- list(MCMCglmm=results1,glmmadmb=results2,zipme=results3)
save("results",file="fit.RData")