# ============================================================================
# MODÈLES Q-LEARNING EMBOÎTÉS POUR DÉCISION AVEC ÉVÉNEMENTS RARES
# ============================================================================

library(tidyverse)
library(DEoptim)
library(numDeriv)

# ============================================================================
# FONCTION GÉNÉRIQUE DE Q-LEARNING
# ============================================================================

qlearning_generic <- function(params, data, model_config, return_negLL = TRUE) {
  
  # Conversion des choix en indices numériques
  if (is.factor(data$button_name) || is.character(data$button_name)) {
    choice_levels <- c("antifragile", "fragile", "robuste", "vulnerable")
    data$choice_idx <- match(as.character(data$button_name), choice_levels)
  } else {
    data$choice_idx <- data$button_name
  }
  
  n_arms <- 4
  n_trials <- nrow(data)
  
  # Extraction des paramètres selon la configuration du modèle
  param_idx <- 1
  
  # ALPHA(S)
  if (model_config$n_alpha == 1) {
    alpha_loss <- alpha_gain <- alpha_BS <- alpha_JP <- plogis(params[param_idx])
    param_idx <- param_idx + 1
  } else if (model_config$n_alpha == 2) {
    alpha_loss <- plogis(params[param_idx])
    alpha_gain <- plogis(params[param_idx + 1])
    alpha_BS <- alpha_loss
    alpha_JP <- alpha_gain
    param_idx <- param_idx + 2
  } else if (model_config$n_alpha == 4) {
    alpha_loss <- plogis(params[param_idx])
    alpha_gain <- plogis(params[param_idx + 1])
    alpha_BS <- plogis(params[param_idx + 2])
    alpha_JP <- plogis(params[param_idx + 3])
    param_idx <- param_idx + 4
  }
  
  # FORGET(S)
  if (model_config$n_forget == 1) {
    forget <- rep(plogis(params[param_idx]), n_arms)
    param_idx <- param_idx + 1
  } else if (model_config$n_forget == 4) {
    forget <- plogis(params[param_idx:(param_idx + 3)])
    param_idx <- param_idx + 4
  }
  
  # LAMBDA(S)
  if (model_config$n_lambda == 1) {
    lambda <- rep(exp(params[param_idx]), n_arms)
    param_idx <- param_idx + 1
  } else if (model_config$n_lambda == 4) {
    lambda <- exp(params[param_idx:(param_idx + 3)])
    param_idx <- param_idx + 4
  }
  
  # RHO(S) - Biais pour événements rares
  if (model_config$has_rho) {
    rho_BS <- params[param_idx]      # BS avoidance
    rho_JP <- params[param_idx + 1]  # JP seeking
    param_idx <- param_idx + 2
  } else {
    rho_BS <- rho_JP <- 0
  }
  
  # Initialisation des Q-values
  Q <- rep(0, n_arms)
  log_lik <- 0
  
  for (t in 1:n_trials) {
    choice <- data$choice_idx[t]
    reward <- data$button_value[t]
    
    # Calcul des valeurs subjectives V(t)
    V <- lambda * Q
    
    # Ajout des biais pour événements rares si le modèle le permet
    if (model_config$has_rho) {
      # Identification des options susceptibles de produire BS/JP
      # antifragile (1) = JP possible, fragile (2) = BS possible
      # vulnerable (4) = BS et JP possibles
      V[1] <- V[1] + rho_JP  # antifragile
      V[2] <- V[2] + rho_BS  # fragile
      V[4] <- V[4] + rho_BS + rho_JP  # vulnerable
    }
    
    # Softmax
    V_max <- max(V)
    exp_V <- exp(V - V_max)
    probs <- exp_V / sum(exp_V)
    probs <- pmax(probs, 1e-10)
    probs <- probs / sum(probs)
    
    # Log-likelihood
    log_lik <- log_lik + log(probs[choice])
    
    # Mise à jour Q-learning
    Q_new <- Q
    
    # Choix de l'alpha approprié
    if (reward == -3000) {
      alpha_used <- alpha_BS
    } else if (reward == 3000) {
      alpha_used <- alpha_JP
    } else if (reward < 0) {
      alpha_used <- alpha_loss
    } else {
      alpha_used <- alpha_gain
    }
    
    # Option choisie : Q(t+1) = Q(t) + alpha * (r(t) - Q(t))
    Q_new[choice] <- Q[choice] + alpha_used * (reward - Q[choice])
    
    # Options non choisies : Q(t+1) = Q(t) * (1 - f)
    not_chosen <- setdiff(1:n_arms, choice)
    Q_new[not_chosen] <- Q[not_chosen] * (1 - forget[not_chosen])
    
    Q <- Q_new
  }
  
  if (return_negLL) {
    return(-log_lik)
  } else {
    return(log_lik)
  }
}

# ============================================================================
# CONFIGURATIONS DES MODÈLES EMBOÎTÉS
# ============================================================================

get_model_configs <- function() {
  list(
    HOMOGENEOUS = list(
      name = "HOMOGENEOUS",
      n_alpha = 1,
      n_forget = 1,
      n_lambda = 1,
      has_rho = FALSE,
      n_params = 3,
      param_names = c("alpha", "forget", "lambda"),
      lower = c(-5, -5, -3),
      upper = c(5, 5, 3)
    ),
    
    GAIN_LOSS = list(
      name = "GAIN_LOSS",
      n_alpha = 2,
      n_forget = 1,
      n_lambda = 1,
      has_rho = FALSE,
      n_params = 4,
      param_names = c("alpha_loss", "alpha_gain", "forget", "lambda"),
      lower = c(-5, -5, -5, -3),
      upper = c(5, 5, 5, 3)
    ),
    
    BIASED = list(
      name = "BIASED",
      n_alpha = 2,
      n_forget = 4,
      n_lambda = 4,
      has_rho = FALSE,
      n_params = 10,
      param_names = c("alpha_loss", "alpha_gain", 
                      "forget_1", "forget_2", "forget_3", "forget_4",
                      "lambda_1", "lambda_2", "lambda_3", "lambda_4"),
      lower = c(-5, -5, rep(-5, 4), rep(-3, 4)),
      upper = c(5, 5, rep(5, 4), rep(3, 4))
    ),
    
    REE_BIASED_SIMPLE = list(
      name = "REE_BIASED_SIMPLE",
      n_alpha = 2,
      n_forget = 1,
      n_lambda = 1,
      has_rho = TRUE,
      n_params = 6,
      param_names = c("alpha_loss", "alpha_gain", "forget", "lambda", 
                      "rho_BS", "rho_JP"),
      lower = c(-5, -5, -5, -3, -10, -10),
      upper = c(5, 5, 5, 3, 10, 10)
    ),
    
    REE_BIASED_COMPLEX = list(
      name = "REE_BIASED_COMPLEX",
      n_alpha = 2,
      n_forget = 4,
      n_lambda = 4,
      has_rho = TRUE,
      n_params = 12,
      param_names = c("alpha_loss", "alpha_gain",
                      "forget_1", "forget_2", "forget_3", "forget_4",
                      "lambda_1", "lambda_2", "lambda_3", "lambda_4",
                      "rho_BS", "rho_JP"),
      lower = c(-5, -5, rep(-5, 4), rep(-3, 4), -10, -10),
      upper = c(5, 5, rep(5, 4), rep(3, 4), 10, 10)
    ),
    
    REE_LEARNING_SIMPLE = list(
      name = "REE_LEARNING_SIMPLE",
      n_alpha = 4,
      n_forget = 1,
      n_lambda = 1,
      has_rho = FALSE,
      n_params = 6,
      param_names = c("alpha_loss", "alpha_gain", "alpha_BS", "alpha_JP",
                      "forget", "lambda"),
      lower = c(-5, -5, -5, -5, -5, -3),
      upper = c(5, 5, 5, 5, 5, 3)
    ),
    
    REE_LEARNING_COMPLEX = list(
      name = "REE_LEARNING_COMPLEX",
      n_alpha = 4,
      n_forget = 4,
      n_lambda = 4,
      has_rho = FALSE,
      n_params = 12,
      param_names = c("alpha_loss", "alpha_gain", "alpha_BS", "alpha_JP",
                      "forget_1", "forget_2", "forget_3", "forget_4",
                      "lambda_1", "lambda_2", "lambda_3", "lambda_4"),
      lower = c(-5, -5, -5, -5, rep(-5, 4), rep(-3, 4)),
      upper = c(5, 5, 5, 5, rep(5, 4), rep(3, 4))
    ),
    
    REE_LEARNING_BIASED_SIMPLE = list(
      name = "REE_LEARNING_BIASED_SIMPLE",
      n_alpha = 4,
      n_forget = 1,
      n_lambda = 1,
      has_rho = TRUE,
      n_params = 8,
      param_names = c("alpha_loss", "alpha_gain", "alpha_BS", "alpha_JP",
                      "forget", "lambda", "rho_BS", "rho_JP"),
      lower = c(-5, -5, -5, -5, -5, -3, -10, -10),
      upper = c(5, 5, 5, 5, 5, 3, 10, 10)
    ),
    
    REE_LEARNING_BIASED_COMPLEX = list(
      name = "REE_LEARNING_BIASED_COMPLEX",
      n_alpha = 4,
      n_forget = 4,
      n_lambda = 4,
      has_rho = TRUE,
      n_params = 14,
      param_names = c("alpha_loss", "alpha_gain", "alpha_BS", "alpha_JP",
                      "forget_1", "forget_2", "forget_3", "forget_4",
                      "lambda_1", "lambda_2", "lambda_3", "lambda_4",
                      "rho_BS", "rho_JP"),
      lower = c(-5, -5, -5, -5, rep(-5, 4), rep(-3, 4), -10, -10),
      upper = c(5, 5, 5, 5, rep(5, 4), rep(3, 4), 10, 10)
    )
  )
}

# ============================================================================
# ESTIMATION POUR UN PARTICIPANT
# ============================================================================

fit_participant <- function(participant_data, model_config, n_runs = 5) {
  
  all_results <- vector("list", n_runs)
  
  for (run in 1:n_runs) {
    set.seed(1000 * as.numeric(factor(model_config$name)) + run)
    
    result <- DEoptim(
      fn = qlearning_generic,
      lower = model_config$lower,
      upper = model_config$upper,
      data = participant_data,
      model_config = model_config,
      control = DEoptim.control(
        itermax = 200,
        trace = FALSE,
        parallelType = 0,
        NP = max(50, model_config$n_params * 10)
      )
    )
    
    all_results[[run]] <- list(
      params = result$optim$bestmem,
      negLL = result$optim$bestval
    )
  }
  
  # Sélection du meilleur run
  all_negLL <- sapply(all_results, function(x) x$negLL)
  best_run <- which.min(all_negLL)
  
  negLL_sd <- sd(all_negLL)
  negLL_range <- max(all_negLL) - min(all_negLL)
  
  params <- all_results[[best_run]]$params
  negLL <- all_results[[best_run]]$negLL
  
  # Calcul de la Hessienne
  hessian_result <- tryCatch({
    numDeriv::hessian(qlearning_generic, params, 
                      data = participant_data, 
                      model_config = model_config)
  }, error = function(e) NULL)
  
  hessian_positive_definite <- FALSE
  param_se <- rep(NA, model_config$n_params)
  
  if (!is.null(hessian_result)) {
    eigenvalues <- eigen(hessian_result, only.values = TRUE)$values
    hessian_positive_definite <- all(eigenvalues > 0)
    
    if (hessian_positive_definite) {
      param_vcov <- tryCatch(solve(hessian_result), error = function(e) NULL)
      if (!is.null(param_vcov)) {
        param_se <- sqrt(diag(param_vcov))
      }
    }
  }
  
  # Création du tibble de résultats
  result_df <- tibble(
    model = model_config$name,
    n_params = model_config$n_params,
    negLL = negLL,
    AIC = 2 * negLL + 2 * model_config$n_params,
    BIC = 2 * negLL + model_config$n_params * log(nrow(participant_data)),
    convergence_sd = negLL_sd,
    convergence_range = negLL_range,
    hessian_positive_definite = hessian_positive_definite,
    converged = negLL_range < 1
  )
  
  # Ajout des paramètres estimés
  for (i in 1:model_config$n_params) {
    result_df[[model_config$param_names[i]]] <- params[i]
    result_df[[paste0("se_", model_config$param_names[i])]] <- param_se[i]
  }
  
  return(result_df)
}

# ============================================================================
# ESTIMATION POUR TOUS LES PARTICIPANTS
# ============================================================================

fit_all_participants_all_models <- function(data, models_to_fit = NULL) {
  
  model_configs <- get_model_configs()
  
  if (!is.null(models_to_fit)) {
    model_configs <- model_configs[models_to_fit]
  }
  
  participants <- unique(data$participant_id)
  
  all_results <- list()
  
  for (model_name in names(model_configs)) {
    cat("\n=== Fitting model:", model_name, "===\n")
    
    model_config <- model_configs[[model_name]]
    
    model_results <- map_df(participants, function(pid) {
      cat("  Participant", pid, "\n")
      
      participant_data <- data %>% 
        filter(participant_id == pid) %>%
        arrange(trial)
      
      fit <- fit_participant(participant_data, model_config)
      fit %>% mutate(participant_id = pid, .before = 1)
    })
    
    all_results[[model_name]] <- model_results
  }
  
  return(all_results)
}

# ============================================================================
# COMPARAISON DES MODÈLES EMBOÎTÉS
# ============================================================================

compare_nested_models <- function(all_results) {
  
  # Comparaison globale
  global_comparison <- map_df(names(all_results), function(model_name) {
    results <- all_results[[model_name]]
    
    tibble(
      model = model_name,
      n_params = unique(results$n_params),
      n_converged = sum(results$converged),
      mean_negLL = mean(results$negLL),
      total_negLL = sum(results$negLL),
      total_AIC = sum(results$AIC),
      total_BIC = sum(results$BIC),
      mean_convergence_range = mean(results$convergence_range)
    )
  }) %>%
    arrange(total_BIC)
  
  cat("\n=== COMPARAISON GLOBALE DES MODÈLES ===\n")
  print(global_comparison)
  
  # Meilleur modèle par participant (BIC)
  best_models_per_participant <- map_df(names(all_results), function(model_name) {
    all_results[[model_name]] %>%
      select(participant_id, model, BIC, converged)
  }) %>%
    group_by(participant_id) %>%
    slice_min(BIC, n = 1) %>%
    ungroup()
  
  cat("\n=== MEILLEUR MODÈLE PAR PARTICIPANT (BIC) ===\n")
  print(table(best_models_per_participant$model))
  
  # Comparaison par paires de modèles emboîtés
  cat("\n=== TESTS LRT POUR MODÈLES EMBOÎTÉS ===\n")
  
  # Exemples de paires emboîtées
  nested_pairs <- list(
    c("HOMOGENEOUS", "GAIN_LOSS"),
    c("GAIN_LOSS", "BIASED"),
    c("GAIN_LOSS", "REE_BIASED_SIMPLE"),
    c("REE_BIASED_SIMPLE", "REE_BIASED_COMPLEX"),
    c("GAIN_LOSS", "REE_LEARNING_SIMPLE"),
    c("REE_LEARNING_SIMPLE", "REE_LEARNING_COMPLEX"),
    c("REE_LEARNING_SIMPLE", "REE_LEARNING_BIASED_SIMPLE"),
    c("REE_LEARNING_BIASED_SIMPLE", "REE_LEARNING_BIASED_COMPLEX")
  )
  
  lrt_results <- map_df(nested_pairs, function(pair) {
    simple_model <- pair[1]
    complex_model <- pair[2]
    
    if (simple_model %in% names(all_results) && complex_model %in% names(all_results)) {
      simple_res <- all_results[[simple_model]]
      complex_res <- all_results[[complex_model]]
      
      # LRT par participant
      lrt_df <- simple_res %>%
        select(participant_id, negLL_simple = negLL, converged_simple = converged) %>%
        left_join(
          complex_res %>% select(participant_id, negLL_complex = negLL, converged_complex = converged),
          by = "participant_id"
        ) %>%
        filter(converged_simple, converged_complex) %>%
        mutate(
          LR_stat = 2 * (negLL_simple - negLL_complex),
          df_diff = unique(complex_res$n_params) - unique(simple_res$n_params),
          p_value = pchisq(LR_stat, df = df_diff, lower.tail = FALSE),
          significant = p_value < 0.05
        )
      
      tibble(
        simple_model = simple_model,
        complex_model = complex_model,
        n_participants = nrow(lrt_df),
        pct_significant = mean(lrt_df$significant) * 100,
        mean_LR = mean(lrt_df$LR_stat),
        median_p = median(lrt_df$p_value)
      )
    }
  })
  
  print(lrt_results)
  
  list(
    global_comparison = global_comparison,
    best_models_per_participant = best_models_per_participant,
    lrt_results = lrt_results,
    all_results = all_results
  )
}

# ============================================================================
# VISUALISATION
# ============================================================================

plot_model_comparison <- function(comparison_results) {
  require(ggplot2)
  require(patchwork)
  
  # 1. Comparaison BIC globale
  p1 <- comparison_results$global_comparison %>%
    mutate(model = fct_reorder(model, total_BIC)) %>%
    ggplot(aes(x = model, y = total_BIC, fill = model)) +
    geom_col() +
    coord_flip() +
    theme_minimal() +
    labs(title = "Comparaison globale des modèles (BIC)",
         subtitle = "Plus bas = meilleur") +
    theme(legend.position = "none")
  
  # 2. Meilleur modèle par participant
  p2 <- comparison_results$best_models_per_participant %>%
    count(model) %>%
    mutate(model = fct_reorder(model, n)) %>%
    ggplot(aes(x = model, y = n, fill = model)) +
    geom_col() +
    coord_flip() +
    theme_minimal() +
    labs(title = "Meilleur modèle par participant",
         y = "Nombre de participants") +
    theme(legend.position = "none")
  
  # 3. Tests LRT
  if (nrow(comparison_results$lrt_results) > 0) {
    p3 <- comparison_results$lrt_results %>%
      mutate(comparison = paste(simple_model, "→", complex_model)) %>%
      ggplot(aes(x = fct_reorder(comparison, pct_significant), 
                 y = pct_significant)) +
      geom_col(fill = "steelblue") +
      geom_hline(yintercept = 50, linetype = "dashed", color = "red") +
      coord_flip() +
      theme_minimal() +
      labs(title = "Tests LRT entre modèles emboîtés",
           y = "% participants avec p < 0.05",
           x = "Comparaison")
  } else {
    p3 <- ggplot() + theme_void()
  }
  
  # 4. Convergence par modèle
  p4 <- map_df(names(comparison_results$all_results), function(model_name) {
    comparison_results$all_results[[model_name]] %>%
      select(model, convergence_range, converged)
  }) %>%
    ggplot(aes(x = model, y = convergence_range, fill = converged)) +
    geom_boxplot() +
    scale_y_log10() +
    coord_flip() +
    theme_minimal() +
    labs(title = "Convergence par modèle",
         y = "Range negLL (log scale)")
  
  (p1 | p2) / (p3 | p4)
}

# ============================================================================
# EXEMPLE D'UTILISATION
# ============================================================================

# Charger vos données
# data <- read_csv("votre_fichier.csv")
# Colonnes requises: participant_id, trial, choice, reward

# Estimation de tous les modèles
# all_results <- fit_all_participants_all_models(data)

# Ou seulement certains modèles
# all_results <- fit_all_participants_all_models(
#   data, 
#   models_to_fit = c("HOMOGENEOUS", "GAIN_LOSS", "REE_BIASED_SIMPLE", 
#                     "REE_LEARNING_SIMPLE", "REE_LEARNING_BIASED_SIMPLE")
# )

# Comparaison
# comparison <- compare_nested_models(all_results)

# Visualisation
# plot_model_comparison(comparison)

# Sauvegarder les résultats
# for (model_name in names(all_results)) {
#   write_csv(all_results[[model_name]], 
#             paste0("results_", model_name, ".csv"))
# }
# write_csv(comparison$global_comparison, "global_comparison.csv")
# write_csv(comparison$best_models_per_participant, "best_models.csv")