# Constrained ordination, concurrent ordination, partial
# Section 4.1. in the paper. 
# Note, that minor differences are to be expected, due to randomness.
set.seed(2025)
library(gllvm)

data("beetle")
beetleEnv <- beetle$X
beetle <- beetle$Y

head(beetle)
head(beetleEnv)

# scale and center the predictors
X <- scale(beetleEnv[,!names(beetleEnv) %in% c("SiteCode", "Landuse", "Grid", "Area", "Samplingyear")])
head(X)

# constrained ordination/rrr model + reproduces Fig. 1
ftConstOrd <- gllvm(y=beetle, X=X, family="negative.binomial", num.RR=2)
summary(ftConstOrd, rotate = TRUE)
par(mfrow=c(2,2))
coefplot(ftConstOrd, which.Xcoef = c("Reprobiom", "Elevation"), mar = c(4,4,2,1),
        mfrow=c(2,2), main="coefplot(ftConstOrd)", cex.lab=1, ind.spp=c(1:34))
ordiplot(ftConstOrd, symbols = TRUE, mar = c(4,4,1,1), main="ordiplot(ftConstOrd)", cex.env = 0.8, jitter = TRUE)
plot(summary(ftConstOrd), mar =c(4,6,2,1), cex.axis = .55, main="plot(summary(ftConstOrd))")
dev.off()

# concurrent ordination + reproduce Fig. 2:
ftConcOrd <- gllvm(y=beetle, X=X, family="ZINB", num.lv.c=2, n.init=5)
#ftConcOrd <- gllvm(y=beetle, X=X, family="negative.binomial", num.lv.c=2, starting.val="zero")
summary(ftConcOrd)
ordiplot(ftConcOrd, symbols = TRUE, ind.spp = 34, jitter=TRUE, biplot=TRUE, xlim=c(-4,4)) # draw only some of the species to reduce clutter

# partial ordination
ftPartOrd <- gllvm(y=beetle, X=X, family="ZINB", num.lv.c=2,
                      formula=~Canopyheight + Reprobiom, #starting.val="random",
                      lv.formula=~Texture + Org + pH + AvailP + AvailK + Moist
                      + Bare + Litter + Bryophyte + Plants.m2 + Stemdensity + Biom_l5
                      + Biom_m5 + Elevation + Management, randomB="P", n.init=5)

round(ftPartOrd$params$LvXcoef,5)
round(ftConcOrd$params$LvXcoef,5)
round(cbind(ftPartOrd$params$LvXcoef, ftConcOrd$params$LvXcoef[-c(11,15),]), digits=5)
# environmental correlation plot, reproduces Fig. 5
corrplot::corrplot(getEnvironCor(ftPartOrd),
                   type = "lower", diag = FALSE, order = "AOE",
                   tl.cex = 0.7, tl.srt = 45, tl.col = "black")

# model selection using information criteria:
lvcount <- c(2,3,4)
model <- c("poisson", "ZIP", "negative.binomial", "ZINB")
#model <- c("negative.binomial", "ZINB")

res <- data.frame(family = rep(model, each=3), num.lv.c = rep(lvcount, 4),
                  AIC = NA, AICc = NA, BIC = NA, logL = NA, df = NA)
idx = 1
for (m in model) {
  for (d in lvcount) {
    ft <- gllvm(y = beetle, X = X, family = m, num.lv.c = d, n.init=3)
    res$AIC[idx] <- AIC(ft)
    res$AICc[idx] <- AICc(ft)
    res$BIC[idx] <- BIC(ft)
    res$logL[idx] <- ft$logL
    res$df[idx] <- summary(ft)$df
    idx <- idx+1
  }
}
print(res)