# =============================================================
# Spy or Not? -- Random Role + Clue Draw
# =============================================================
# Draws a suspect's true role (Loyal/Spy) using the prior, then
# draws a clue conditional on that role using the likelihoods
# from the handout's clue table.
# =============================================================

set.seed(3257)  # remove/replace with a fixed seed for reproducibility

# ---- Priors ----
priors <- c(Loyal = 0.7, Spy = 0.3)

# ---- Conditional probabilities: P(Clue | Role) ----
# Rows = clues, columns = role
clue_probs <- matrix(
  c(0.6, 0.2,   # Clue A: "I was helping the team the whole time."
    0.3, 0.5,   # Clue B: "I don't really remember what happened."
    0.1, 0.3),  # Clue C: "Trust me, it wasn't me."
  nrow = 3, byrow = TRUE,
  dimnames = list(c("A", "B", "C"), c("Loyal", "Spy"))
)

clue_text <- c(
  A = "I was helping the team the whole time.",
  B = "I don't really remember what happened.",
  C = "Trust me, it wasn't me."
)

# ---- Function: draw one suspect's role + clue ----
draw_suspect <- function(priors, clue_probs, clue_text) {
  # Step 1: draw true role from the prior
  role <- sample(names(priors), size = 1, prob = priors)

  # Step 2: draw a clue conditional on that role
  # (clue_probs[, role] gives P(Clue | role) for each clue)
  clue_weights <- clue_probs[, role]
  clue_id <- sample(rownames(clue_probs), size = 1, prob = clue_weights)

  list(
    role = role,
    clue_id = clue_id,
    clue_text = clue_text[[clue_id]]
  )
}

# ---- Run a single draw ----
result <- draw_suspect(priors, clue_probs, clue_text)

cat("Clue given to class: \"", result$clue_text, "\" (Clue ", result$clue_id, ")\n", sep = "")
cat("(True role -- keep hidden until reveal: ", result$role, ")\n", sep = "")

# ---- Optional: simulate many suspects at once (e.g. one per group) ----
n_suspects <- 5
suspects <- replicate(n_suspects, draw_suspect(priors, clue_probs, clue_text), simplify = FALSE)

cat("\n--- Batch draw for", n_suspects, "suspects ---\n")
for (i in seq_along(suspects)) {
  s <- suspects[[i]]
  cat(sprintf("Suspect %d | Clue %s: \"%s\"  [role: %s]\n",
              i, s$clue_id, s$clue_text, s$role))
}

# ---- Optional: verify empirical draw frequencies match theory ----
# Run many trials to sanity-check the sampling matches the priors/likelihoods
n_trials <- 10000
sim <- replicate(n_trials, draw_suspect(priors, clue_probs, clue_text), simplify = FALSE)
roles_drawn <- sapply(sim, function(x) x$role)
clues_drawn <- sapply(sim, function(x) x$clue_id)

cat("\n--- Sanity check over", n_trials, "trials ---\n")
cat("Empirical role frequencies (should approx. match priors 0.7/0.3):\n")
print(round(prop.table(table(roles_drawn)), 3))

cat("\nEmpirical P(Clue) marginal frequencies:\n")
print(round(prop.table(table(clues_drawn)), 3))

cat("\nEmpirical P(Role | Clue) -- compare to the answer key's worked posteriors:\n")
print(round(prop.table(table(clues_drawn, roles_drawn), margin = 1), 3))

############################
# ---- Function: simulate one suspect across n_rounds ----
simulate_suspect <- function(suspect_id, true_role, n_rounds,
                              prior_spy, clue_probs) {
  post_spy <- prior_spy       # running posterior, starts at the prior
  clues_drawn <- character(n_rounds)
  posterior_trace <- numeric(n_rounds + 1)
  posterior_trace[1] <- prior_spy   # round 0 = the prior itself

  for (r in seq_len(n_rounds)) {
    clue_id <- draw_clue(true_role, clue_probs)
    post_spy <- bayes_update(post_spy, clue_id, clue_probs)
    clues_drawn[r] <- clue_id
    posterior_trace[r + 1] <- post_spy
  }

  data.frame(
    suspect = suspect_id,
    true_role = true_role,
    round = 0:n_rounds,
    clue = c(NA, clues_drawn),
    posterior_spy = posterior_trace
  )
}

# ---- Simulate several suspects ----
n_suspects <- 6
n_rounds   <- 5

# Randomly assign true roles to suspects using the prior
true_roles <- sample(c("Loyal", "Spy"), size = n_suspects,
                      replace = TRUE, prob = c(prior_loyal, prior_spy))

results_list <- lapply(seq_len(n_suspects), function(i) {
  simulate_suspect(
    suspect_id = paste0("Suspect ", i),
    true_role  = true_roles[i],
    n_rounds   = n_rounds,
    prior_spy  = prior_spy,
    clue_probs = clue_probs
  )
})

results <- do.call(rbind, results_list)

# ---- Print the clue sequence + posterior trace per suspect ----
cat("=== Clue sequences and posterior trajectories ===\n\n")
for (i in seq_len(n_suspects)) {
  sub <- results[results$suspect == paste0("Suspect ", i), ]
  cat(sprintf("Suspect %d [TRUE ROLE: %s]\n", i, true_roles[i]))
  cat("  Clues:      ", paste(na.omit(sub$clue), collapse = " -> "), "\n")
  cat("  P(Spy) path:", paste(round(sub$posterior_spy, 3), collapse = " -> "), "\n\n")
}

# ---- Plot: posterior trajectory per suspect, faceted ----
results$suspect <- factor(results$suspect,
                           levels = paste0("Suspect ", seq_len(n_suspects)))

p_facet <- ggplot(results, aes(x = round, y = posterior_spy, color = true_role)) +
  geom_hline(yintercept = prior_spy, linetype = "dashed", color = "gray50") +
  geom_line(linewidth = 1) +
  geom_point(size = 2) +
  facet_wrap(~ suspect, ncol = 3) +
  scale_color_manual(values = c(Loyal = "#1D6F9F", Spy = "#A63A50"),
                      name = "True Role") +
  scale_y_continuous(limits = c(0, 1)) +
  labs(
    title = "Sequential Bayesian Updating: P(Spy | Clues) Over Rounds",
    subtitle = "Dashed line = prior P(Spy) = 0.3",
    x = "Round (number of clues observed)",
    y = expression(P(Spy ~ "|" ~ clue[1] * "," ~ ... * "," ~ clue[n]))
  ) +
  theme_minimal(base_size = 13) +
  theme(legend.position = "bottom",
        plot.title = element_text(face = "bold"))

print(p_facet)
#ggsave("spy_posterior_facets.png", p_facet, width = 10, height = 6, dpi = 300)