# ============================================================
# Analysis: Shaded Lights Experiment — 3-Light Pilot
# ============================================================
# Sources Structural_3lights.R   → DOfile_realexp_3lights.R
#         Structural_3lights (2).R → DOfile_realexp_3lights (2).R
# Stacks them vertically with continuous subject numbering.
# Output: Data_3L — one row per subject × treatment × guess
# ============================================================
#
# ------------------------------------------------------------
# DATA_3L — VARIABLE CODEBOOK
# ------------------------------------------------------------
#
# UNIT OF OBSERVATION
#   One row per subject × treatment × guess (trial).
#   Each subject makes 16 guesses per treatment across 5 treatments
#   → 80 rows per subject.
#
# ── IDENTIFIERS ─────────────────────────────────────────────
#   subject_id      : Integer. Sequential anonymous ID.
#   prolific_id     : Character. Anonymised Prolific ID (24-char hex).
#   session         : Factor. "3L(1)" or "3L(2)".
#
# ── TREATMENT ───────────────────────────────────────────────
#   treatment       : Character. Logical rule governing the machine.
#                     Five levels: AND, OR, INHIBIT, EITHER, JOINT.
#                     The rule depends only on Red and Blue lights;
#                     Green light is irrelevant to all rules.
#   round_order     : Integer (1–5). Position of this treatment in the
#                     subject's randomised sequence.
#
# ── TRIAL-LEVEL VARIABLES ───────────────────────────────────
#   Guess_Number    : Integer (1–16). Index within the treatment.
#   Light_Config    : Character. Observable light state formatted as
#                     "(red,blue,green)"; 8 possible values.
#   red_light       : Integer (0/1). State of the red light.
#   blue_light      : Integer (0/1). State of the blue light.
#   green_light     : Integer (0/1). State of the green light.
#                     Green does NOT affect Machine_CorrectP.
#   Guess           : Integer (0/1). Subject's prediction.
#   Machine_CorrectP: Integer (0/1). Optimal prediction given the rule
#                     (based on red and blue only).
#
# ── TREATMENT-LEVEL VARIABLES ───────────────────────────────
#   time_machine    : Numeric. Seconds on the prediction page.
#   Notes           : Character. Free-text notes.
#   certainty       : Numeric. Self-reported confidence (0–100).
#   difficulty      : Numeric. Self-reported difficulty (1–10).
#   difficulty_certainty : Numeric. Combined rating (0–100).
#   predicted_correct_self : Numeric. Subject's own accuracy estimate.
#   max_correct     : Integer (0–16). Correct guesses in treatment.
#   trueextracted   : Integer (0/1). 1 if subject demonstrably learned
#                     the true rule: binom.test(max_correct, 16, 0.5)
#                     p-value ≤ 0.01 AND max_correct > 8.
#
# ── SUBJECT-LEVEL VARIABLES ─────────────────────────────────
#   STRATEGY, COMMENTS, num_wrong, final_payment,
#   payment_machine, education, age, time_total_min,
#   original_color  : same definitions as DOfile_realexp_3lights.R.
#
# ── RULE-VALUE VARIABLES ────────────────────────────────────
#   val_<rule>      : Numeric. Proportion of trials in this treatment's
#                     dataset that <rule> would predict correctly.
#                     Rules: all named single- and two-variable rules
#                     from Structural_3lights.R (38 strategies).
#
# ── PER-CONFIG ACCURACY OF THE TRUE RULE ────────────────────
#   ACC_000 … ACC_111 : Numeric. Accuracy of the true rule on each of
#                       the 8 Light_Config values within this treatment.
#
# ── LIGHT-CONFIG FREQUENCIES ────────────────────────────────
#   Frequency_000 … Frequency_111 : Numeric. Proportion of trials in
#                   the treatment dataset at each Light_Config.
#
# ── STRUCTURAL (MIXTURE-MODEL) POSTERIORS ───────────────────
#   post_<rule>     : Numeric. Bayesian posterior probability that this
#                     subject × treatment is generated by <rule>.
#                     From Structural_3lights.R mixture model.
#   Rule_used       : Character. Rule with highest posterior.
#   posterior       : Numeric (2 d.p.). Posterior of Rule_used.
#
# ── ANALYSIS VARIABLES ──────────────────────────────────────
#   correlationhigh : Integer (0/1).
#       1 – DIFFICULT (single-variable heuristics most competitive):
#             3L(1): AND, INHIBIT
#             3L(2): JOINT, OR, EITHER
#       0 – EASY:
#             3L(1): OR, EITHER, JOINT
#             3L(2): AND, INHIBIT
# ------------------------------------------------------------

library(tidyverse)


# ---- 1. Source structural scripts (runs DOfile then stratEst estimation) ------
# Each script is sourced into its own environment to prevent name collisions.

script_dir <- dirname(rstudioapi::getActiveDocumentContext()$path)

#env_3L1 <- new.env(parent = globalenv())
#source(file.path(script_dir, "Structural_modelselection.R"), local = env_3L1)

env_3L1 <- new.env(parent = globalenv())
source(file.path(script_dir, "Structural_3lights.R"), local = env_3L1)

#env_3L2 <- new.env(parent = globalenv())
#source(file.path(script_dir, "Structural_modelselection(2).R"), local = env_3L2)

env_3L2 <- new.env(parent = globalenv())
source(file.path(script_dir, "Structural_3lights (2).R"), local = env_3L2)

df_3L1 <- env_3L1$df_final_post   # session 1
df_3L2 <- env_3L2$df_final_post   # session 2


# ---- 2. Add session identifiers
df_3L1 <- df_3L1 %>% mutate(session = "3L(1)")
df_3L2 <- df_3L2 %>% mutate(session = "3L(2)")


# ---- 3. Continue subject numbering across sessions
# 3L(1) subjects keep their original IDs (1 … N₁).
# 3L(2) subject IDs are offset by the number of subjects in 3L(1)
# so that numbering is continuous and never restarts from 1.
n_subjects_3L1 <- n_distinct(df_3L1$subject_id)
df_3L2 <- df_3L2 %>%
  mutate(subject_id = subject_id + n_subjects_3L1)


# ---- 4. Stack vertically → Data_3L
Data_3L <- bind_rows(df_3L1, df_3L2) %>%
  mutate(session = factor(session, levels = c("3L(1)", "3L(2)"))) %>%
  arrange(session, subject_id, round_order, Guess_Number)


# ---- 5. correlationhigh variable
# correlationhigh = 1 for session × treatment combinations labelled DIFFICULT
# in __init__.py (single-variable heuristics most competitive with true rule):
#   3L(1): AND (freq2), INHIBIT (freq15)
#   3L(2): JOINT (freq9_A), OR (freq5), EITHER (freq7)
# correlationhigh = 0 for EASY treatments:
#   3L(1): OR (freq3), EITHER (freq6), JOINT (freq9)
#   3L(2): AND (freq1), INHIBIT (freq12)

Data_3L <- Data_3L |>
  mutate(correlationhigh = as.integer(
    (session == "3L(1)" & treatment %in% c("AND", "INHIBIT")) |
    (session == "3L(2)" & treatment %in% c("JOINT", "OR", "EITHER"))
  ))


# ---- Shared labels and colours (used throughout) 

corr_labels  <- c("0" = "Low correlation",  "1" = "High correlation")
corr_colours <- c("0" = "#2166ac", "1" = "#d6604d")
corr_fills   <- c("0" = "#2166ac", "1" = "#d6604d")

treatment_levels <- c("AND", "OR", "INHIBIT", "EITHER", "JOINT")


# ---- 2. Subject_acc: guess-weighted accuracy by light-config frequency -------
# For each guess the contribution is:
#   ACC_lightconfig       if Guess == Machine_CorrectP (subject chose optimally)
#   1 - ACC_lightconfig   otherwise
# Subject_acc is the mean over the 16 guesses within a (subject, treatment).

# Reset graphics device to clear any invalid state from sourced scripts
tryCatch(grDevices::dev.off(), error = function(e) invisible(NULL))

subject_stats <- Data_3L |>
  mutate(
    acc_val = case_when(
      Light_Config == "(0,0,0)" ~ ACC_000,
      Light_Config == "(0,0,1)" ~ ACC_001,
      Light_Config == "(0,1,0)" ~ ACC_010,
      Light_Config == "(0,1,1)" ~ ACC_011,
      Light_Config == "(1,0,0)" ~ ACC_100,
      Light_Config == "(1,0,1)" ~ ACC_101,
      Light_Config == "(1,1,0)" ~ ACC_110,
      Light_Config == "(1,1,1)" ~ ACC_111
    ),
    guess_acc = if_else(Guess == Machine_CorrectP, acc_val, 1 - acc_val)
  ) |>
  group_by(subject_id, treatment, correlationhigh) |>
  summarise(
    Subject_acc = mean(guess_acc, na.rm = TRUE),
    max_correct = first(max_correct),
    .groups = "drop"
  )


# ---- 3. CDF plots by statistical properties (correlationhigh): Aggregate ----

p_max <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh)) |>
  ggplot(aes(x = max_correct, colour = correlationhigh)) +
  stat_ecdf(linewidth = 0.8) +
  scale_colour_manual(values = corr_colours, labels = corr_labels) +
  labs(
    x      = "Number of correct guesses (max_correct)",
    y      = "Cumulative proportion",
    colour = NULL
  ) +
  theme_bw() +
  theme(legend.position = "bottom")

print(p_max)


# ---- 8. CDF plots by statistical properties (correlationhigh): By treatment --

p_max_bytreat <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh),
         treatment = factor(treatment, levels = treatment_levels)) |>
  ggplot(aes(x = max_correct, colour = correlationhigh)) +
  stat_ecdf(linewidth = 0.8) +
  scale_colour_manual(values = corr_colours, labels = corr_labels) +
  facet_wrap(~ treatment, nrow = 1) +
  labs(
    x      = "Number of correct guesses (max_correct)",
    y      = "Cumulative proportion",
    colour = NULL
  ) +
  theme_bw() +
  theme(legend.position = "bottom")

print(p_max_bytreat)


# ---- 4. Bar plots: fraction of subjects per max_correct bin -----------------

bar_data <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh)) |>
  group_by(correlationhigh, max_correct) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup()

means_data <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh)) |>
  group_by(correlationhigh) |>
  summarise(mean_mc = mean(max_correct, na.rm = TRUE), .groups = "drop")

p_bar <- ggplot(bar_data, aes(x = max_correct, y = frac, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", width = 0.7, alpha = 0.85) +
  geom_vline(data = means_data,
             aes(xintercept = mean_mc, colour = correlationhigh),
             linewidth = 0.9, linetype = "dashed") +
  scale_fill_manual(values  = corr_fills,  labels = corr_labels) +
  scale_colour_manual(values = corr_colours, labels = corr_labels, guide = "none") +
  scale_x_continuous(breaks = 0:16) +
  labs(
    x    = "Number of correct guesses (max_correct)",
    y    = "Fraction of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(legend.position = "bottom")

print(p_bar)

# By treatment — 10 panels: 5 treatments × 2 correlationhigh
bar_data_treat <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh),
         treatment = factor(treatment, levels = treatment_levels)) |>
  group_by(treatment, correlationhigh, max_correct) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(treatment, correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup()

means_data_treat <- subject_stats |>
  mutate(correlationhigh = factor(correlationhigh),
         treatment = factor(treatment, levels = treatment_levels)) |>
  group_by(treatment, correlationhigh) |>
  summarise(mean_mc = mean(max_correct, na.rm = TRUE), .groups = "drop")

p_bar_treat <- ggplot(bar_data_treat,
                      aes(x = max_correct, y = frac, fill = correlationhigh)) +
  geom_bar(stat = "identity", width = 0.8, alpha = 0.85) +
  geom_vline(data = means_data_treat,
             aes(xintercept = mean_mc, colour = correlationhigh),
             linewidth = 0.9, linetype = "dashed") +
  scale_fill_manual(values  = corr_fills,  labels = corr_labels) +
  scale_colour_manual(values = corr_colours, labels = corr_labels, guide = "none") +
  scale_x_continuous(breaks = seq(0, 16, 4)) +
  facet_grid(correlationhigh ~ treatment,
             labeller = labeller(correlationhigh = corr_labels)) +
  labs(
    x    = "Number of correct guesses (max_correct)",
    y    = "Fraction of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(legend.position = "none")

print(p_bar_treat)


# ---- 10. Bar plots: fraction of subjects by trueextracted × correlationhigh --

te_stats <- Data_3L |>
  distinct(subject_id, treatment, correlationhigh, trueextracted)

bar_data_te <- te_stats |>
  mutate(correlationhigh = factor(correlationhigh)) |>
  group_by(correlationhigh, trueextracted) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup()

means_data_te <- te_stats |>
  mutate(correlationhigh = factor(correlationhigh)) |>
  group_by(correlationhigh) |>
  summarise(mean_te = mean(trueextracted, na.rm = TRUE), .groups = "drop")

p_bar_te <- ggplot(bar_data_te, aes(x = trueextracted, y = frac, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", width = 0.5, alpha = 0.85) +
  geom_vline(data = means_data_te,
             aes(xintercept = mean_te, colour = correlationhigh),
             linewidth = 0.9, linetype = "dashed") +
  scale_fill_manual(values  = corr_fills,  labels = corr_labels) +
  scale_colour_manual(values = corr_colours, labels = corr_labels, guide = "none") +
  scale_x_continuous(breaks = 0:1, labels = c("Not extracted", "Extracted")) +
  labs(
    x    = "True rule extracted (trueextracted)",
    y    = "Fraction of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(legend.position = "bottom")

print(p_bar_te)

# By treatment
bar_data_te_treat <- te_stats |>
  mutate(correlationhigh = factor(correlationhigh),
         treatment = factor(treatment, levels = treatment_levels)) |>
  group_by(treatment, correlationhigh, trueextracted) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(treatment, correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup()

means_data_te_treat <- te_stats |>
  mutate(correlationhigh = factor(correlationhigh),
         treatment = factor(treatment, levels = treatment_levels)) |>
  group_by(treatment, correlationhigh) |>
  summarise(mean_te = mean(trueextracted, na.rm = TRUE), .groups = "drop")

p_bar_te_treat <- ggplot(bar_data_te_treat,
                         aes(x = trueextracted, y = frac, fill = correlationhigh)) +
  geom_bar(stat = "identity", width = 0.5, alpha = 0.85) +
  geom_vline(data = means_data_te_treat,
             aes(xintercept = mean_te, colour = correlationhigh),
             linewidth = 0.9, linetype = "dashed") +
  scale_fill_manual(values  = corr_fills,  labels = corr_labels) +
  scale_colour_manual(values = corr_colours, labels = corr_labels, guide = "none") +
  scale_x_continuous(breaks = 0:1, labels = c("Not extracted", "Extracted")) +
  facet_grid(correlationhigh ~ treatment,
             labeller = labeller(correlationhigh = corr_labels)) +
  labs(
    x    = "True rule extracted (trueextracted)",
    y    = "Fraction of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(legend.position = "none")

print(p_bar_te_treat)


# ---- 5. Suboptimal extraction and randomizing ------------------------------

# 11a. All 256 deterministic rules mapped to Light_Config values.
# expand.grid with p000 as LSB: row i+1 corresponds to integer i.
rule_table_3L <- expand.grid(
  p000 = 0:1, p001 = 0:1, p010 = 0:1, p011 = 0:1,
  p100 = 0:1, p101 = 0:1, p110 = 0:1, p111 = 0:1
) |>
  as_tibble() |>
  mutate(rule_id = row_number()) |>
  pivot_longer(-rule_id, names_to = "config_key", values_to = "rule_pred") |>
  mutate(Light_Config = case_when(
    config_key == "p000" ~ "(0,0,0)",
    config_key == "p001" ~ "(0,0,1)",
    config_key == "p010" ~ "(0,1,0)",
    config_key == "p011" ~ "(0,1,1)",
    config_key == "p100" ~ "(1,0,0)",
    config_key == "p101" ~ "(1,0,1)",
    config_key == "p110" ~ "(1,1,0)",
    config_key == "p111" ~ "(1,1,1)"
  )) |>
  select(-config_key)

# 11b. Find which rule_id corresponds to the true rule for each treatment
true_rule_ids <- Data_3L |>
  distinct(treatment, Light_Config, Machine_CorrectP) |>
  left_join(rule_table_3L, by = "Light_Config", relationship = "many-to-many") |>
  group_by(treatment, rule_id) |>
  summarise(is_true = all(Machine_CorrectP == rule_pred), .groups = "drop") |>
  filter(is_true) |>
  select(treatment, true_rule_id = rule_id)

# 11c. Count matches against all 255 non-true rules per subject × treatment.
# suboptimalextracted = 1 if any non-true rule passes binom.test
# (p <= 0.01, count > 8; 16 guesses, strictly above chance).
suboptimalextracted_df <- Data_3L |>
  select(subject_id, treatment, correlationhigh, Light_Config, Guess) |>
  left_join(rule_table_3L, by = "Light_Config", relationship = "many-to-many") |>
  left_join(true_rule_ids, by = "treatment") |>
  filter(rule_id != true_rule_id) |>
  group_by(subject_id, treatment, correlationhigh, rule_id) |>
  summarise(n_match = sum(Guess == rule_pred), .groups = "drop") |>
  rowwise() |>
  mutate(passes = binom.test(n_match, 16, 0.5)$p.value <= 0.01 & n_match > 8) |>
  ungroup() |>
  group_by(subject_id, treatment, correlationhigh) |>
  summarise(suboptimalextracted = as.integer(any(passes)), .groups = "drop")

# 11d. Combine with trueextracted and define randomizing
extraction_stats <- Data_3L |>
  distinct(subject_id, treatment, correlationhigh, trueextracted) |>
  left_join(suboptimalextracted_df,
            by = c("subject_id", "treatment", "correlationhigh")) |>
  mutate(randomizing = as.integer(trueextracted == 0 & suboptimalextracted == 0))

# 11e. Frequency table by correlationhigh
tbl_corr <- extraction_stats |>
  group_by(correlationhigh) |>
  summarise(
    N        = n(),
    true_ext = sum(trueextracted,       na.rm = TRUE),
    subopt   = sum(suboptimalextracted, na.rm = TRUE),
    rand     = sum(randomizing,         na.rm = TRUE),
    .groups = "drop"
  ) |>
  mutate(
    true_ext_pct = round(true_ext / N * 100, 1),
    subopt_pct   = round(subopt   / N * 100, 1),
    rand_pct     = round(rand     / N * 100, 1)
  )

print(tbl_corr)

# 11f. Frequency table by treatment
tbl_treat <- extraction_stats |>
  group_by(treatment) |>
  summarise(
    N        = n(),
    true_ext = sum(trueextracted,       na.rm = TRUE),
    subopt   = sum(suboptimalextracted, na.rm = TRUE),
    rand     = sum(randomizing,         na.rm = TRUE),
    .groups = "drop"
  ) |>
  mutate(
    true_ext_pct = round(true_ext / N * 100, 1),
    subopt_pct   = round(subopt   / N * 100, 1),
    rand_pct     = round(rand     / N * 100, 1)
  )

print(tbl_treat)

# 11g. Frequency table by treatment × correlationhigh
tbl_treat_corr <- extraction_stats |>
  group_by(treatment, correlationhigh) |>
  summarise(
    N        = n(),
    true_ext = sum(trueextracted,       na.rm = TRUE),
    subopt   = sum(suboptimalextracted, na.rm = TRUE),
    rand     = sum(randomizing,         na.rm = TRUE),
    .groups = "drop"
  ) |>
  mutate(
    true_ext_pct = round(true_ext / N * 100, 1),
    subopt_pct   = round(subopt   / N * 100, 1),
    rand_pct     = round(rand     / N * 100, 1),
    correlationhigh = factor(correlationhigh,
                             levels = c(0, 1),
                             labels = c("Low", "High"))
  ) |>
  arrange(treatment, correlationhigh)

print(tbl_treat_corr)

# 11g (wide). Same table pivoted: one row per treatment, % columns by Low/High
tbl_treat_corr_wide <- tbl_treat_corr |>
  select(treatment, correlationhigh, N,
         true_ext_pct, subopt_pct, rand_pct) |>
  pivot_wider(
    names_from  = correlationhigh,
    values_from = c(N, true_ext_pct, subopt_pct, rand_pct),
    names_glue  = "{.value}_{correlationhigh}"
  ) |>
  select(treatment,
         N_Low,            N_High,
         true_ext_pct_Low, true_ext_pct_High,
         subopt_pct_Low,   subopt_pct_High,
         rand_pct_Low,     rand_pct_High)

print(tbl_treat_corr_wide)


# ---- 6. Mixed-effects regression: trueextracted ~ correlationhigh -----------
library(lme4)

lmm_maxcorrect <- lmer(
  trueextracted ~ correlationhigh + (1 | subject_id),
  data = te_stats,
  REML = TRUE
)

summary(lmm_maxcorrect)

# Variance decomposition
vc <- as.data.frame(VarCorr(lmm_maxcorrect))
total_var <- sum(vc$vcov)
cat("\n--- Variance decomposition ---\n")
cat(sprintf("Subject intercept variance : %.3f  (%.1f%%)\n",
            vc$vcov[vc$grp == "subject_id"],
            100 * vc$vcov[vc$grp == "subject_id"] / total_var))
cat(sprintf("Residual variance          : %.3f  (%.1f%%)\n",
            vc$vcov[vc$grp == "Residual"],
            100 * vc$vcov[vc$grp == "Residual"] / total_var))
cat(sprintf("Fixed effect (correlationhigh): %.3f  (SE = %.3f)\n",
            fixef(lmm_maxcorrect)["correlationhigh"],
            sqrt(diag(vcov(lmm_maxcorrect)))["correlationhigh"]))


# ---- 7. Rule_used: posterior table and bar chart ---------------------------

RULE_POST_THRESHOLD <- 0.0   # set > 0 to keep only subjects with posterior >= threshold

rule_obs <- Data_3L |>
  filter(is.na(posterior) | posterior >= RULE_POST_THRESHOLD) |>
  distinct(subject_id, treatment, correlationhigh, Rule_used) |>
  mutate(
    Rule_used       = if_else(str_starts(Rule_used, "s_"), "Other", Rule_used),
    correlationhigh = factor(correlationhigh,
                             levels = c(0, 1),
                             labels = c("Low", "High"))
  )

# Proportion table (treatment × correlationhigh × Rule_used)
tbl_rule <- rule_obs |>
  group_by(treatment, correlationhigh, Rule_used) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(treatment, correlationhigh) |>
  mutate(pct = round(n / sum(n) * 100, 1)) |>
  ungroup() |>
  arrange(treatment, correlationhigh, desc(pct))

print(tbl_rule)

# Wide version: one row per treatment × Rule_used, columns = Low / High %
tbl_rule_wide <- tbl_rule |>
  select(treatment, correlationhigh, Rule_used, pct) |>
  pivot_wider(names_from  = correlationhigh,
              values_from = pct,
              values_fill = 0) |>
  arrange(treatment, desc(Low + High))

print(tbl_rule_wide)

# Bar chart: proportion of Rule_used by treatment × correlationhigh
rule_plot_data <- tbl_rule |>
  mutate(treatment = factor(treatment, levels = treatment_levels))

p_rule <- ggplot(rule_plot_data,
                 aes(x = reorder(Rule_used, -pct), y = pct, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", alpha = 0.85) +
  scale_fill_manual(values = c("Low" = "#2166ac", "High" = "#d6604d")) +
  facet_wrap(~ treatment, nrow = 1) +
  labs(
    x    = "Rule used",
    y    = "% of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(
    axis.text.x      = element_text(angle = 45, hjust = 1, size = 7),
    legend.position  = "bottom"
  )

print(p_rule)


# ---- 7b. Rule_used: collapsed 3-category bar chart -------------------------
# Categories:
#   True Extracted – Rule_used matches the treatment's true rule
#   Random         – Rule_used is "random" or "always_p"
#   Other Rule     – everything else (including s_* → "Other")

true_rule_map_3L <- c(
  AND     = "AND_rb",
  OR      = "OR_rb",
  INHIBIT = "INHIBIT_rb",
  EITHER  = "EITHER_rb",
  JOINT   = "JOINT_rb"
)

rule_cat_obs <- rule_obs |>
  mutate(
    true_rule = true_rule_map_3L[as.character(treatment)],
    category  = case_when(
      Rule_used %in% c("random", "always_p") ~ "Random",
      Rule_used == true_rule                  ~ "True Extracted",
      TRUE                                    ~ "Other Rule"
    ),
    category = factor(category, levels = c("True Extracted", "Other Rule", "Random"))
  )

rule_cat_plot_data <- rule_cat_obs |>
  mutate(treatment = factor(treatment, levels = treatment_levels)) |>
  group_by(treatment, correlationhigh, category) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(treatment, correlationhigh) |>
  mutate(pct = round(n / sum(n) * 100, 1)) |>
  ungroup()

p_rule_cat <- ggplot(rule_cat_plot_data,
                     aes(x = category, y = pct, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", alpha = 0.85) +
  scale_fill_manual(values = c("Low" = "#2166ac", "High" = "#d6604d")) +
  facet_wrap(~ treatment, nrow = 1) +
  labs(
    x    = "Category",
    y    = "% of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(
    axis.text.x     = element_text(angle = 45, hjust = 1, size = 8),
    legend.position = "bottom"
  )

print(p_rule_cat)


# ---- 7c. Rule_used: collapsed 3-category, all treatments pooled ------------

rule_cat_pooled_data <- rule_cat_obs |>
  group_by(correlationhigh, category) |>
  summarise(n = n(), .groups = "drop") |>
  group_by(correlationhigh) |>
  mutate(pct = round(n / sum(n) * 100, 1)) |>
  ungroup()

p_rule_cat_pooled <- ggplot(rule_cat_pooled_data,
                            aes(x = category, y = pct, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", alpha = 0.85) +
  scale_fill_manual(values = c("Low" = "#2166ac", "High" = "#d6604d")) +
  labs(
    x     = "Category",
    y     = "% of subjects",
    fill  = NULL,
    title = "All treatments pooled"
  ) +
  theme_bw() +
  theme(
    axis.text.x     = element_text(angle = 45, hjust = 1, size = 8),
    legend.position = "bottom"
  )

print(p_rule_cat_pooled)


# ---- 8. Prior shares from structural estimation ----------------------------
# ESTIMATION_RESULTS contains one population-level share per strategy per
# treatment — these are the MLE priors, i.e. the mixing weights before any
# subject-level data update the posteriors.

prior_long <- bind_rows(
  env_3L1$ESTIMATION_RESULTS |> mutate(session = "3L(1)"),
  env_3L2$ESTIMATION_RESULTS |> mutate(session = "3L(2)")
) |>
  select(-loglike, -num_par) |>
  pivot_longer(c(-treatment, -session), names_to = "Rule", values_to = "share") |>
  filter(!is.na(share), share > 0) |>
  mutate(
    Rule = if_else(str_starts(Rule, "s_"), "Other", Rule),
    correlationhigh = factor(
      as.integer(
        (session == "3L(1)" & treatment %in% c("AND", "INHIBIT")) |
        (session == "3L(2)" & treatment %in% c("JOINT", "OR", "EITHER"))
      ),
      levels = c(0, 1),
      labels = c("Low", "High")
    ),
    treatment = factor(treatment, levels = treatment_levels)
  ) |>
  group_by(treatment, session, correlationhigh, Rule) |>
  summarise(share = sum(share), .groups = "drop") |>
  mutate(pct = round(share * 100, 1))

# Table (treatment × session × Rule, sorted by share)
tbl_prior <- prior_long |>
  arrange(treatment, session, desc(pct))

print(tbl_prior)

# Wide version
tbl_prior_wide <- tbl_prior |>
  select(treatment, session, correlationhigh, Rule, pct) |>
  pivot_wider(names_from  = correlationhigh,
              values_from = pct,
              values_fill = 0) |>
  arrange(treatment, session, desc(Low + High))

print(tbl_prior_wide)

# Bar chart: rules with prior share >= 5%
rules_above5_pr <- prior_long |>
  filter(pct >= 5) |>
  distinct(Rule)

p_prior_top <- prior_long |>
  semi_join(rules_above5_pr, by = "Rule") |>
  ggplot(aes(x = reorder(Rule, -pct), y = pct, fill = correlationhigh)) +
  geom_bar(stat = "identity", position = "dodge", alpha = 0.85) +
  scale_fill_manual(values = c("Low" = "#2166ac", "High" = "#d6604d")) +
  facet_wrap(~ treatment, nrow = 1, scales = "free_x") +
  labs(
    x    = "Rule (prior)",
    y    = "Prior share (%)",
    fill = NULL
  ) +
  theme_bw() +
  theme(
    axis.text.x      = element_text(angle = 45, hjust = 1, size = 7),
    legend.position  = "bottom"
  )

print(p_prior_top)


# ---- 9. Types (confidence based on predicted_correct_self) -----------------

# Treatment-level variables (one value per subject × treatment)
subj_treat_ec <- Data_3L |>
  distinct(subject_id, treatment, correlationhigh,
           time_machine, predicted_correct_self, trueextracted)

# Subject-level averages across all treatments
subj_avgs_ec <- subj_treat_ec |>
  group_by(subject_id) |>
  summarise(
    avg_time_machine           = mean(time_machine,           na.rm = TRUE),
    avg_predicted_correct_self = mean(predicted_correct_self, na.rm = TRUE),
    .groups = "drop"
  )

# effort: above own avg time; confidence_2: above own avg predicted_correct_self
effort_conf2 <- subj_treat_ec |>
  left_join(subj_avgs_ec, by = "subject_id") |>
  mutate(
    effort       = as.integer(time_machine           > avg_time_machine),
    confidence_2 = as.integer(predicted_correct_self > avg_predicted_correct_self),
    factor_ec2   = factor(
      case_when(
        effort == 0 & confidence_2 == 0 ~ 1L,
        effort == 1 & confidence_2 == 0 ~ 2L,
        effort == 0 & confidence_2 == 1 ~ 3L,
        effort == 1 & confidence_2 == 1 ~ 4L
      ),
      levels = 1:4,
      labels = c("Low effort\nLow conf",
                 "High effort\nLow conf",
                 "Low effort\nHigh conf",
                 "High effort\nHigh conf")
    )
  )

# Distribution of factor_ec2 by treatment × correlationhigh
p_factor_ec2 <- effort_conf2 |>
  mutate(
    treatment       = factor(treatment, levels = treatment_levels),
    correlationhigh = factor(correlationhigh,
                             levels = c(0, 1),
                             labels = c("Low", "High"))
  ) |>
  group_by(correlationhigh, treatment, factor_ec2) |>
  summarise(
    n         = n(),
    prop_true = sum(trueextracted, na.rm = TRUE) / n(),
    .groups   = "drop"
  ) |>
  group_by(correlationhigh, treatment) |>
  mutate(frac = n / sum(n)) |>
  ungroup() |>
  ggplot(aes(x = factor_ec2, y = frac, fill = factor_ec2)) +
  geom_bar(stat = "identity", alpha = 0.85, width = 0.7) +
  geom_text(aes(label = scales::percent(prop_true, accuracy = 1)),
            vjust = -0.4, size = 2.5) +
  scale_fill_brewer(palette = "Set2") +
  facet_grid(correlationhigh ~ treatment,
             labeller = labeller(correlationhigh = corr_labels)) +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1),
                     expand = expansion(mult = c(0, 0.18))) +
  labs(
    x    = NULL,
    y    = "Fraction of subjects",
    fill = NULL
  ) +
  theme_bw() +
  theme(
    axis.text.x     = element_text(size = 7),
    legend.position = "bottom"
  )

print(p_factor_ec2)

# Aggregated across treatments — by correlationhigh only
p_factor_ec2_agg <- effort_conf2 |>
  mutate(correlationhigh = factor(correlationhigh,
                                  levels = c(0, 1),
                                  labels = c("Low", "High"))) |>
  group_by(correlationhigh, factor_ec2) |>
  summarise(
    n         = n(),
    prop_true = sum(trueextracted, na.rm = TRUE) / n(),
    .groups   = "drop"
  ) |>
  group_by(correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup() |>
  ggplot(aes(x = factor_ec2, y = frac, fill = factor_ec2)) +
  geom_bar(stat = "identity", alpha = 0.85, width = 0.7) +
  geom_text(aes(label = scales::percent(prop_true, accuracy = 1)),
            vjust = -0.4, size = 2.5) +
  scale_fill_brewer(palette = "Set2") +
  facet_wrap(~ correlationhigh, nrow = 1,
             labeller = labeller(correlationhigh = corr_labels)) +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1),
                     expand = expansion(mult = c(0, 0.18))) +
  labs(x = NULL, y = "Fraction of subjects", fill = NULL) +
  theme_bw() +
  theme(axis.text.x = element_text(size = 7), legend.position = "bottom")

print(p_factor_ec2_agg)


# ---- 10. Subject level types ------------------------------------------------
#
# For each subject × correlationhigh cell, average time_machine and
# predicted_correct_self across all treatments in that cell.
# Then compare each subject's cell average to the full group average
# (across all subjects within that correlationhigh level) to assign
# above/below-average status for both measures.

# Step 1: subject × correlationhigh averages
subj_corr_ec <- subj_treat_ec |>
  group_by(subject_id, correlationhigh) |>
  summarise(
    avg_time_corr  = mean(time_machine,           na.rm = TRUE),
    avg_conf_corr  = mean(predicted_correct_self, na.rm = TRUE),
    prop_true_corr = mean(trueextracted,           na.rm = TRUE),
    .groups = "drop"
  )

# Step 2: full correlationhigh-level averages (benchmark for classification)
corr_benchmarks <- subj_corr_ec |>
  group_by(correlationhigh) |>
  summarise(
    corr_avg_time = mean(avg_time_corr, na.rm = TRUE),
    corr_avg_conf = mean(avg_conf_corr, na.rm = TRUE),
    .groups = "drop"
  )

# Step 3: classify subjects and build 4-level factor
effort_conf_corr <- subj_corr_ec |>
  left_join(corr_benchmarks, by = "correlationhigh") |>
  mutate(
    effort_corr     = as.integer(avg_time_corr > corr_avg_time),
    confidence_corr = as.integer(avg_conf_corr > corr_avg_conf),
    factor_ec_corr  = factor(
      case_when(
        effort_corr == 0 & confidence_corr == 0 ~ 1L,
        effort_corr == 1 & confidence_corr == 0 ~ 2L,
        effort_corr == 0 & confidence_corr == 1 ~ 3L,
        effort_corr == 1 & confidence_corr == 1 ~ 4L
      ),
      levels = 1:4,
      labels = c("Low effort\nLow conf",
                 "High effort\nLow conf",
                 "Low effort\nHigh conf",
                 "High effort\nHigh conf")
    )
  )

# Step 4: plot — distribution and trueextracted rate by correlationhigh
p_factor_ec_corr <- effort_conf_corr |>
  mutate(correlationhigh = factor(correlationhigh,
                                  levels = c(0, 1),
                                  labels = c("Low", "High"))) |>
  group_by(correlationhigh, factor_ec_corr) |>
  summarise(
    n          = n(),
    prop_true  = mean(prop_true_corr, na.rm = TRUE),
    .groups    = "drop"
  ) |>
  group_by(correlationhigh) |>
  mutate(frac = n / sum(n)) |>
  ungroup() |>
  ggplot(aes(x = factor_ec_corr, y = frac, fill = factor_ec_corr)) +
  geom_bar(stat = "identity", alpha = 0.85, width = 0.7) +
  geom_text(aes(label = scales::percent(prop_true, accuracy = 1)),
            vjust = -0.4, size = 2.5) +
  scale_fill_brewer(palette = "Set2") +
  facet_wrap(~ correlationhigh, nrow = 1,
             labeller = labeller(correlationhigh = corr_labels)) +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1),
                     expand = expansion(mult = c(0, 0.18))) +
  labs(
    x       = NULL,
    y       = "Fraction of subjects",
    fill    = NULL,
    title   = "Effort × Confidence (within-correlationhigh avg benchmark)",
    caption = "Effort: avg time > corr. group avg. Confidence: avg predicted_correct_self > corr. group avg.\nBar labels = mean trueextracted rate within cell."
  ) +
  theme_bw() +
  theme(
    axis.text.x     = element_text(size = 7),
    legend.position = "bottom"
  )

print(p_factor_ec_corr)


# ---- 10b. Type consistency across correlationhigh levels --------------------
#
# For each subject, check whether they land in the same effort × confidence
# group in Low vs High correlationhigh. Since every subject has both Low and
# High treatments within their session, all subjects appear in both columns.

# Wide: one row per subject, one column per correlationhigh level
type_wide <- effort_conf_corr |>
  select(subject_id, correlationhigh, factor_ec_corr) |>
  mutate(corr_label = recode(as.character(correlationhigh),
                             "0" = "Low", "1" = "High")) |>
  select(-correlationhigh) |>
  pivot_wider(names_from  = corr_label,
              values_from = factor_ec_corr,
              names_prefix = "type_")

# Consistency summary: subjects with both Low and High assignments
consistency_summary <- type_wide |>
  filter(!is.na(type_Low) & !is.na(type_High)) |>
  summarise(
    n_subjects = n(),
    n_same     = sum(type_Low == type_High),
    pct_same   = round(100 * n_same / n_subjects, 1)
  )

cat("\n================================================================\n")
cat(" TYPE CONSISTENCY: Low vs High correlationhigh\n")
cat("================================================================\n")
cat(sprintf("  Subjects with both levels : %d\n",   consistency_summary$n_subjects))
cat(sprintf("  Same group in both levels : %d (%.1f%%)\n",
            consistency_summary$n_same, consistency_summary$pct_same))
cat("================================================================\n\n")

# Cross-tab heatmap: type at Low (x-axis) vs type at High (y-axis)
cross_tab_corr <- type_wide |>
  filter(!is.na(type_Low) & !is.na(type_High)) |>
  count(type_Low, type_High) |>
  group_by(type_Low) |>
  mutate(pct_row = n / sum(n)) |>
  ungroup()

p_type_consistency <- cross_tab_corr |>
  ggplot(aes(x = type_Low, y = type_High, fill = n)) +
  geom_tile(colour = "white", linewidth = 0.6) +
  geom_text(aes(label = sprintf("%d\n(%.0f%%)", n, 100 * pct_row)), size = 3) +
  scale_fill_gradient(low = "#f7f7f7", high = "#2166ac") +
  labs(
    x       = "Type at Low correlationhigh",
    y       = "Type at High correlationhigh",
    fill    = "N subjects",
    title   = "Type consistency across correlationhigh levels (Low vs High)",
    caption = "Cell count + % of row total. Diagonal = same type in both conditions."
  ) +
  theme_bw() +
  theme(axis.text.x = element_text(size = 7),
        axis.text.y = element_text(size = 7))

print(p_type_consistency)


# ---- 11. Variance decomposition: trueextracted / other_rule / random --------
#
# Outcomes (from section 13b rule_cat_obs):
#   trueextracted – subject inferred the true rule         (category == "True Extracted")
#   other_rule    – subject inferred a different rule      (category == "Other Rule")
#   random        – subject classified as random/always_p  (category == "Random")
#
# Predictors:
#   subject_id      – individual fixed effects (absorbs between-subject heterogeneity)
#   correlationhigh – correlation condition (Low / High)
#   treatment       – specific logical rule (AND, OR, EITHER, …)
#
# OLS is used for straightforward variance decomposition on binary outcomes.
# Type I (sequential) SS entered: subject_id → correlationhigh → treatment

reg_data_cat <- rule_cat_obs |>
  mutate(
    trueextracted = as.integer(category == "True Extracted"),
    other_rule    = as.integer(category == "Other Rule"),
    random        = as.integer(category == "Random"),
    subject_id    = factor(subject_id),
    treatment     = factor(treatment)
  )

# Helper: Type I ANOVA → variance decomposition table
# subject_id spans many levels; its rows are collapsed into one summary row.
anova_decomp <- function(model) {
  a      <- anova(model)
  ss     <- a$`Sum Sq`
  ss_res <- tail(ss, 1)
  ss_tot <- sum(ss)
  full <- tibble(
    Predictor        = rownames(a),
    df               = a$Df,
    SS               = ss,
    F_stat           = a$`F value`,
    p_value          = a$`Pr(>F)`,
    `eta2_%`         = 100 * ss / ss_tot,
    `partial_eta2_%` = 100 * ifelse(rownames(a) == "Residuals",
                                    NA_real_,
                                    ss / (ss + ss_res))
  )
  subj_rows  <- full |> filter(grepl("^subject_id", Predictor))
  other_rows <- full |> filter(!grepl("^subject_id", Predictor))
  if (nrow(subj_rows) > 0) {
    subj_summary <- tibble(
      Predictor        = "subject_id",
      df               = sum(subj_rows$df),
      SS               = sum(subj_rows$SS),
      F_stat           = NA_real_,
      p_value          = NA_real_,
      `eta2_%`         = sum(subj_rows$`eta2_%`),
      `partial_eta2_%` = 100 * sum(subj_rows$SS) /
                           (sum(subj_rows$SS) + tail(full$SS, 1))
    )
    full <- bind_rows(subj_summary, other_rows)
  }
  full |> mutate(
    SS               = round(SS,               4),
    F_stat           = round(F_stat,           3),
    p_value          = round(p_value,          4),
    `eta2_%`         = round(`eta2_%`,         2),
    `partial_eta2_%` = round(`partial_eta2_%`, 2)
  )
}

# Fit three OLS models (subject_id entered first)
outcomes <- c("trueextracted", "other_rule", "random")
models_cat <- setNames(
  lapply(outcomes, function(y)
    lm(as.formula(paste(y, "~ subject_id + correlationhigh + treatment")),
       data = reg_data_cat)),
  outcomes
)

for (y in outcomes) {
  m   <- models_cat[[y]]
  tbl <- anova_decomp(m)
  sm  <- summary(m)

  cat("\n================================================================\n")
  cat(sprintf(" VARIANCE DECOMPOSITION — %s\n", toupper(gsub("_", " ", y))))
  cat("================================================================\n")
  cat(sprintf("  R² = %.3f | Adj R² = %.3f | N = %d\n",
              sm$r.squared, sm$adj.r.squared, nobs(m)))
  cat("----------------------------------------------------------------\n")
  cat(" VARIANCE EXPLAINED BY PREDICTOR  (η²  and  partial η²  in %)\n")
  cat("----------------------------------------------------------------\n")
  print(as.data.frame(tbl), row.names = FALSE)
  cat("----------------------------------------------------------------\n")
  cat(" eta2 = SS/SS_total; partial eta2 = SS/(SS + SS_residual)\n")
  cat(" Type I SS: subject_id → correlationhigh → treatment\n")
  cat("================================================================\n")
}


# ---- 11b. Variance decomposition (cost types) --------------------------------
# factor_ec_corr classifies each subject × correlationhigh cell as one of:
#   1 – Low effort  / Low conf
#   2 – High effort / Low conf
#   3 – Low effort  / High conf
#   4 – High effort / High conf
# Type I SS entered: factor_ec_corr → correlationhigh → treatment

reg_data_type <- reg_data_cat |>
  left_join(
    effort_conf_corr |>
      mutate(subject_id      = factor(subject_id),
             correlationhigh = factor(correlationhigh,
                                      levels = c(0, 1),
                                      labels = c("Low", "High"))) |>
      select(subject_id, correlationhigh, factor_ec_corr),
    by = c("subject_id", "correlationhigh")
  )

models_type <- setNames(
  lapply(outcomes, function(y)
    lm(as.formula(paste(y, "~ factor_ec_corr + correlationhigh + treatment")),
       data = reg_data_type)),
  outcomes
)

cat("\n################################################################\n")
cat(" 18b  SUBJECT TYPE (effort × confidence) INSTEAD OF SUBJECT ID\n")
cat("################################################################\n")

for (y in outcomes) {
  m   <- models_type[[y]]
  tbl <- anova_decomp(m)
  sm  <- summary(m)

  cat("\n================================================================\n")
  cat(sprintf(" VARIANCE DECOMPOSITION — %s\n", toupper(gsub("_", " ", y))))
  cat("================================================================\n")
  cat(sprintf("  R² = %.3f | Adj R² = %.3f | N = %d\n",
              sm$r.squared, sm$adj.r.squared, nobs(m)))
  cat("----------------------------------------------------------------\n")
  cat(" VARIANCE EXPLAINED BY PREDICTOR  (η²  and  partial η²  in %)\n")
  cat("----------------------------------------------------------------\n")
  print(as.data.frame(tbl), row.names = FALSE)
  cat("----------------------------------------------------------------\n")
  cat(" eta2 = SS/SS_total; partial eta2 = SS/(SS + SS_residual)\n")
  cat(" Type I SS: factor_ec_corr → correlationhigh → treatment\n")
  cat("================================================================\n")

  coef_tbl <- as.data.frame(sm$coefficients)
  coef_tbl <- data.frame(
    Term      = rownames(coef_tbl),
    Estimate  = round(coef_tbl[, "Estimate"],   4),
    Std_Error = round(coef_tbl[, "Std. Error"], 4),
    t_value   = round(coef_tbl[, "t value"],    3),
    p_value   = round(coef_tbl[, "Pr(>|t|)"],  4)
  )
  cat("\n REGRESSION COEFFICIENTS\n")
  cat("----------------------------------------------------------------\n")
  print(coef_tbl, row.names = FALSE)
  cat("----------------------------------------------------------------\n")
}


# ---- 11c. ICC of trueextracted across correlation conditions ----------------
# Unit of observation: subject.
# Value per subject × condition: proportion of treatments where they extracted
# the true rule (prop_te). ICC treats Low and High as fixed raters and asks
# how consistently subjects are ranked across both conditions.
#
# model = "twoway"      : conditions are fixed (not a random sample)
# type  = "consistency" : ranks matter, not absolute level shifts
# unit  = "single"      : one condition = one observation per subject
#install.packages("irr")
library(irr)

icc_data <- rule_cat_obs |>
  group_by(subject_id, correlationhigh) |>
  summarise(prop_te = mean(category == "True Extracted"), .groups = "drop") |>
  pivot_wider(names_from = correlationhigh, values_from = prop_te) |>
  select(-subject_id)

icc_result <- irr::icc(icc_data, model = "twoway", type = "consistency", unit = "single")

cat("\n================================================================\n")
cat(" ICC — CONSISTENCY OF TRUE EXTRACTION ACROSS CORRELATION CONDITIONS\n")
cat("================================================================\n")
cat(sprintf("  ICC(2,1) consistency = %.3f  [95%% CI: %.3f – %.3f]\n",
            icc_result$value, icc_result$lbound, icc_result$ubound))
cat(sprintf("  F(%d, %d) = %.3f,  p = %.4f\n",
            icc_result$df1, icc_result$df2, icc_result$Fvalue, icc_result$p.value))
cat(sprintf("  N subjects = %d,  conditions = %s\n",
            icc_result$subjects,
            paste(names(icc_data), collapse = " / ")))
cat("================================================================\n")


# ---- 12. Subject filter for posterior + Strategy cost (MAP posteriors) ------
#
# For each (session x treatment) cell D with k strategies (rho > 0),
# set up the k(k-1)/2 pairwise equations:
#   v^D(pi_i) - v^D(pi_j) = rho^D(pi_i)*kappa(pi_i) - rho^D(pi_j)*kappa(pi_j)
# This is the linear system  A*kappa = b,  rank k-1.
# The unique minimum-norm solution (SVD pseudoinverse) requires no normalisation.
#
# v^D(pi)   = val_<strategy> from Data_3L (session-specific accuracy).
# rho^D(pi) = posterior MAP share (Rule_used) per (session x treatment) cell,
#             computed on subjects passing the KAPPA_ACC_THRESHOLD filter.

KAPPA_ACC_THRESHOLD <- 0.00   # e.g. 0.80, 0.85, 0.90

Data_3L_19 <- Data_3L |>
  filter(!is.na(posterior), posterior >= KAPPA_ACC_THRESHOLD)

cat(sprintf(
  "Section 19 subject filter (posterior >= %.0f%%): %d of %d subjects retained\n",
  KAPPA_ACC_THRESHOLD * 100,
  n_distinct(Data_3L_19$subject_id),
  n_distinct(Data_3L$subject_id)
))


# ---- 12a. v^D(pi) from val_ columns in Data_3L ------------------------------

val_long_19 <- Data_3L |>
  mutate(session = as.character(session)) |>
  group_by(session, treatment) |>
  summarise(across(starts_with("val_"), first), .groups = "drop") |>
  pivot_longer(starts_with("val_"),
               names_to  = "strategy_val",
               values_to = "v") |>
  mutate(strategy = str_remove(strategy_val, "^val_")) |>
  select(session, treatment, strategy, v)

cat(sprintf("val_long_19: %d rows | sessions: %s\n",
            nrow(val_long_19),
            paste(sort(unique(val_long_19$session)), collapse = ", ")))


# ---- 19b. Posterior shares rho^D(pi) from MAP assignment (Rule_used) --------

shares_long_19 <- Data_3L_19 |>
  mutate(session = as.character(session)) |>
  distinct(subject_id, session, treatment, Rule_used) |>
  filter(!is.na(Rule_used), Rule_used != "always_p") |>
  count(session, treatment, strategy = Rule_used, name = "n_subjects") |>
  group_by(session, treatment) |>
  mutate(rho = n_subjects / sum(n_subjects)) |>
  filter(rho > 0) |>
  mutate(rho = rho / sum(rho)) |>   # re-normalise after dropping zero-share strategies
  ungroup() |>
  select(session, treatment, strategy, rho)

cat(sprintf("shares_long_19: %d rows | sessions: %s\n",
            nrow(shares_long_19),
            paste(sort(unique(shares_long_19$session)), collapse = ", ")))


# ---- 19c. Join v and rho per (session x treatment x strategy) ---------------

cost_data_v_19 <- shares_long_19 |>
  inner_join(val_long_19, by = c("session", "treatment", "strategy")) |>
  filter(!is.na(v))

cat(sprintf("cost_data_v_19: %d rows | sessions: %s\n",
            nrow(cost_data_v_19),
            paste(sort(unique(cost_data_v_19$session)), collapse = ", ")))


# ---- 19d. Solver: minimum-norm + non-negative shift via SVD pseudoinverse ---
# For k strategies in one D cell, builds the k(k-1)/2 x k design matrix A
# and RHS b, then returns kappa = A^+ b shifted to enforce kappa >= 0.
#
# The null space of A is spanned by (1/rho_1, ..., 1/rho_k): adding alpha/rho_i
# to every kappa_i leaves all residuals unchanged (A*n = 0).  The minimum
# non-negative shift is alpha* = max(0, -min(kappa * rho)).

solve_kappa_cell <- function(strats, v_vec, rho_vec) {
  k <- length(strats)
  if (k == 1L)
    return(tibble(strategy = strats, kappa = NA_real_))

  pairs <- combn(k, 2)
  np    <- ncol(pairs)

  A <- matrix(0, np, k)
  b <- numeric(np)
  for (p in seq_len(np)) {
    i      <- pairs[1, p]; j <- pairs[2, p]
    A[p, i] <-  rho_vec[i]
    A[p, j] <- -rho_vec[j]
    b[p]    <- v_vec[i] - v_vec[j]
  }

  sv  <- svd(A)
  tol <- max(dim(A)) * max(sv$d) * .Machine$double.eps
  nz  <- sv$d > tol
  d_inv <- 1 / sv$d[nz]
  A_plus <- sv$v[, nz, drop = FALSE] %*%
              (d_inv * t(sv$u[, nz, drop = FALSE]))

  kappa <- as.numeric(A_plus %*% b)

  # Shift along null direction to enforce kappa >= 0
  alpha <- max(0, -min(kappa * rho_vec))
  kappa <- kappa + alpha / rho_vec

  tibble(strategy = strats, kappa = kappa)
}


# ---- 19e. Solve per (session x treatment) cell ------------------------------

kappa_solved_19 <- cost_data_v_19 |>
  group_by(session, treatment) |>
  summarise(
    kappa_df = list(solve_kappa_cell(strategy, v, rho)),
    .groups  = "drop"
  ) |>
  unnest(kappa_df)

corr_lookup_19 <- Data_3L |>
  distinct(session = as.character(session), treatment, correlationhigh)

cost_data_19 <- cost_data_v_19 |>
  left_join(kappa_solved_19, by = c("session", "treatment", "strategy")) |>
  left_join(corr_lookup_19,  by = c("session", "treatment")) |>
  mutate(
    correlationhigh = factor(correlationhigh,
                             levels = c(0L, 1L),
                             labels = c("Low", "High"))
  )


# ---- 19f. Results table -----------------------------------------------------

cat(sprintf(
  "\nSection 19 -- %d rows | sessions: %s | %d strategies\n",
  nrow(cost_data_19),
  paste(sort(unique(cost_data_19$session)), collapse = ", "),
  n_distinct(cost_data_19$strategy)
))

cat("\n=== Section 19: Strategy cost table (MAP posteriors, min-norm pairwise) ===\n")
print(
  cost_data_19 |>
    select(session, treatment, correlationhigh, strategy, rho, v, kappa) |>
    arrange(treatment, session, desc(rho)) |>
    mutate(across(c(rho, v, kappa), ~ round(.x, 4)))
)


# ---- 19g. Scatter: estimated kappa grouped by strategy family ----------------

scatter_data_19g <- cost_data_19 |>
  filter(!is.na(kappa)) |>
  mutate(
    strategy_group = case_when(
      str_starts(strategy, "s_")                            ~ "Other",
      strategy == "random"                                  ~ "Random",
      strategy %in% c("always_0", "always_1",
                      "red",  "not_red",
                      "blue", "not_blue",
                      "green", "not_green")                 ~ "1-variable",
      str_starts(strategy, "not_INHIBIT")                   ~ "not-INHIBIT",
      str_detect(strategy, "^AND")                          ~ "AND",
      str_detect(strategy, "^OR")                           ~ "OR",
      str_detect(strategy, "^EITHER")                       ~ "EITHER",
      str_detect(strategy, "^JOINT")                        ~ "JOINT",
      str_detect(strategy, "^INHIBIT")                      ~ "INHIBIT",
      str_detect(strategy, "^NOR")                          ~ "NOR",
      str_detect(strategy, "^NAND")                         ~ "NAND",
      TRUE                                                  ~ strategy
    )
  ) |>
  group_by(strategy_group) |>
  filter({
    iqr <- IQR(kappa)
    med <- median(kappa)
    kappa >= med - 1.5 * iqr & kappa <= med + 1.5 * iqr
  }) |>
  ungroup()

scatter_summary_19g <- scatter_data_19g |>
  group_by(strategy_group) |>
  summarise(
    mean_kappa = mean(kappa),
    se_kappa   = sd(kappa) / sqrt(n()),
    .groups    = "drop"
  )

ggplot(scatter_data_19g, aes(x = strategy_group, y = kappa)) +
  geom_hline(yintercept = 0, linetype = "dashed", colour = "grey60") +
  geom_jitter(aes(colour = treatment, shape = session),
              width = 0.15, height = 0, size = 2.5, alpha = 0.85) +
  geom_point(data = scatter_summary_19g,
             aes(x = strategy_group, y = mean_kappa),
             colour = "red", size = 3.5, inherit.aes = FALSE) +
  geom_errorbar(data = scatter_summary_19g,
                aes(x = strategy_group,
                    ymin = mean_kappa - se_kappa,
                    ymax = mean_kappa + se_kappa),
                colour = "red", width = 0.25, linewidth = 0.9,
                inherit.aes = FALSE) +
  scale_colour_brewer(palette = "Set1") +
  labs(
    x      = "Strategy family",
    y      = expression(hat(kappa)),
    colour = "Treatment",
    shape  = "Session",
    title  = "Estimated strategy cost by strategy family (MAP posteriors)"
  ) +
  theme_bw() +
  theme(axis.text.x = element_text(angle = 40, hjust = 1))


# ============================================================
# 20. Strategy cost (kappa) estimation — structural prior shares
# ============================================================
# Identical in structure to Section 19, but rho^D(pi) comes from the
# structural-estimation prior (ESTIMATION_RESULTS) rather than from
# individual MAP assignments.  No subject-level accuracy filter needed.
#
# v^D(pi)   = val_<strategy> from Data_3L (same as Section 19).
# rho^D(pi) = population share from ESTIMATION_RESULTS.
# ============================================================

# ---- 20a. v^D(pi) — reuse val_long_19 (same sessions / treatments)

val_long_20 <- val_long_19   # already computed in 19a

cat(sprintf("val_long_20: %d rows | sessions: %s\n",
            nrow(val_long_20),
            paste(sort(unique(val_long_20$session)), collapse = ", ")))


# ---- 20b. Structural shares rho^D(pi) (both sessions, no always_p) ----------

shares_long_20 <- bind_rows(
  env_3L1$ESTIMATION_RESULTS |> mutate(session = "3L(1)"),
  env_3L2$ESTIMATION_RESULTS |> mutate(session = "3L(2)")
) |>
  select(-loglike, -num_par) |>
  pivot_longer(-c(treatment, session), names_to = "strategy", values_to = "rho") |>
  filter(!is.na(rho), rho > 0, strategy != "always_p")

cat(sprintf("shares_long_20: %d rows | sessions: %s\n",
            nrow(shares_long_20),
            paste(sort(unique(shares_long_20$session)), collapse = ", ")))


# ---- 20c. Join v and rho per (session x treatment x strategy) ---------------

cost_data_v_20 <- shares_long_20 |>
  inner_join(val_long_20, by = c("session", "treatment", "strategy")) |>
  filter(!is.na(v))

cat(sprintf("cost_data_v_20: %d rows | sessions: %s\n",
            nrow(cost_data_v_20),
            paste(sort(unique(cost_data_v_20$session)), collapse = ", ")))


# ---- 20d. Solver: minimum-norm solution via SVD pseudoinverse ---------------
# Reuses solve_kappa_cell() defined in Section 19d (without non-negative shift,
# matching the original Analysis_3lights.R section 10 approach).

solve_kappa_cell_minnorm <- function(strats, v_vec, rho_vec) {
  k <- length(strats)
  if (k == 1L)
    return(tibble(strategy = strats, kappa = NA_real_))

  pairs <- combn(k, 2)
  np    <- ncol(pairs)

  A <- matrix(0, np, k)
  b <- numeric(np)
  for (p in seq_len(np)) {
    i      <- pairs[1, p]; j <- pairs[2, p]
    A[p, i] <-  rho_vec[i]
    A[p, j] <- -rho_vec[j]
    b[p]    <- v_vec[i] - v_vec[j]
  }

  sv  <- svd(A)
  tol <- max(dim(A)) * max(sv$d) * .Machine$double.eps
  nz  <- sv$d > tol
  d_inv <- 1 / sv$d[nz]
  A_plus <- sv$v[, nz, drop = FALSE] %*%
              (d_inv * t(sv$u[, nz, drop = FALSE]))

  tibble(strategy = strats, kappa = as.numeric(A_plus %*% b))
}


# ---- 20e. Solve per (session x treatment) cell ------------------------------

kappa_solved_20 <- cost_data_v_20 |>
  group_by(session, treatment) |>
  summarise(
    kappa_df = list(solve_kappa_cell_minnorm(strategy, v, rho)),
    .groups  = "drop"
  ) |>
  unnest(kappa_df)

cost_data_20 <- cost_data_v_20 |>
  left_join(kappa_solved_20, by = c("session", "treatment", "strategy")) |>
  mutate(
    correlationhigh = factor(
      as.integer(
        (session == "3L(1)" & treatment %in% c("AND", "INHIBIT")) |
        (session == "3L(2)" & treatment %in% c("JOINT", "OR", "EITHER"))
      ),
      levels = c(0L, 1L), labels = c("Low", "High")
    )
  )


# ---- 20f. Results table -----------------------------------------------------

cat(sprintf(
  "\nSection 20 -- %d cells | sessions: %s | %d strategies\n",
  nrow(cost_data_20),
  paste(sort(unique(cost_data_20$session)), collapse = ", "),
  n_distinct(cost_data_20$strategy)
))

cat("\n=== Section 20: Strategy cost table (structural prior, min-norm pairwise) ===\n")
print(
  cost_data_20 |>
    select(session, treatment, correlationhigh, strategy, rho, v, kappa) |>
    arrange(treatment, session, desc(rho)) |>
    mutate(across(c(rho, v, kappa), ~ round(.x, 4)))
)


# ---- 20g. Scatter: estimated kappa grouped by strategy family ---------------

scatter_data_20g <- cost_data_20 |>
  filter(!is.na(kappa)) |>
  mutate(
    strategy_group = case_when(
      str_starts(strategy, "s_")                            ~ "Other",
      strategy == "random"                                  ~ "Random",
      strategy %in% c("always_0", "always_1",
                      "red",  "not_red",
                      "blue", "not_blue",
                      "green", "not_green")                 ~ "1-variable",
      str_starts(strategy, "not_INHIBIT")                   ~ "not-INHIBIT",
      str_detect(strategy, "^AND")                          ~ "AND",
      str_detect(strategy, "^OR")                           ~ "OR",
      str_detect(strategy, "^EITHER")                       ~ "EITHER",
      str_detect(strategy, "^JOINT")                        ~ "JOINT",
      str_detect(strategy, "^INHIBIT")                      ~ "INHIBIT",
      str_detect(strategy, "^NOR")                          ~ "NOR",
      str_detect(strategy, "^NAND")                         ~ "NAND",
      TRUE                                                  ~ strategy
    )
  ) |>
  group_by(strategy_group) |>
  filter({
    iqr <- IQR(kappa)
    med <- median(kappa)
    kappa >= med - 1.5 * iqr & kappa <= med + 1.5 * iqr
  }) |>
  ungroup()

scatter_summary_20g <- scatter_data_20g |>
  group_by(strategy_group) |>
  summarise(
    mean_kappa = mean(kappa),
    se_kappa   = sd(kappa) / sqrt(n()),
    .groups    = "drop"
  )

ggplot(scatter_data_20g, aes(x = strategy_group, y = kappa)) +
  geom_hline(yintercept = 0, linetype = "dashed", colour = "grey60") +
  geom_jitter(aes(colour = treatment, shape = session),
              width = 0.15, height = 0, size = 2.5, alpha = 0.85) +
  geom_point(data = scatter_summary_20g,
             aes(x = strategy_group, y = mean_kappa),
             colour = "red", size = 3.5, inherit.aes = FALSE) +
  geom_errorbar(data = scatter_summary_20g,
                aes(x = strategy_group,
                    ymin = mean_kappa - se_kappa,
                    ymax = mean_kappa + se_kappa),
                colour = "red", width = 0.25, linewidth = 0.9,
                inherit.aes = FALSE) +
  scale_colour_brewer(palette = "Set1") +
  labs(
    x      = "Strategy family",
    y      = expression(hat(kappa)),
    colour = "Treatment",
    shape  = "Session",
    title  = "Estimated strategy cost by strategy family (structural prior shares)"
  ) +
  theme_bw() +
  theme(axis.text.x = element_text(angle = 40, hjust = 1))


# ---- 21. Scatter: P(rule) vs. distance from highest-accuracy rule -----------
#
# Unit of observation: one (session × treatment × strategy) triple.
# Every strategy with non-zero posterior share (MAP, section 19) in a given
# cell is one point.
#
# Y axis — rho: posterior MAP share from shares_long_19.
# X axis — distance: max accuracy in the cell minus this strategy's accuracy.
#   distance = max(val_*) − val_{strategy}
#   = 0   → strategy is as accurate as the best possible rule
#   > 0   → strategy is accuracy-inferior

# 21a. Best accuracy per (session × treatment) across all strategies
max_val_21 <- val_long_19 |>
  group_by(session, treatment) |>
  summarise(max_val = max(v, na.rm = TRUE), .groups = "drop")

# 21b. Join rho and val for every strategy present in each cell,
#      then compute distance from the best rule.
scatter_21 <- shares_long_19 |>
  inner_join(val_long_19,  by = c("session", "treatment", "strategy")) |>
  inner_join(max_val_21,   by = c("session", "treatment")) |>
  mutate(
    distance = max_val - v,
    correlationhigh = factor(
      as.integer(
        (session == "3L(1)" & treatment %in% c("AND", "INHIBIT")) |
        (session == "3L(2)" & treatment %in% c("JOINT", "OR", "EITHER"))
      ),
      levels = c(0L, 1L),
      labels = c("Low", "High")
    )
  )

cat(sprintf("scatter_21: %d points | sessions: %s | %d strategies\n",
            nrow(scatter_21),
            paste(sort(unique(scatter_21$session)), collapse = ", "),
            n_distinct(scatter_21$strategy)))

cat("\n=== Section 21: Rule share and distance from best rule ===\n")
print(
  scatter_21 |>
    select(session, treatment, strategy, rho, v, max_val, distance) |>
    arrange(treatment, session, distance) |>
    mutate(across(c(rho, v, max_val, distance), ~ round(.x, 4)))
)

# 21c. OLS regression: rho ~ distance
lm_21 <- lm(rho ~ distance, data = scatter_21)
sm_21 <- summary(lm_21)

cat("\n=== Section 21: OLS regression  rho ~ distance ===\n")
cat(sprintf("  Intercept : %.4f  (SE = %.4f,  t = %.3f,  p = %.4f)\n",
            coef(sm_21)[1, 1], coef(sm_21)[1, 2],
            coef(sm_21)[1, 3], coef(sm_21)[1, 4]))
cat(sprintf("  distance  : %.4f  (SE = %.4f,  t = %.3f,  p = %.4f)\n",
            coef(sm_21)[2, 1], coef(sm_21)[2, 2],
            coef(sm_21)[2, 3], coef(sm_21)[2, 4]))
cat(sprintf("  R\u00b2 = %.3f  |  Adj R\u00b2 = %.3f  |  N = %d\n",
            sm_21$r.squared, sm_21$adj.r.squared, nobs(lm_21)))

# 21d. Plot with best-fit line
p_21 <- ggplot(scatter_21,
               aes(x = distance, y = rho)) +
  geom_vline(xintercept = 0, linetype = "dashed", colour = "grey60") +
  geom_point(aes(colour = treatment, shape = session),
             size = 2.5, alpha = 0.85) +
  geom_smooth(method = "lm", formula = y ~ x,
              colour = "black", linewidth = 0.9, se = TRUE) +
  scale_colour_brewer(palette = "Set1") +
  scale_x_continuous(
    name   = "Distance from highest-accuracy rule\n(max accuracy \u2212 strategy accuracy)",
    limits = c(0, NA),
    expand = expansion(mult = c(0.02, 0.08))
  ) +
  scale_y_continuous(
    name   = "Posterior share \u03c1  [P(rule extracted) in treatment \u00d7 session]",
    limits = c(0, 1)
  ) +
  labs(
    title    = "Rule extraction probability vs. accuracy gap to best rule",
    subtitle = sprintf("\u03b2(distance) = %.3f, p = %.4f, R\u00b2 = %.3f",
                       coef(lm_21)[["distance"]],
                       coef(sm_21)[2, 4],
                       sm_21$r.squared),
    colour   = "Treatment",
    shape    = "Session"
  ) +
  theme_bw() +
  theme(
    legend.position = "right",
    axis.title      = element_text(size = 9),
    plot.subtitle   = element_text(size = 8, colour = "grey30")
  )

print(p_21)


# ---- 22. Scatter: P(rule) vs. distance — suboptimal rules only (distance > 0) ----
#
# Same as section 21 but restricted to strategies whose accuracy is strictly
# below the best rule in the cell (distance > 0).  Dropping the true-rule
# points (distance = 0) lets the regression capture how the share of
# *suboptimal* strategies varies with their accuracy gap.

scatter_22 <- scatter_21 |>
  filter(distance > 0)

cat(sprintf("\nscatter_22: %d points (distance > 0) | %d strategies\n",
            nrow(scatter_22), n_distinct(scatter_22$strategy)))

# 22a. OLS regression on suboptimal-rule points only
lm_22 <- lm(rho ~ distance, data = scatter_22)
sm_22 <- summary(lm_22)

cat("\n=== Section 22: OLS regression  rho ~ distance  (distance > 0 only) ===\n")
cat(sprintf("  Intercept : %.4f  (SE = %.4f,  t = %.3f,  p = %.4f)\n",
            coef(sm_22)[1, 1], coef(sm_22)[1, 2],
            coef(sm_22)[1, 3], coef(sm_22)[1, 4]))
cat(sprintf("  distance  : %.4f  (SE = %.4f,  t = %.3f,  p = %.4f)\n",
            coef(sm_22)[2, 1], coef(sm_22)[2, 2],
            coef(sm_22)[2, 3], coef(sm_22)[2, 4]))
cat(sprintf("  R\u00b2 = %.3f  |  Adj R\u00b2 = %.3f  |  N = %d\n",
            sm_22$r.squared, sm_22$adj.r.squared, nobs(lm_22)))

# 22b. Plot
p_22 <- ggplot(scatter_22,
               aes(x = distance, y = rho)) +
  geom_point(aes(colour = treatment, shape = session),
             size = 2.5, alpha = 0.85) +
  geom_smooth(method = "lm", formula = y ~ x,
              colour = "black", linewidth = 0.9, se = TRUE) +
  scale_colour_brewer(palette = "Set1") +
  scale_x_continuous(
    name   = "Distance from highest-accuracy rule\n(max accuracy \u2212 strategy accuracy)",
    expand = expansion(mult = c(0.02, 0.08))
  ) +
  scale_y_continuous(
    name   = "Posterior share \u03c1  [P(rule extracted) in treatment \u00d7 session]",
    limits = c(0, 1)
  ) +
  labs(
    title    = "Rule extraction probability vs. accuracy gap — suboptimal rules only",
    subtitle = sprintf("\u03b2(distance) = %.3f, p = %.4f, R\u00b2 = %.3f  (distance > 0 only)",
                       coef(lm_22)[["distance"]],
                       coef(sm_22)[2, 4],
                       sm_22$r.squared),
    colour   = "Treatment",
    shape    = "Session"
  ) +
  theme_bw() +
  theme(
    legend.position = "right",
    axis.title      = element_text(size = 9),
    plot.subtitle   = element_text(size = 8, colour = "grey30")
  )

print(p_22)