# Distributional tests: sampler state distributions vs exact probabilities,
# and consistency of the distribution functions (IsingStateProb,
# IsingSumLikelihood, LinTransform).
library("IsingSampler")

set.seed(1234)

## Fixed 4-node model (0/1 coding) ------------------------------------------
N <- 4
graph <- matrix(0, N, N)
graph[upper.tri(graph)] <- c(0.5, -0.4, 0.3, 0.2, -0.1, 0.6)
graph <- graph + t(graph)
thresholds <- c(-0.2, 0.1, -0.5, 0.3)
beta <- 1
responses <- c(0L, 1L)

## Exact distribution over all 2^4 states:
lik <- IsingLikelihood(graph, thresholds, beta, responses)
stopifnot(nrow(lik) == 2^N)
stopifnot(abs(sum(lik$Probability) - 1) < 1e-10)

# Key identifying each state (columns 2:(N+1) of the IsingLikelihood output
# hold the state configurations):
stateKey <- function(states) apply(states, 1, paste, collapse = "")
exactKey <- stateKey(lik[, -1])
exactP <- lik$Probability

## 1. Agreement of the three samplers with the exact distribution -----------
# n = 5000 samples per method; empirical state frequencies must be within
# 0.05 (max absolute deviation) of the exact probabilities. Note that
# method = "MH" is approximate (finite chain, nIter = 200), hence the
# generous tolerance; "direct" and "CFTP" are exact samplers and only carry
# Monte Carlo error.
n <- 5000
maxDev <- function(sample) {
  freq <- table(factor(stateKey(sample), levels = exactKey)) / nrow(sample)
  max(abs(as.numeric(freq) - exactP))
}

devs <- c(
  direct = maxDev(IsingSampler(n, graph, thresholds, beta, responses = responses,
                               method = "direct")),
  CFTP = maxDev(IsingSampler(n, graph, thresholds, beta, responses = responses,
                             method = "CFTP")),
  MH = maxDev(IsingSampler(n, graph, thresholds, beta, nIter = 200,
                           responses = responses, method = "MH"))
)
print(round(devs, 4))
stopifnot(all(devs < 0.05))

## 2. IsingStateProb sums to 1 over the full state space --------------------
allStates <- as.matrix(do.call(expand.grid, replicate(N, responses,
                                                      simplify = FALSE)))
pAll <- IsingStateProb(lapply(seq_len(nrow(allStates)),
                              function(i) allStates[i, ]),
                       graph, thresholds, beta, responses)
stopifnot(length(pAll) == 2^N)
stopifnot(abs(sum(pAll) - 1) < 1e-10)

## 3. LinTransform round-trip invariance ------------------------------------
# Transform the 0/1 parameters to -1/1 coding; state s (0/1) corresponds to
# state 2*s - 1 (-1/1). Both codings must assign identical probabilities to
# corresponding states.
tr <- LinTransform(graph, thresholds, from = c(0L, 1L), to = c(-1L, 1L))
lik2 <- IsingLikelihood(tr$graph, tr$thresholds, beta, responses = c(-1L, 1L))

# expand.grid enumerates both codings in the same order and 0 -> -1, 1 -> 1
# preserves that order; verify the state mapping explicitly:
stopifnot(all(as.matrix(lik2[, -1]) == 2 * as.matrix(lik[, -1]) - 1))
stopifnot(isTRUE(all.equal(lik$Probability, lik2$Probability,
                           tolerance = 1e-12)))

## 4. IsingSumLikelihood probabilities sum to 1 ------------------------------
# IsingSumLikelihood returns a data frame with columns 'Sum' (sum score) and
# 'P' (probability of that sum score); the probabilities must sum to 1.
sl <- IsingSumLikelihood(graph, thresholds, beta, responses)
stopifnot(is.data.frame(sl), all(c("Sum", "P") %in% names(sl)))
stopifnot(nrow(sl) == N + 1) # sum scores 0:4 for 0/1 responses
stopifnot(all(sl$P >= 0))
stopifnot(abs(sum(sl$P) - 1) < 1e-10)

# The sum-score distribution must also match the aggregated exact state
# distribution:
sumExact <- tapply(exactP, rowSums(lik[, -1]), sum)
stopifnot(isTRUE(all.equal(unname(sl$P), unname(as.numeric(sumExact)),
                           tolerance = 1e-10)))

cat("test-distribution.R: all tests passed\n")
