library(caTools)
library(fsbrain)

raw_data <- read.csv(paste(PATH, "data/3Lmap.csv", sep=""))

#####################
# Data processing
#####################

# create a variable called "Not_Country_UK"
raw_data$Not_Country_UK <- 1 - raw_data$UK

# Compute the disutility from the utility (EQ5DINDEX)
raw_data$Disutility <- 1 - raw_data$EQ5DIndex3L
raw_data$EQ5DINDEX <- raw_data$EQ5DIndex3L
raw_data$Disutility[raw_data$Disutility < 0] <- 0.000001

# create variables "treated" and "diagnosed"
raw_data$diagnosed[raw_data$Group_Dx_Unclean %in% c(1, 2)] <- 1
raw_data$diagnosed[raw_data$Group_Dx_Unclean %in% c(3, 4)] <- 0
raw_data$treated[raw_data$Group_Dx_Unclean %in% c(1, 3)] <- 1
raw_data$treated[raw_data$Group_Dx_Unclean %in% c(2, 4)] <- 0

targets <- c("Disutility", "EQ5DINDEX")

factor_features <- c(
  'female',
  'married', 
  'degree_4years',
  'employed',
  'retired',
  'current_smoker',
  'former_smoker',
  'heavy_drinker',
  'low_moderate_drinker', 
  'Combined_DP_PTS_AX',
  'pain',
  'treated',
  'diagnosed',
  'Not_Country_UK'
)

cont_features <-  c(
  'ISI_Score', 
  'deage', 
  'BMI_R',
  'CCI'
)

features <- c(factor_features, cont_features)
columns_to_keep <- c(features, targets)

processed_df <- raw_data %>% select(one_of(columns_to_keep))

# Filling missing data (BMI_R) with average value
processed_df_std <- processed_df
processed_df_std$BMI_R[is.na(processed_df_std$BMI_R)] <- mean(processed_df_std$BMI_R, na.rm = TRUE)

# standardisation of the full data
processed_df_std <- processed_df_std %>% mutate_at(cont_features, ~scale(.) %>% as.vector())

# create variable "treated_ISI_Score"
processed_df_std$treated_ISI_Score <- processed_df_std$treated * processed_df_std$ISI_Score

# dump the processed data and standardised data
write.csv(processed_df_std, paste(DATA_PATH, "full_dataset_std.csv", sep="/"), row.names=FALSE, na='')


#####################
# Data splitting
#####################

# perform the train/ test split for the 10 datasets with 50-50

NUM_RUNS <- 100
RUNS <- c(1:NUM_RUNS)

for (RUN in RUNS) {
  # make this example reproducible
  set.seed(as.numeric(RUN))
  sample <- sample.split(processed_df, SplitRatio = 0.5)
  train  <- subset(processed_df, sample == TRUE)
  test   <- subset(processed_df, sample == FALSE)
  
  # Filling missing data (BMI_R) with average value
  train$BMI_R[is.na(train$BMI_R)] <- mean(train$BMI_R, na.rm = TRUE)
  test$BMI_R[is.na(test$BMI_R)] <- mean(test$BMI_R, na.rm = TRUE)
  
  # standardisation
  train <- train %>% mutate_at(cont_features, ~scale(.) %>% as.vector())
  test <- test %>% mutate_at(cont_features, ~scale(.) %>% as.vector())
  
  train$treated_ISI_Score <- train$treated * train$ISI_Score
  test$treated_ISI_Score <- test$treated * test$ISI_Score
  
  write.csv(train, paste(DATA_PATH, "/runs/splits_", NUM_RUNS, "/train_", RUN, ".csv", sep=""), row.names=FALSE, na='')
  write.csv(test, paste(DATA_PATH, "/runs/splits_", NUM_RUNS, "/test_", RUN, ".csv", sep=""), row.names=FALSE, na='')
}