library(ggplot2)
library(tidyverse)


collapse_numeric <- function(data, orig_column_name, new_column_name, num_bins) {
  cut_points <- quantile(data[[orig_column_name]],
                         probs = seq(0, 1, length.out = num_bins + 1), na.rm = TRUE)
  data %>% mutate(!!new_column_name := cut(.data[[orig_column_name]], breaks=cut_points, right=FALSE, include.lowest = TRUE))
}


add_iter_level_variables <- function(user_iter_df, num_score_bins) {
  user_iter_df$dist_to_user_centroid <- user_iter_df$dist_to_user_centroid_junkFiltered
  user_iter_df <- user_iter_df %>%
    mutate(across(c(flag_topical_transition_tsr_userlvl, flag_topical_transition_cosine_userlvl,
                    flag_topical_transition_tsr_tasklvl, flag_topical_transition_cosine_tasklvl),
           ~as.factor(.x)))
                  
  # add z-score or performance within image-iteration
  user_iter_df <- user_iter_df %>% group_by(referenceImageName, iteration) %>%
    mutate(z_cosine_sim_mean = scale(cosine_sim_mean)[,1],
           z_dream_sim_mean = scale(dream_sim_mean)[,1],
           pct_rank_cosine_sim_mean = round(cume_dist(cosine_sim_mean)*100),
           pct_rank_dream_sim_mean = round(cume_dist(dream_sim_mean)*100))
  
  # add diff between iterations, prev iteration outcomes, etc..
  user_iter_df <- user_iter_df %>%
    group_by(referenceImageName, userId) %>%
    arrange(iteration) %>%
    mutate(prev_iteration = lag(iteration),
           cosine_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(cosine_sim_mean), NA),
           cosine_sim_observed_z_prev = ifelse(prev_iteration==iteration-1, lag(cosine_sim_observed_z), NA),
           z_cosine_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(z_cosine_sim_mean), NA),
           pct_rank_cosine_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(pct_rank_cosine_sim_mean), NA),
           dream_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(dream_sim_mean), NA),
           dream_sim_observed_z_prev = ifelse(prev_iteration==iteration-1, lag(dream_sim_observed_z), NA),
           z_dream_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(z_dream_sim_mean), NA),
           pct_rank_dream_sim_mean_prev = ifelse(prev_iteration==iteration-1, lag(pct_rank_dream_sim_mean), NA),
           prev_prompt_length_words = ifelse(prev_iteration==iteration-1, lag(prompt_length_words), NA),
           prev_prompt_length_chars = ifelse(prev_iteration==iteration-1, lag(prompt_length_chars), NA),
           prompt_length_words_diff = prompt_length_words - prev_prompt_length_words,
           prompt_length_chars_diff = prompt_length_chars - prev_prompt_length_chars) %>%
    ungroup()
  
  # convert observed z scores to brackets
  user_iter_df <- user_iter_df %>%
    group_by(referenceImageName, iteration) %>%
    mutate(cosine_sim_observed_z_prev_bracket=as.factor(cut(cosine_sim_observed_z_prev, breaks=c(-Inf, -0.45, 0.45, Inf))),
           dream_sim_observed_z_prev_bracket=as.factor(cut(dream_sim_observed_z_prev, breaks=c(-Inf, -0.45, 0.45, Inf)))) %>%
    ungroup() %>%
    group_by(cosine_sim_observed_z_prev_bracket) %>%
    mutate(cosine_sim_observed_z_prev_bracket_mean=as.factor(round(mean(cosine_sim_observed_z_prev, na.rm=T), 2))) %>%
    ungroup() %>%
    group_by(dream_sim_observed_z_prev_bracket) %>%
    mutate(dream_sim_observed_z_prev_bracket_mean=as.factor(round(mean(dream_sim_observed_z_prev, na.rm=T), 2))) %>%
    ungroup()

  user_iter_df <- collapse_numeric(user_iter_df, "prompt_length_words", "prompt_length_words4", 4)
  factor_cols <- c("referenceImageName", "userId", "referenceImageCategory", "dalle_version")
  user_iter_df <- user_iter_df %>% mutate(across(all_of(factor_cols), ~factor(., ordered = FALSE)))
  return(user_iter_df)
}


prepare_user_data <- function(user_iter_df) {
  user_df <- user_iter_df %>%
    group_by(referenceImageName, userId) %>%
    summarise(
      max_cosine_sim_mean = max(cosine_sim_mean),
      mean_cosine_sim_mean = mean(cosine_sim_mean),
      mean_cosine_sim_observed_z = mean(cosine_sim_observed_z),
      iteration_cosine_max = which.max(cosine_sim_mean),
      
      max_dream_sim_mean = max(dream_sim_mean),
      mean_dream_sim_mean = mean(dream_sim_mean),
      mean_dream_sim_observed_z = mean(dream_sim_observed_z),
      iteration_dream_max = which.max(dream_sim_mean),
      
      dalle_version=first(dalle_version),
      unique_versions=n_distinct(dalle_version),
    ) %>% 
    ungroup()
  assertthat::assert_that(all(user_df$unique_versions == 1))
  return(user_df)
}

prepare_nlp_user_agg_data <- function(user_iter_df) {
  # aggregate nlp features by user across iterations
  nlp_df <- user_iter_df %>%
    group_by(referenceImageName, userId) %>%
    summarise(embedding_var=var(dist_to_user_centroid),
              
              n_total_topics_tsr_tasklvl_mean=first(n_total_topics_tsr_tasklvl),
              unique_n_total_topics_tsr_tasklvl_mean=n_distinct(n_total_topics_tsr_tasklvl),
              n_total_topics_tsr_userlvl_mean=first(n_total_topics_tsr_userlvl),
              unique_n_total_topics_tsr_userlvl_mean=n_distinct(n_total_topics_tsr_userlvl),
              n_total_topics_cosine_tasklvl_mean=first(n_total_topics_cosine_tasklvl),
              unique_n_total_topics_cosine_tasklvl_mean=n_distinct(n_total_topics_cosine_tasklvl),
              n_total_topics_cosine_userlvl_mean=first(n_total_topics_cosine_userlvl),
              unique_n_total_topics_cosine_userlvl_mean=n_distinct(n_total_topics_cosine_userlvl),
              
              mean_tsr_prev_prompt=mean(tsr_prev_prompt, na.rm=TRUE),
              mean_cosine_sim_prev_prompt=mean(cosine_sim_prev_prompt, na.rm=TRUE),
              mean_contains_prev_prompt=mean(contains_prev_prompt, na.rm=TRUE),
              mean_prompt_length_chars=mean(prompt_length_chars),
              mean_prompt_length_words=mean(prompt_length_words)) %>%
    ungroup()
  
  nlp_df$log_embedding_var <- log(nlp_df$embedding_var)
  nlp_df <- collapse_numeric(nlp_df, "mean_prompt_length_chars", "mean_prompt_length_chars4", 4)
  nlp_df <- collapse_numeric(nlp_df, "mean_prompt_length_words", "mean_prompt_length_words4", 4)
  # ensure there was a single topical transition value per user
  assertthat::assert_that(all(nlp_df$unique_n_total_topics_tsr_tasklvl_mean == 1))
  assertthat::assert_that(all(nlp_df$unique_n_total_topics_tsr_userlvl_mean == 1))
  assertthat::assert_that(all(nlp_df$unique_n_total_topics_cosine_tasklvl_mean == 1))
  assertthat::assert_that(all(nlp_df$unique_n_total_topics_cosine_userlvl_mean == 1))
  assertthat::assert_that(!any(is.na(nlp_df)))
  nlp_df <- nlp_df %>%
    select(-c(unique_n_total_topics_tsr_tasklvl_mean, unique_n_total_topics_tsr_userlvl_mean,
              unique_n_total_topics_cosine_tasklvl_mean, unique_n_total_topics_cosine_userlvl_mean))
              
  return(nlp_df)
}

prepare_survey_data <- function(user_iter_df) {
  skill_vars <- c("Reading Comprehension", "Active Listening", "Writing", "Speaking", "Critical Thinking",                       
                  "Social Perceptiveness", "Coordination", "Instructing", "Programming",
                  "Judgment and Decision Making", "Systems Evaluation", "Science", "Active Learning",
                  "Learning Strategies", "Monitoring", "Complex Problem Solving", "Operations Analysis",
                  "Technology Design", "Troubleshooting", "Quality Control Analysis", "Systems Analysis")
  srvy_vars <- c('programming_skill', 'program_frequency', 'GPT_usage', 'GenAI_outlook', 'degree',
                  'imagery_skill', 'write_image_use', 'write_instr_usage', 'instruction_skill',
                  'age')
  srvy_vars_cat <- c('programming_skill3', 'program_frequency3', 'GPT_usage4', 'GenAI_outlook3', 'degree3',
                     'imagery_skill3', 'write_image_use3', 'write_instr_usage3', 'instruction_skill3',
                     'age5')
  all_cols <- c(skill_vars, srvy_vars, srvy_vars_cat, 'gender')
  survey_df <- user_iter_df %>% group_by(userId) %>% summarise(across(all_of(all_cols), ~ first(na.omit(.x))))
  return(survey_df)
}

prepare_first_score_cols <- function(user_iter_df) {
  # add binned first score to each user-image
  score_cols <- c("cosine_sim_mean", "dream_sim_mean",
                  "z_cosine_sim_mean", "z_dream_sim_mean",
                  "pct_rank_cosine_sim_mean", "pct_rank_dream_sim_mean",
                  "cosine_sim_observed_z", "dream_sim_observed_z")
  
  first_score_info <- user_iter_df %>%
    select(userId, referenceImageName, iteration, all_of(score_cols)) %>% 
    group_by(userId) %>%
    arrange(iteration) %>%
    summarise(across(all_of(score_cols), ~ ifelse(min(iteration)==0, first(.x), NA), .names='first_{.col}'),
              referenceImageName=first(referenceImageName)) %>%
    ungroup() %>%
    group_by(referenceImageName) %>%
    mutate(first_cosine_sim_mean_bracket=as.factor(cut(first_cosine_sim_mean, breaks=c(-Inf, -0.45, 0.45, Inf))),
           first_dream_sim_mean_bracket=as.factor(cut(first_dream_sim_mean, breaks=c(-Inf, -0.45, 0.45, Inf))),
           first_cosine_sim_observed_z_bracket=as.factor(cut(first_cosine_sim_observed_z, breaks=c(-Inf, -0.45, 0.45, Inf))),
           first_dream_sim_observed_z_bracket=as.factor(cut(first_dream_sim_observed_z, breaks=c(-Inf, -0.45, 0.45, Inf)))) %>%
    ungroup() %>%
    group_by(first_cosine_sim_mean_bracket) %>%
    mutate(first_cosine_sim_mean_bracket_mean=as.factor(round(mean(first_cosine_sim_mean), 2))) %>%
    ungroup() %>%
    group_by(first_dream_sim_mean_bracket) %>%
    mutate(first_dream_sim_mean_bracket_mean=as.factor(round(mean(first_dream_sim_mean), 2))) %>%
    ungroup() %>%
    group_by(first_cosine_sim_observed_z_bracket) %>%
    mutate(first_cosine_sim_observed_z_bracket_mean=as.factor(round(mean(first_cosine_sim_observed_z), 2))) %>%
    ungroup() %>%
    group_by(first_dream_sim_observed_z_bracket) %>%
    mutate(first_dream_sim_observed_z_bracket_mean=as.factor(round(mean(first_dream_sim_observed_z), 2)))
  
  
  # set mid level as base level for first_cosine_sim_bracket
  first_score_info$first_cosine_sim_mean_bracket <- relevel(first_score_info$first_cosine_sim_mean_bracket, ref = 2)
  first_score_info$first_dream_sim_mean_bracket <- relevel(first_score_info$first_dream_sim_mean_bracket, ref = 2)
  
  duplicate_rows <- duplicated(first_score_info[c('referenceImageName', 'userId')])
  if(any(duplicate_rows)) {
    cat('Error: duplicate rows with the same refImage and user in first iteration')
  }
  return(first_score_info)
}


add_user_level_variables <- function(user_df) {
  # convert observed mean z scores to brackets
  user_df <- user_df %>%
    group_by(referenceImageName) %>%
    mutate(mean_cosine_sim_observed_z_bracket=as.factor(cut(mean_cosine_sim_observed_z, breaks=c(-Inf, -0.45, 0.45, Inf))),
           mean_dream_sim_observed_z_bracket=as.factor(cut(mean_dream_sim_observed_z, breaks=c(-Inf, -0.45, 0.45, Inf)))) %>%
    ungroup() %>%
    group_by(mean_cosine_sim_observed_z_bracket) %>%
    mutate(mean_cosine_sim_observed_z_bracket_mean=as.factor(round(mean(mean_cosine_sim_observed_z, na.rm=T), 2))) %>%
    ungroup() %>%
    group_by(mean_dream_sim_observed_z_bracket) %>%
    mutate(mean_dream_sim_observed_z_bracket_mean=as.factor(round(mean(mean_dream_sim_observed_z, na.rm=T), 2))) %>%
    ungroup()
  
  # add z-score of cosine sim and dream sim scores among all other subjects
  user_df <- user_df %>%
    group_by(referenceImageName) %>%
    mutate(z_mean_cosine_sim_mean = scale(mean_cosine_sim_mean)[,1],
           z_mean_dream_sim_mean = scale(mean_dream_sim_mean)[,1],
           pct_rank_mean_cosine_sim_mean = round(cume_dist(mean_cosine_sim_mean)*100),
           pct_rank_max_cosine_sim_mean = round(cume_dist(max_cosine_sim_mean)*100),
           pct_rank_first_cosine_sim_mean = round(cume_dist(first_cosine_sim_mean)*100),
           
           pct_rank_mean_dream_sim_mean = round(cume_dist(mean_dream_sim_mean)*100),
           pct_rank_max_dream_sim_mean = round(cume_dist(max_dream_sim_mean)*100),
           pct_rank_first_dream_sim_mean = round(cume_dist(first_dream_sim_mean)*100)) %>%
    ungroup()
  
  # convert cols to factor
  factor_cols <- c("referenceImageName", "userId")
  user_df <- user_df %>% mutate(across(all_of(factor_cols), ~factor(., ordered = FALSE)))
  return(user_df)
  
}

generate_data_from_user_iter <- function(user_iter_df) {
  user_iter_df <- add_iter_level_variables(user_iter_df, num_score_bins=3)
  # data per user: generate all user level features from base user_iter_df
  user_df <- prepare_user_data(user_iter_df)
  nlp_df <- prepare_nlp_user_agg_data(user_iter_df)
  survey_df <- prepare_survey_data(user_iter_df)
  first_score_info <- prepare_first_score_cols(user_iter_df)
  
  cat(paste('Original size of user_df:', nrow(user_df), '\n'))
  user_df <- inner_join(user_df, nlp_df, by=c("userId", "referenceImageName"))
  cat(paste('Size of user_df after merging with nlp data:', nrow(user_df), '\n')) 
  user_df <- inner_join(user_df, survey_df, by=c("userId"))
  cat(paste('Size of user_df after merging with survey data:', nrow(user_df), '\n'))
  user_df <- user_df %>% inner_join(first_score_info, by=c("referenceImageName", "userId"))
  cat(paste('Size of user_df after merging with first score data:', nrow(user_df), '\n')) 
  user_df <- add_user_level_variables(user_df)
  cat(paste('Size of user_df after adding extra variables:', nrow(user_df), '\n'))
  return(list(user_iter_df=user_iter_df, user_df=user_df))
}
generate_data_from_user_iter_file <- function(user_iter_file) {
  user_iter_df <- read_csv(user_iter_file) %>% filter(junk_flag==FALSE)
  return(generate_data_from_user_iter(user_iter_df))
}