# Script: app_ono.R
# setwd("~/Dropbox/OptimizingSI/Analysis/ono")

options(error = NULL)
library(shiny)
library(ggplot2)
library(strategize)
library(dplyr)

# Custom plotting function for optimal strategy distributions
plot_factor <- function(pi_star_list, pi_star_se_list, factor_name, zStar = 1.96) {
  probs <- lapply(pi_star_list, function(x) x[[factor_name]])
  ses <- lapply(pi_star_se_list, function(x) x[[factor_name]])
  levels <- names(probs[[1]])
  n_strategies <- length(probs)
  
  # Create data frame for plotting
  df <- do.call(rbind, lapply(1:n_strategies, function(i) {
    data.frame(
      Strategy = if (n_strategies == 1) "Optimal" else c("Democrat", "Republican")[i],
      Level = levels,
      Probability = probs[[i]],
      SE = ses[[i]]
    )
  }))
  
  # Plot with ggplot2
  p <- ggplot(df, aes(x = Level, y = Probability, fill = Strategy)) +
    geom_bar(stat = "identity", position = position_dodge(width = 0.9), width = 0.8) +
    geom_errorbar(aes(ymin = Probability - zStar * SE, ymax = Probability + zStar * SE),
                  position = position_dodge(width = 0.9), width = 0.25) +
    labs(title = paste("Optimal Distribution for", factor_name),
         x = "Level", y = "Probability") +
    theme_minimal() +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "top") +
    scale_fill_manual(values = c("Democrat" = "#89cff0", "Republican" = "red", "Optimal" = "black"))
  
  return(p)
}

# UI Definition
ui <- fluidPage(
  titlePanel("Exploring strategize with the candidate choice conjoint data"),
  
  sidebarLayout(
    sidebarPanel(
      h4("Analysis Options"),
      radioButtons("case_type", "Case Type:",
                   choices = c("Average", "Adversarial"),
                   selected = "Average"),
      conditionalPanel(
        condition = "input.case_type == 'Average'",
        selectInput("respondent_group", "Respondent Group:",
                    choices = c("All", "Democrat", "Independent", "Republican"),
                    selected = "All")
      ),
      # Add a single numeric input for lambda
      numericInput("lambda_input", "Lambda (regularization):", 
                   value = 0.01, min = 1e-6, max = 10, step = 0.01),
      actionButton("compute", "Compute Results", class = "btn-primary"),
      hr(),
      h4("Visualization"),
      selectInput("factor", "Select Factor to Display:",
                  choices = NULL),
      hr(),
      h5("Instructions:"),
      p("1. Select a case type and, for Average case, a respondent group."),
      p("2. Specify the single lambda to be used by strategize."),
      p("3. Click 'Compute Results' to generate optimal strategies."),
      p("4. Choose a factor to view its distribution.")
    ),
    
    mainPanel(
      tabsetPanel(
        tabPanel("Optimal Strategy Plot",
                 plotOutput("strategy_plot", height = "600px")),
        tabPanel("Q Value",
                 verbatimTextOutput("q_value"),
                 p("Q represents the estimated outcome (e.g., selection probability) under the optimal strategy, with 95% confidence interval.")),
        tabPanel("About",
                 h3("About This App"),
                 p("This Shiny app explores the `strategize` package using Ono experimental data. It computes optimal strategies for Average (optimizing for a respondent group) and Adversarial (optimizing for both parties in competition) cases on the fly."),
                 p("**Average Case**: Optimizes candidate characteristics for a selected respondent group."),
                 p("**Adversarial Case**: Finds equilibrium strategies for Democrats and Republicans, identified by 'Pro-life' stance.")
        )
      )
    )
  )
)

# Server Definition
server <- function(input, output, session) {
  # Load data
  load("Processed_OnoData.RData")
  Primary2016 <- read.csv("PrimaryCandidates2016 - Sheet1.csv")
  
  # Update factor choices dynamically
  observe({
    if (input$case_type == "Average") {
      factors <- colnames(FACTOR_MAT_FULL)[!colnames(FACTOR_MAT_FULL) %in% c("Office")]
    } else {
      factors <- colnames(FACTOR_MAT_FULL)[!colnames(FACTOR_MAT_FULL) %in% c("Office", "Party.affiliation", "Party.competition")]
    }
    updateSelectInput(session, "factor", choices = factors, selected = factors[1])
  })
  
  # Reactive computation triggered by button
  result <- eventReactive(input$compute, {
    withProgress(message = "Computing optimal strategies...", value = 0, {
      # Increment progress
      incProgress(0.2, detail = "Preparing data...")
      
      # Common hyperparameters (mirroring QRun_Apps.R)
      params <- list(
        nSGD = 1000L,
        batch_size = 50L,
        penalty_type = "KL",
        nFolds = 3L,
        use_optax = TRUE,
        compute_se = FALSE, # Set to FALSE for quicker results 
        conf_level = 0.95,
        conda_env = "strategize",
        conda_env_required = TRUE
      )
      
      # Grab the single user-chosen lambda
      my_lambda <- input$lambda_input
      
      if (input$case_type == "Average") {
        # Subset data for Average case
        if (input$respondent_group == "All") {
          indices <- which( my_data$Office == "President" )
        } else {
          indices <- which(my_data_FULL$R_Partisanship == input$respondent_group & 
                             my_data$Office == "President")
        }
        
        FACTOR_MAT <- FACTOR_MAT_FULL[indices, 
                      !colnames(FACTOR_MAT_FULL) %in% c("Office","Party.affiliation","Party.competition")]
        Yobs <- Yobs_FULL[indices]
        X <- X_FULL[indices, ]
        log_pr_w <- log_pr_w_FULL[indices]
        pair_id <- pair_id_FULL[indices]
        assignmentProbList <- assignmentProbList_FULL[!names(assignmentProbList_FULL) %in% "Office"]
        
        incProgress(0.4, detail = "Running strategize...")
        
        # Compute with strategize using a single lambda
        Qoptimized <- strategize(
          Y = Yobs,
          W = FACTOR_MAT,
          X = X,
          pair_id = pair_id,
          
          p_list = assignmentProbList[colnames(FACTOR_MAT)],
          lambda = my_lambda,
          diff = TRUE, 
          adversarial = FALSE,
          K = 1L, # Base analysis
          nSGD = params$nSGD,
          penalty_type = params$penalty_type,
          folds = params$nFolds,
          use_optax = params$use_optax,
          compute_se = params$compute_se,
          conf_level = params$conf_level,
          conda_env = params$conda_env,
          conda_env_required = params$conda_env_required
        )
      } else { # Adversarial case
        # Use full data, drop specific factors
        DROP_FACTORS <- c("Office", "Party.affiliation", "Party.competition")
        FACTOR_MAT <- FACTOR_MAT_FULL[, !colnames(FACTOR_MAT_FULL) %in% DROP_FACTORS]
        Yobs <- Yobs_FULL
        X <- X_FULL
        log_pr_w <- log_pr_w_FULL
        assignmentProbList <- assignmentProbList_FULL[!names(assignmentProbList_FULL) %in% DROP_FACTORS]
        
        # Prepare slate_list (simplified from QRun_Apps.R)
        incProgress(0.3, detail = "Preparing slate data...")
        FactorOptions <- apply(FACTOR_MAT, 2, table)
        prior_alpha <- 10
        Primary_D <- Primary2016[Primary2016$Party == "Democratic", colnames(FACTOR_MAT)]
        Primary_R <- Primary2016[Primary2016$Party == "Republican", colnames(FACTOR_MAT)]
        
        Primary_D_slate <- lapply(colnames(Primary_D), function(col) {
          posterior_alpha <- FactorOptions[[col]]; posterior_alpha[] <- prior_alpha
          Empirical_ <- table(Primary_D[[col]])
          Empirical_ <- Empirical_[names(Empirical_) != "Unclear"]
          posterior_alpha[names(Empirical_)] <- posterior_alpha[names(Empirical_)] + Empirical_
          prop.table(posterior_alpha)
        })
        names(Primary_D_slate) <- colnames(Primary_D)
        
        Primary_R_slate <- lapply(colnames(Primary_R), function(col) {
          posterior_alpha <- FactorOptions[[col]]; posterior_alpha[] <- prior_alpha
          Empirical_ <- table(Primary_R[[col]])
          Empirical_ <- Empirical_[names(Empirical_) != "Unclear"]
          posterior_alpha[names(Empirical_)] <- posterior_alpha[names(Empirical_)] + Empirical_
          prop.table(posterior_alpha)
        })
        names(Primary_R_slate) <- colnames(Primary_R)
        
        slate_list <- list("Democratic" = Primary_D_slate, "Republican" = Primary_R_slate)
        
        # subset data 
        indices <- which( my_data$R_Partisanship %in% c("Republican","Democrat") & 
                           my_data$Office == "President" )
        FACTOR_MAT <- FACTOR_MAT_FULL[indices, 
                                      !colnames(FACTOR_MAT_FULL) %in% c("Office",
                                                                        "Party.competition",
                                                                        "Party.affiliation")]
        Yobs <- Yobs_FULL[indices]
        my_data_red <- my_data_FULL[indices,]
        pair_id <- pair_id_FULL[indices]
        cluster_var <- cluster_var_FULL[ indices ]
        my_data_red$Party.affiliation_clean <- ifelse(my_data_red$Party.affiliation == "Republican Party",
                                                      yes = "Republican", no = ifelse(my_data_red$Party.affiliation == "Democratic Party",
                                                                                      yes = "Democrat",no = "Independent"))
        
        # subset cols 
        assignmentProbList <- assignmentProbList_FULL[colnames(FACTOR_MAT)]
        slate_list$Democratic <- slate_list$Democratic[names(assignmentProbList)]
        slate_list$Republican <- slate_list$Republican[names(assignmentProbList)]
        
        incProgress(0.4, detail = "Running strategize...")
        
        # Compute with strategize using a single lambda
        Qoptimized <- strategize(
          Y = Yobs,
          W = FACTOR_MAT,
          #X = X,
          X = NULL,
          p_list = assignmentProbList,
          slate_list = slate_list,
          varcov_cluster_variable = cluster_var,
          competing_group_variable_respondent = my_data_red$R_Partisanship,
          competing_group_variable_candidate = my_data_red$Party.affiliation_clean,
          competing_group_competition_variable_candidate = my_data_red$Party.competition,
          pair_id = pair_id,
          respondent_id = my_data_red$respondentIndex,
          respondent_task_id = my_data_red$task,
          profile_order = my_data_red$profile,
          
          lambda = my_lambda,
          diff = TRUE, 
          force_gaussian = FALSE,
          adversarial = TRUE,
          nFolds_glm = 3L, 
          K = 1L,
          nMonte_adversarial = 100L,
          nSGD = params$nSGD,
          penalty_type = params$penalty_type,
          use_optax = params$use_optax,
          compute_se = params$compute_se,
          conf_level = params$conf_level,
          conda_env = params$conda_env,
          conda_env_required = params$conda_env_required
        )
        
        # Identify Democrat vs Republican based on "Pro-life" stance
        prolife_probs <- c(Qoptimized$pi_star_point$k1$Position.on.abortion["Pro-life"],
                           Qoptimized$pi_star_point$k2$Position.on.abortion["Pro-life"])
        which_repub <- which.max(prolife_probs)
        if (which_repub == 1) {
          # Swap
          Qoptimized$pi_star_point <- list(k1 = Qoptimized$pi_star_point$k2, k2 = Qoptimized$pi_star_point$k1)
          Qoptimized$pi_star_se <- list(k1 = Qoptimized$pi_star_se$k2, k2 = Qoptimized$pi_star_se$k1)
        }
      }
      
      incProgress(0.8, detail = "Finalizing results...")
      return(Qoptimized)
    })
  })
  
  # Render strategy plot
  output$strategy_plot <- renderPlot({
    req(result())
    factor_name <- input$factor
    pi_star_list <- result()$pi_star_point
    pi_star_se_list <- result()$pi_star_se
    plot_factor(pi_star_list, pi_star_se_list, factor_name)
  })
  
  # Render Q value
  output$q_value <- renderText({
    req(result())
    q_point <- result()$Q_point_mEst
    q_se <- result()$Q_se_mEst
    paste("Estimated Q Value: ", sprintf("%.3f ± %.3f", q_point, 1.96 * q_se))
  })
}

# Run the app
shinyApp(ui, server)