# Simulate from IRTree models and fit them with lme4.
# Part 1: ratings from 1 to 5 generated by a three-node tree (take a side?
#   which side? how far?) whose node abilities are correlated. Fit one theta for
#   all nodes and one theta per node; compare with the truth.
# Part 2: a test with skipped items, where lower-ability respondents skip more.
#   Score the skips as wrong, as missing, or with a skip node, and see whose
#   abilities come out biased.
# nAGQ = 0 makes glmer() skip its slowest step, so the fits take seconds; delete
# ", nAGQ = 0" in your own R session for the full Laplace fit. Adapted from
# ben-domingue/252: c8/irtree_sim.R.
library(lme4)
set.seed(51)

## Part 1: a three-node tree
np <- 400                                       # respondents (try 2000 locally)
ni <- 10                                        # items
sds <- c(side = 0.8, direction = 1.7, extreme = 1.2)          # SDs of the node thetas
R <- matrix(c(1, 0.6, 0.3,  0.6, 1, 0.2,  0.3, 0.2, 1), 3)   # their correlations
theta <- matrix(rnorm(np * 3), np) %*% chol(R) %*% diag(sds)
b <- cbind(side = rnorm(ni, -1.2, 0.5), direction = rnorm(ni, -0.9, 0.8),
           extreme = rnorm(ni, 0.5, 0.5))       # node difficulties, item by item
x <- sapply(1:ni, function(i) {
  side <- rbinom(np, 1, plogis(theta[, 1] - b[i, 1]))
  up   <- rbinom(np, 1, plogis(theta[, 2] - b[i, 2]))
  ext  <- rbinom(np, 1, plogis(theta[, 3] - b[i, 3]))
  ifelse(side == 0, 3, ifelse(up == 1, 4 + ext, 2 - ext))  # walk down the tree
})

# Recode: one row per node reached.
long <- data.frame(id = rep(1:np, ni), item = rep(sprintf("i%02d", 1:ni), each = np), x = c(x))
sided <- long[long$x != 3, ]
tree <- rbind(data.frame(long,  node = "side",      resp = as.integer(long$x != 3)),
              data.frame(sided, node = "direction", resp = as.integer(sided$x > 3)),
              data.frame(sided, node = "extreme",   resp = as.integer(sided$x %in% c(1, 5))))
tree$node <- factor(tree$node, levels = c("side", "direction", "extreme"))

one <- glmer(resp ~ 0 + item:node + (1 | id), tree, binomial, nAGQ = 0)
per <- glmer(resp ~ 0 + item:node + (0 + node | id), tree, binomial, nAGQ = 0)

cat("Share of responses in each category:\n")
print(round(prop.table(table(x)), 2))
cat("\nAIC, one theta and a theta per node:", round(AIC(one)), round(AIC(per)), "\n")
vc <- VarCorr(per)$id
cat("\nSDs of the node thetas, true and estimated:\n")
print(round(rbind(true = sds, estimated = attr(vc, "stddev")), 2))
pairs <- rbind(c(2, 1), c(3, 1), c(3, 2))
cat("\nCorrelations between nodes, true and estimated:\n")
cm <- rbind(true = R[pairs], estimated = attr(vc, "correlation")[pairs])
colnames(cm) <- c("side-direction", "side-extreme", "direction-extreme")
print(round(cm, 2))
cat("\nHow well each model's estimates track the true content (direction) theta:",
    "\n  one theta:", round(cor(ranef(one)$id[, 1], theta[, 2]), 2),
    "\n  theta per node:", round(cor(ranef(per)$id[, "nodedirection"], theta[, 2]), 2), "\n")

## Part 2: skipped items
np2 <- 600                                      # respondents
ni2 <- 15                                       # right/wrong items
ability <- rnorm(np2)
skipper <- -0.6 * ability + sqrt(1 - 0.6^2) * rnorm(np2)  # lower ability, more skipping
skipped <- matrix(rbinom(np2 * ni2, 1, plogis(1.2 * skipper - 2)), np2)
correct <- matrix(rbinom(np2 * ni2, 1, plogis(outer(ability, seq(-2, 2, length = ni2), "-"))), np2)
d <- data.frame(id = rep(1:np2, ni2), item = rep(sprintf("q%02d", 1:ni2), each = np2),
                skip = c(skipped), y = c(correct))

as_wrong <- glmer(ifelse(skip == 1, 0, y) ~ 0 + item + (1 | id), d, binomial, nAGQ = 0)
as_missing <- glmer(y ~ 0 + item + (1 | id), d[d$skip == 0, ], binomial, nAGQ = 0)
answered <- d[d$skip == 0, ]
st <- rbind(data.frame(d, node = "answer", resp = 1 - d$skip),
            data.frame(answered, node = "correct", resp = answered$y))
st$node <- factor(st$node, levels = c("answer", "correct"))
skip_node <- glmer(resp ~ 0 + item:node + (0 + node | id), st, binomial, nAGQ = 0)

est <- function(m, col = 1) {
  r <- ranef(m)$id
  setNames(r[, col], rownames(r))[as.character(1:np2)]
}
theta_hat <- cbind(wrong = est(as_wrong), missing = est(as_missing),
                   "skip node" = est(skip_node, "nodecorrect"))
n_skip <- rowSums(skipped)
most <- n_skip >= quantile(n_skip, 0.9)         # the tenth who skipped most
cat("\nShare of items skipped:", round(mean(skipped), 2),
    "; the", sum(most), "respondents who skipped most skipped", round(mean(n_skip[most]), 1), "of", ni2, "\n")
cat("Mean error in estimated ability (estimate - truth):\n")
print(round(rbind("skipped most" = colMeans(theta_hat[most, ] - ability[most], na.rm = TRUE),
                  "everyone else" = colMeans(theta_hat[!most, ] - ability[!most], na.rm = TRUE)), 2))
cat("Estimated correlation of the answer and correct nodes:",
    round(attr(VarCorr(skip_node)$id, "correlation")[1, 2], 2), "(true 0.6)\n")