# Tests for the 'constrain' argument of IsingSampler() (GitHub issue #1):
# constraints must be honored by all three sampling methods ("MH", "CFTP",
# "direct"), and for "direct" the constrained samples must follow the
# conditional distribution given the clamped nodes.
# Plain stopifnot tests, no package dependencies beyond IsingSampler itself.

library("IsingSampler")

set.seed(1234)

## Small test model -----------------------------------------------------------
N <- 4
graph <- matrix(0, N, N)
graph[upper.tri(graph)] <- runif(N * (N - 1) / 2, -0.5, 1)
graph <- graph + t(graph)
thresholds <- runif(N, -1, 1)
responses <- c(0L, 1L)

## 1) Clamping node 1 to each response option, all methods --------------------
nSamp <- 50
for (method in c("MH", "CFTP", "direct")) {
  for (val in responses) {
    constrain <- matrix(NA_real_, nSamp, N)
    constrain[, 1] <- val
    samples <- IsingSampler(nSamp, graph, thresholds, method = method,
                            responses = responses, constrain = constrain)
    samples <- as.matrix(samples)
    stopifnot(nrow(samples) == nSamp, ncol(samples) == N)
    # The constrained column must equal the clamped value in every sample:
    stopifnot(all(samples[, 1] == val))
    # Unconstrained columns must still contain valid response options:
    stopifnot(all(as.vector(samples[, -1]) %in% responses))
  }
}

## 2) Mixed constraints: different rows constrained differently (direct) ------
constrain <- matrix(NA_real_, 60, N)
constrain[1:20, 1] <- 0   # rows 1-20: node 1 clamped to 0
constrain[21:40, 1] <- 1  # rows 21-40: node 1 clamped to 1
constrain[21:40, 3] <- 0  # rows 21-40: node 3 also clamped to 0
                          # rows 41-60: unconstrained
samples <- as.matrix(IsingSampler(60, graph, thresholds, method = "direct",
                                  constrain = constrain))
stopifnot(all(samples[1:20, 1] == 0))
stopifnot(all(samples[21:40, 1] == 1))
stopifnot(all(samples[21:40, 3] == 0))
stopifnot(all(as.vector(samples) %in% responses))

## 3) Conditional distribution is respected for method = "direct" -------------
# 3-node model: compare empirical state frequencies of constrained direct
# samples against the exact conditional probabilities obtained by restricting
# IsingLikelihood() output to the constrained subset.
set.seed(5678)
N3 <- 3
graph3 <- matrix(0, N3, N3)
graph3[upper.tri(graph3)] <- c(0.5, -0.3, 0.8)
graph3 <- graph3 + t(graph3)
thresholds3 <- c(-0.2, 0.1, 0.3)

lik <- IsingLikelihood(graph3, thresholds3, beta = 1, responses = c(0L, 1L))
# Restrict to states with node 1 == 1 and renormalize:
sub <- lik[lik[, 2] == 1, ]
condP <- sub$Probability / sum(sub$Probability)
condKey <- apply(sub[, -1], 1, paste, collapse = ",")

nBig <- 20000
constrain3 <- matrix(NA_real_, nBig, N3)
constrain3[, 1] <- 1
samples3 <- as.matrix(IsingSampler(nBig, graph3, thresholds3, method = "direct",
                                   constrain = constrain3))
stopifnot(all(samples3[, 1] == 1))
empKey <- apply(samples3, 1, paste, collapse = ",")
empFreq <- table(factor(empKey, levels = condKey)) / nBig
# All conditional states should be covered and frequencies close to theory:
stopifnot(max(abs(as.numeric(empFreq) - condP)) < 0.02)

## 4) Validation of the 'constrain' argument ----------------------------------
# Wrong dimensions must give an error, for every method:
for (method in c("MH", "CFTP", "direct")) {
  res <- try(IsingSampler(10, graph, thresholds, method = method,
                          constrain = matrix(NA_real_, 5, N)), silent = TRUE)
  stopifnot(inherits(res, "try-error"))
}
# Values not in 'responses' must give an error:
badConstrain <- matrix(NA_real_, 10, N)
badConstrain[, 1] <- 2
res <- try(IsingSampler(10, graph, thresholds, constrain = badConstrain),
           silent = TRUE)
stopifnot(inherits(res, "try-error"))

cat("All constrain tests passed.\n")
