# =============================================================================
# Analysis_3lights.R — MLE strategy estimation with stratEst
#                      for the Shaded Lights Experiment (3 Lights)
#
# OBSERVABLE STATES (8 total — one per unique (Red, Blue, Green) configuration):
#   State 1: (0,0,0) — all off
#   State 2: (0,0,1) — green on only
#   State 3: (0,1,0) — blue on only
#   State 4: (0,1,1) — blue + green on
#   State 5: (1,0,0) — red on only
#   State 6: (1,0,1) — red + green on
#   State 7: (1,1,0) — red + blue on
#   State 8: (1,1,1) — all on
#
# DETERMINISTIC STRATEGIES — complete enumeration:
# With 8 binary states there are 2^8 = 256 deterministic strategies.
# Each is identified by an 8-character binary string, where the k-th character
# (1-indexed) is the action in state k: "0" = predict no sound, "1" = predict sound.
# E.g. "s_00000011" predicts sound only in states 7 and 8 → AND(red,blue).
#
# Named strategies (integer j = decimal value of action bits, LSB = state 1):
#
# ── Single-variable (third light irrelevant) ──────────────────────────────────
#   always_0    j=0  : (0,0,0,0,0,0,0,0) — never predict sound
#   always_1    j=255: (1,1,1,1,1,1,1,1) — always predict sound
#   green       j=170: sound iff green = 1                (red/blue irrelevant)
#   not_green   j=85 : sound iff green = 0
#   blue        j=204: sound iff blue  = 1                (red/green irrelevant)
#   not_blue    j=51 : sound iff blue  = 0
#   red         j=240: sound iff red   = 1                (blue/green irrelevant)
#   not_red     j=15 : sound iff red   = 0
#
# ── Two-variable rules (one light irrelevant) — for each of the three pairs ───
#
#   Pair (red, blue) — green irrelevant:
#   NOR_rb      j=3  : sound iff red=0  AND  blue=0
#   AND_rb      j=192: sound iff red=1  AND  blue=1       [true AND treatment]
#   INHIBIT_rb  j=48 : sound iff red=1  AND  blue=0
#   INHIBIT_br  j=12 : sound iff blue=1 AND  red=0
#   EITHER_rb   j=60 : sound iff red ≠ blue               [true EITHER/XOR treatment]
#   JOINT_rb    j=195: sound iff red =  blue (XNOR)       [true JOINT treatment]
#   OR_rb       j=252: sound iff red=1  OR   blue=1       [true OR treatment]
#   NAND_rb     j=63 : sound unless red=1 AND blue=1
#   not_INHIBIT_rb j=207: sound unless red=1 AND blue=0
#   not_INHIBIT_br j=243: sound unless blue=1 AND red=0
#
#   Pair (red, green) — blue irrelevant:
#   NOR_rg      j=5  : sound iff red=0  AND  green=0
#   AND_rg      j=160: sound iff red=1  AND  green=1
#   INHIBIT_rg  j=80 : sound iff red=1  AND  green=0
#   INHIBIT_gr  j=10 : sound iff green=1 AND red=0
#   EITHER_rg   j=90 : sound iff red ≠ green
#   JOINT_rg    j=165: sound iff red =  green
#   OR_rg       j=250: sound iff red=1  OR   green=1
#   NAND_rg     j=95 : sound unless red=1 AND green=1
#   not_INHIBIT_rg j=175: sound unless red=1 AND green=0
#   not_INHIBIT_gr j=245: sound unless green=1 AND red=0
#
#   Pair (blue, green) — red irrelevant:
#   NOR_bg      j=17 : sound iff blue=0  AND  green=0
#   AND_bg      j=136: sound iff blue=1  AND  green=1
#   INHIBIT_bg  j=68 : sound iff blue=1  AND  green=0
#   INHIBIT_gb  j=34 : sound iff green=1 AND  blue=0
#   EITHER_bg   j=102: sound iff blue ≠ green
#   JOINT_bg    j=153: sound iff blue =  green
#   OR_bg       j=238: sound iff blue=1  OR   green=1
#   NAND_bg     j=119: sound unless blue=1 AND green=1
#   not_INHIBIT_bg j=187: sound unless blue=1 AND green=0
#   not_INHIBIT_gb j=221: sound unless green=1 AND blue=0
#
# ── Three-variable rules (all three lights matter) ────────────────────────────
#   NOR3        j=1  : sound iff all off             (0,0,0) only
#   green_alone j=2  : sound iff green only on       (0,0,1) only
#   blue_alone  j=4  : sound iff blue only on        (0,1,0) only
#   blue_green  j=8  : sound iff blue+green only     (0,1,1) only
#   red_alone   j=16 : sound iff red only on         (1,0,0) only
#   red_green   j=32 : sound iff red+green only      (1,0,1) only
#   red_blue    j=64 : sound iff red+blue only       (1,1,0) only
#   AND3        j=128: sound iff all on              (1,1,1) only
#   majority    j=232: sound iff ≥ 2 lights on       (states 4,6,7,8)
#   minority    j=23 : sound iff ≤ 1 light on        (states 1,2,3,5)
#   ODD_parity  j=150: sound iff odd number of lights on (XOR of all three)
#   EVEN_parity j=105: sound iff even number of lights on (XNOR of all three)
#   NAND3       j=127: sound unless all on
#   OR3         j=254: sound iff at least one on
#   All remaining j values: named "s_XXXXXXXX" (binary action string)
#
# Additionally:
#   random   : 1-state; P(predict sound) = 0.5 fixed
#   always_p : 1-state; P(predict sound) = p estimated via MLE
# =============================================================================


# ============================================================
# 0) Packages
# ============================================================

library(dplyr)
library(tidyr)
library(stringr)
library(stratEst)


# ============================================================
# 1) Generate dataset (source cleaning script)
# ============================================================

script_dir <- dirname(rstudioapi::getActiveDocumentContext()$path)
source(file.path(script_dir, "DOfile_realexp_3lights.R"))

# Rename to stratEst conventions
df_strat <- df_final %>%
  rename(
    id     = subject_id,
    period = Guess_Number,
    input  = Light_Config
  ) %>%
  mutate(
    choice    = as.character(Guess),
    treatment = as.character(treatment),
    game      = 1L
  ) %>%
  select(id, period, choice, input, treatment, game) %>%
  arrange(id, treatment, period)

cat("Dataset ready:", nrow(df_strat), "rows |",
    n_distinct(df_strat$id), "subjects |",
    n_distinct(df_strat$treatment), "treatments\n")
cat("Choice values:", paste(sort(unique(df_strat$choice)), collapse = ", "), "\n")
cat("Input values: ", paste(sort(unique(df_strat$input)),  collapse = ", "), "\n")


# ============================================================
# 2) Strategy building blocks
# ============================================================

choice_labels <- c("0", "1")

# 8 light configurations as input levels (match Light_Config in data exactly)
input_vals_8 <- c("(0,0,0)", "(0,0,1)", "(0,1,0)", "(0,1,1)",
                  "(1,0,0)", "(1,0,1)", "(1,1,0)", "(1,1,1)")
inputs_strat_8 <- c(NA, input_vals_8)

# Shared transition matrix for all 8-state strategies.
# States map 1-to-1 with Light_Config values (state k ↔ input_vals_8[k]).
# From any state, observing input x always moves the automaton to state(x).
# Layout: for each of the 8 states (rows), list next-state per non-NA input.
# Length = 8 states × 8 non-NA inputs = 64 values.
tr_8state <- as.numeric(rep(1:8, times = 8))

# Helper: build a deterministic 8-state strategy from an 8-element action vector.
# actions[k] in {0,1}: 0 = predict no sound in state k; 1 = predict sound.
# State ordering matches input_vals_8: k=1 → (0,0,0), ..., k=8 → (1,1,1).
# prob.choices layout: c(P("0"|s1), P("1"|s1), ..., P("0"|s8), P("1"|s8)).
make_det_strategy_8 <- function(actions) {
  stopifnot(length(actions) == 8, all(actions %in% c(0L, 1L)))
  pc <- as.double(unlist(lapply(actions, function(a) c(1 - a, a))))
  stratEst.strategy(
    choices      = choice_labels,
    inputs       = inputs_strat_8,
    num.states   = 8,
    tr.inputs    = tr_8state,
    prob.choices = pc
  )
}


# ============================================================
# 3) Build all 256 deterministic strategies
# ============================================================

# For integer j in 0:255, the action in state k (1-indexed) is bit (k-1) of j,
# i.e. as.integer(intToBits(j))[k].  intToBits() returns bits LSB-first, which
# aligns perfectly with the state ordering: state 1=(0,0,0) ↔ bit 0 of j.
#
# Systematic name: "s_" + 8-char binary string, character k = action in state k.
# E.g. j=192 → bits c(0,0,0,0,0,0,1,1) → name "s_00000011" → AND(red,blue).

all_det <- setNames(
  lapply(0:255, function(j) {
    actions <- as.integer(intToBits(j))[1:8]
    make_det_strategy_8(actions)
  }),
  sapply(0:255, function(j)
    paste0("s_", paste(as.integer(intToBits(j))[1:8], collapse = ""))
  )
)

cat("Built", length(all_det), "deterministic strategies\n")

# ── Rename interpretable strategies with meaningful labels ────────────────────
# Map: integer j → meaningful name (state-order binary bits, LSB = state 1)

meaningful_aliases <- c(
  # ── Single-variable ──────────────────────────────────────────────────────
  always_0         =   0L,   # (0,0,0,0,0,0,0,0)
  always_1         = 255L,   # (1,1,1,1,1,1,1,1)
  green            = 170L,   # sound iff g=1  (01010101)
  not_green        =  85L,   # sound iff g=0  (10101010)
  blue             = 204L,   # sound iff b=1  (00110011)
  not_blue         =  51L,   # sound iff b=0  (11001100)
  red              = 240L,   # sound iff r=1  (00001111)
  not_red          =  15L,   # sound iff r=0  (11110000)

  # ── Two-variable: (red, blue) — green irrelevant ─────────────────────────
  NOR_rb           =   3L,   # sound iff r=0 AND b=0         (11000000)
  INHIBIT_br       =  12L,   # sound iff b=1 AND r=0         (00110000) ← blue_alone equiv.
  INHIBIT_rb       =  48L,   # sound iff r=1 AND b=0         (00001100) [true INHIBIT treatment]
  EITHER_rb        =  60L,   # sound iff r≠b                 (00110011) [true EITHER treatment]
  NAND_rb          =  63L,   # sound unless r=1 AND b=1      (11111100)
  AND_rb           = 192L,   # sound iff r=1 AND b=1         (00000011) [true AND treatment]
  OR_rb            = 252L,   # sound iff r=1 OR b=1          (00111111) [true OR treatment]
  JOINT_rb         = 195L,   # sound iff r=b (XNOR)          (11000011) [true JOINT treatment]
  not_INHIBIT_rb   = 207L,   # sound unless r=1 AND b=0      (11110011)
  not_INHIBIT_br   = 243L,   # sound unless b=1 AND r=0      (11001111)

  # ── Two-variable: (red, green) — blue irrelevant ─────────────────────────
  NOR_rg           =   5L,   # sound iff r=0 AND g=0         (10100000)
  INHIBIT_gr       =  10L,   # sound iff g=1 AND r=0         (01010000)
  AND_rg           = 160L,   # sound iff r=1 AND g=1         (00000101)
  OR_rg            = 250L,   # sound iff r=1 OR g=1          (01011111)
  INHIBIT_rg       =  80L,   # sound iff r=1 AND g=0         (00001010)
  EITHER_rg        =  90L,   # sound iff r≠g                 (01011010)
  JOINT_rg         = 165L,   # sound iff r=g                 (10100101)
  NAND_rg          =  95L,   # sound unless r=1 AND g=1      (11111010)
  not_INHIBIT_rg   = 175L,   # sound unless r=1 AND g=0      (11110101)
  not_INHIBIT_gr   = 245L,   # sound unless g=1 AND r=0      (10101111)

  # ── Two-variable: (blue, green) — red irrelevant ─────────────────────────
  NOR_bg           =  17L,   # sound iff b=0 AND g=0         (10001000)
  INHIBIT_gb       =  34L,   # sound iff g=1 AND b=0         (01000100)
  AND_bg           = 136L,   # sound iff b=1 AND g=1         (00001001)
  OR_bg            = 238L,   # sound iff b=1 OR g=1          (01110111)
  INHIBIT_bg       =  68L,   # sound iff b=1 AND g=0         (00100010)
  EITHER_bg        = 102L,   # sound iff b≠g                 (01100110)
  JOINT_bg         = 153L,   # sound iff b=g                 (10011001)
  NAND_bg          = 119L,   # sound unless b=1 AND g=1      (11110110)
  not_INHIBIT_bg   = 187L,   # sound unless b=1 AND g=0      (11011101)
  not_INHIBIT_gb   = 221L,   # sound unless g=1 AND b=0      (10111011)

  # ── Three-variable (all lights matter) ───────────────────────────────────
  NOR3             =   1L,   # sound iff all off              (0,0,0) only
  green_alone      =   2L,   # sound iff green only on        (0,0,1) only
  blue_alone       =   4L,   # sound iff blue only on         (0,1,0) only
  blue_green_only  =   8L,   # sound iff blue+green only      (0,1,1) only
  red_alone        =  16L,   # sound iff red only on          (1,0,0) only
  red_green_only   =  32L,   # sound iff red+green only       (1,0,1) only
  red_blue_only    =  64L,   # sound iff red+blue only        (1,1,0) only
  AND3             = 128L,   # sound iff all on               (1,1,1) only
  minority         =  23L,   # sound iff ≤ 1 light on         (states 1-3, 5)
  majority         = 232L,   # sound iff ≥ 2 lights on        (states 4, 6-8)
  ODD_parity       = 150L,   # sound iff odd count on (XOR3)  (states 2,3,5,8)
  EVEN_parity      = 105L,   # sound iff even count on (XNOR3)(states 1,4,6,7)
  NAND3            = 127L,   # sound unless all on
  OR3              = 254L    # sound iff at least one on
)

# Verify that each alias points to a distinct j value
stopifnot(!anyDuplicated(meaningful_aliases))

# Apply meaningful names: rename the corresponding entries in all_det
for (nice_name in names(meaningful_aliases)) {
  j     <- meaningful_aliases[[nice_name]]
  old   <- paste0("s_", paste(as.integer(intToBits(j))[1:8], collapse = ""))
  idx   <- match(old, names(all_det))
  if (!is.na(idx)) names(all_det)[idx] <- nice_name
}

cat("Named interpretable strategies:\n")
cat(paste(names(meaningful_aliases), collapse = ", "), "\n")
cat("Remaining strategies carry systematic s_XXXXXXXX names.\n")


# ============================================================
# 4) Additional strategies: random and always_p
# ============================================================

# For the 1-state strategies, any input transitions back to state 1.
# Use inputs_strat_8 to remain compatible with 8-input data.
random <- stratEst.strategy(
  choices      = choice_labels,
  inputs       = inputs_strat_8,
  num.states   = 1,
  tr.inputs    = rep(1, 8),       # stay in state 1 regardless of input
  prob.choices = c(0.5, 0.5)
)

always_p <- stratEst.strategy(
  choices      = choice_labels,
  inputs       = inputs_strat_8,
  num.states   = 1,
  tr.inputs    = rep(1, 8),
  prob.choices = c(NA_real_, NA_real_)   # p estimated via MLE
)

cat("Built 2 additional strategies: random (p = 0.5 fixed) and always_p (p estimated)\n")


# ============================================================
# 5) Assemble full strategy list (256 deterministic + 2 stochastic)
# ============================================================

strategies_all <- c(
  all_det,
  list(random = random, always_p = always_p)
)

cat("Total strategies:", length(strategies_all), "\n")
stopifnot(!any(sapply(strategies_all, is.null)))
stopifnot(!anyDuplicated(names(strategies_all)))


# ============================================================
# 6) Estimation: loop over all treatments
# ============================================================

all_treatments <- sort(unique(df_strat$treatment))
results_list   <- vector("list", length(all_treatments))

for (idx in seq_along(all_treatments)) {
  TREAT <- all_treatments[idx]
  cat("\n========== Treatment:", TREAT,
      "(", idx, "/", length(all_treatments), ") ==========\n")

  df_t <- df_strat %>%
    filter(treatment == TREAT) %>%
    mutate(treatment = as.factor(treatment))

  cat("  Rows:", nrow(df_t), "| Subjects:", n_distinct(df_t$id), "\n")

  model_t <- tryCatch(
    stratEst.model(
      data       = df_t,
      strategies = strategies_all,
      inner.runs = 10,#100
      outer.runs = 1, #2
      inner.max  = 10 #100
    ),
    error = function(e) {
      cat("  ERROR:", conditionMessage(e), "\n")
      NULL
    }
  )

  if (!is.null(model_t)) {
    cat("  dim(model_t$shares):", dim(model_t$shares), "\n")
    cat("  sum(model_t$shares):", sum(model_t$shares), "\n")
    shares_vec        <- as.numeric(model_t$shares)
    names(shares_vec) <- names(strategies_all)

    results_list[[idx]] <- data.frame(
      treatment      = TREAT,
      loglike        = model_t$loglike,
      num_par        = model_t$num.par,
      t(shares_vec),
      check.names    = FALSE,
      stringsAsFactors = FALSE
    )
    cat("  Log-likelihood:", round(model_t$loglike, 3), "\n")
    cat("  Summary:\n")
    print(summary(model_t, plot.shares = FALSE))
  } else {
    results_list[[idx]] <- data.frame(
      treatment = TREAT, loglike = NA_real_, num_par = NA_integer_
    )
  }
}

# Combine all treatments into one results table
ESTIMATION_RESULTS <- bind_rows(results_list) %>%
  arrange(treatment)

# Drop strategy columns that are zero across all treatments
strat_cols <- setdiff(names(ESTIMATION_RESULTS), c("treatment", "loglike", "num_par"))
nonzero_cols <- strat_cols[colSums(ESTIMATION_RESULTS[, strat_cols], na.rm = TRUE) > 0]
ESTIMATION_RESULTS <- ESTIMATION_RESULTS[, c("treatment", "loglike", "num_par", nonzero_cols)]

cat("\n=== ESTIMATION_RESULTS (all treatments) ===\n")
print(ESTIMATION_RESULTS %>% mutate(across(where(is.numeric), ~round(., 4))))

# Save results
output_path <- file.path(script_dir, "Structural_Results_3lights.csv")
write.csv(ESTIMATION_RESULTS, output_path, row.names = FALSE)
cat("Saved to", output_path, "\n")


# ============================================================
# 7) Per-treatment share plots (non-zero strategies only)
# ============================================================

for (TREAT in sort(unique(ESTIMATION_RESULTS$treatment))) {
  row      <- ESTIMATION_RESULTS[ESTIMATION_RESULTS$treatment == TREAT, nonzero_cols]
  shares   <- as.numeric(row)
  names(shares) <- nonzero_cols
  shares   <- sort(shares[shares > 0], decreasing = TRUE)

  if (length(shares) == 0) next

  old_mar <- par(mar = c(10, 4, 3, 1))
  barplot(
    shares,
    main   = paste("Estimated shares —", TREAT),
    ylab   = "share",
    ylim   = c(0, 1),
    las    = 2,
    cex.names = 0.75
  )
  par(old_mar)
}