library(twopartm)
library(aldvmm) 
library(readxl)
library(dplyr)

DATA_PATH = ..
RESULTS_PATH = ..

data <- read.csv(paste(DATA_PATH, "full_dataset_std.csv", sep="/")) # contains the full standardised dataset


#####################
# Data preparation
#####################

target <- "Disutility"

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, "treated_ISI_Score")
model_formula <- as.formula(paste("Disutility", paste(features, collapse = " + "), sep = " ~ "))

# turning the categorical features to factors
data_factor <- data %>% mutate_at(factor_features, ~(factor(.)))


#####################
# Models
#####################

#####################
# GLM
#####################

glm <- glm(model_formula, family = Gamma(link = "log"), data = data_factor)

results_glm <- data.frame(
  mae = calc_mae(glm$fitted.values, glm$y),
  mse = calc_mse(glm$fitted.values, glm$y),
  rsq = calc_rsq(glm$fitted.values, glm$y),
  ll = -logLik(glm),
  aic = AIC(glm),
  bic = BIC(glm)
)

# save the metrics
write.csv(results_glm, paste(RESULTS_PATH,"metrics", "glm_metrics.csv", sep="/"), row.names=FALSE, na='')

# dump the coef
write.csv(data.frame(glm$coefficients), paste(RESULTS_PATH, "coef", "glm_coef.csv", sep="/"), row.names=TRUE, na='')


# Covariance matrix
glm_cov <- vcov(glm)

order <- c(
  '(Intercept)',
  'deage',
  'female1',
  'married1', 
  'degree_4years1',
  'employed1',
  'retired1',
  'current_smoker1',
  'former_smoker1',
  'heavy_drinker1',
  'low_moderate_drinker1', 
  'BMI_R',
  'CCI',
  'Combined_DP_PTS_AX1',
  'pain1',
  'treated1',
  'diagnosed1',
  'ISI_Score',
  'treated_ISI_Score',
  'Not_Country_UK1'
)

reordered_cov <- glm_cov[order, order]
write.csv(data.frame(glm_cov), paste(RESULTS_PATH, "coef", "glm_cov.csv", sep="/"), row.names=TRUE, na='')

# Cholesky matrix
write.csv(data.frame(t(chol(glm_cov))), paste(RESULTS_PATH, "coef", "glm_cholesky.csv", sep="/"), row.names=TRUE, na='')

# dump the coef SEs
write.csv(data.frame(sqrt(diag(glm_cov))), paste(RESULTS_PATH, "coef", "glm_coef_se.csv", sep="/"), row.names=TRUE, na='')



#####################
# OLS
#####################

ols <- lm(model_formula, data = data_factor)

results_ols <- data.frame(
  mae = mean(abs(ols$residuals)),
  mse = mean(ols$residuals ^ 2),
  rsq = summary(ols)$adj.r.squared,
  ll = -logLik(ols),
  aic = AIC(ols),
  bic = BIC(ols)
)

# save the metrics
write.csv(results_ols, paste(RESULTS_PATH, "metrics", "ols_metrics.csv", sep="/"), row.names=FALSE, na='')

# dump the coef
write.csv(data.frame(ols$coefficients), paste(RESULTS_PATH, "coef", "ols_coef.csv", sep="/"), row.names=TRUE, na='')
# dump the coef SEs
write.csv(data.frame(diag(vcov(ols))), paste(RESULTS_PATH, "coef", "ols_coef_se.csv", sep="/"), row.names=TRUE, na='')



#####################
# CLAD
#####################
 
data_factor$yc <- rep(1.0,length(data_factor$EQ5DINDEX))
clad <- crq (Curv(dither(EQ5DINDEX, type = "right", value = 0.00005), yc, "right") ~ 
               ISI_Score + 
               deage + 
               female + 
               married + 
               degree_4years + 
               employed + 
               retired + 
               current_smoker + 
               former_smoker + 
               heavy_drinker + 
               low_moderate_drinker + 
               BMI_R + 
               CCI + 
               Combined_DP_PTS_AX + 
               pain + 
               treated + 
               diagnosed + 
               treated_ISI_Score + 
               Not_Country_UK, 
             data=data_factor,
             tau=0.5, 
             method="Powell"
)


results_clad <- data.frame(
  mae = mean(abs(clad$residuals)),
  mse = mean(clad$residuals ^ 2),
  rsq = calc_rsq(predict(clad, data_factor), data_factor$EQ5DINDEX)
)

# save the metrics
write.csv(results_clad, paste(RESULTS_PATH, "metrics", "clad_metrics.csv", sep="/"), row.names=FALSE, na='')

# dump the coef
write.csv(data.frame(clad$coefficients), paste(RESULTS_PATH, "coef", "clad_coef.csv", sep="/"), row.names=TRUE, na='')

# dump the coef SEs
# need to figure out how to compute them as it's not given in the package

#####################
# Two part models - NOT USED
#####################

model_formula_2part <- as.formula(paste("Disutility_scaled", paste(features, collapse = " + "), sep = " ~ "))

twopm_glm <-  tpm(formula_part1 = model_formula_2part, 
              formula_part2 = model_formula_2part,
              data = data_factor,
              link_part1 = "logit",
              family_part2 = Gamma(link = "log")
)

y <- data_factor$Disutility_scaled
y_pred <- predict(twopm_glm, data_factor)

results_twopm_glm <- data.frame(
  mae = calc_mae(y_pred, y),
  mse = calc_mse(y_pred, y),
  rsq = calc_rsq(y_pred, y)
)

twopm_ols <-  tpm(formula_part1 = model_formula_2part, 
                  formula_part2 = model_formula_2part,
                  data = data_factor,
                  link_part1 = "logit"
)

y_pred <- predict(twopm_ols, data_factor)

results_twopm_ols <- data.frame(
  mae = calc_mae(y_pred, y),
  mse = calc_mse(y_pred, y),
  rsq = calc_rsq(y_pred, y)
)

#####################
# ALDVMM
#####################

my_opt = "BFGS"  # possibilities: "Nelder-Mead", "BFGS", "CG", "L-BFGS-B", "nlminb", "Rcgmin", "Rvmmin"
my_init = "constant" # possibilities: "zero", "random", "constant", "sann"

target <- "EQ5DINDEX"
model_formula_aldvmm <- as.formula(paste(target, paste(features, collapse = " + "), sep = " ~ "))

# fitting ALDVMMM with 3 components with 5 different seeds
NUM_COMP <- c(3) 
NUM_SEED <- c(1:5) 

for (num_comp in NUM_COMP){
  for (num_seed in NUM_SEED){
  set.seed(num_seed)
    print(num_seed)
  res_file_name = paste("aldvmm_", num_comp, "_", my_opt, "_", my_init, ".csv", sep="")
  
  fit <- aldvmm(formula = model_formula_aldvmm,
                data = data_factor,
                psi = c(-0.594, 1),
                ncmp = num_comp,
                init.method = my_init,
                optim.method = my_opt
  )
  
  assign(paste("fit", num_comp, num_seed, sep="_"), fit)
  
  results <- data.frame(
      mae = fit$gof$mae,
      mse = fit$gof$mse,
      rsq = calc_rsq(fit$pred$yhat, fit$pred$y),
      ll = -fit$gof$ll,
      aic = fit$gof$aic,
      bic = fit$gof$bic
  )
  
  # save the metrics
  write.csv(results, paste(RESULTS_PATH, "metrics", paste("aldvmm_", num_comp, "_", num_seed, "_metrics", ".csv", sep=""), sep="/"), row.names=FALSE, na='')
  
  # dump the coef
  write.csv(data.frame(fit$coef), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_coef", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
  
  # Covariance matrix
  write.csv(data.frame(fit$cov), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_cov", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
  # dump the coef SEs
  write.csv(data.frame(fit$se), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_coef_se", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
  }
}

  # fitting ALDVMMM with 2 & 4 components
  NUM_COMP <- c(2, 4) 
  num_seed <- 1 
  
  for (num_comp in NUM_COMP){
    set.seed(num_seed)
    res_file_name = paste("aldvmm_", num_comp, "_", my_opt, "_", my_init, ".csv", sep="")
    
    fit <- aldvmm(formula = model_formula_aldvmm,
                  data = data_factor,
                  psi = c(-0.594, 1),
                  ncmp = num_comp,
                  init.method = my_init,
                  optim.method = my_opt
    )
    
    assign(paste("fit", num_comp, num_seed, sep="_"), fit)
    
    results <- data.frame(
      mae = fit$gof$mae,
      mse = fit$gof$mse,
      rsq = calc_rsq(fit$pred$yhat, fit$pred$y),
      ll = -fit$gof$ll,
      aic = fit$gof$aic,
      bic = fit$gof$bic
    )
    
    # save the metrics
    write.csv(results, paste(RESULTS_PATH, "metrics", paste("aldvmm_", num_comp, "_", num_seed, "_metrics", ".csv", sep=""), sep="/"), row.names=FALSE, na='')
    
    # dump the coef
    write.csv(data.frame(fit$coef), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_coef", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
  
    # Covariance matrix
    write.csv(data.frame(fit$cov), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_cov", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
    # dump the coef SEs
    write.csv(data.frame(fit$se), paste(RESULTS_PATH, "coef", paste("aldvmm_", num_comp, "_", num_seed, "_coef_se", ".csv", sep=""), sep="/"), row.names=TRUE, na='')
}

#########################################
# try different optimisation for ALDVMM 5 

num_comp <- 5

fit <- aldvmm(formula = model_formula_aldvmm,
              data = data_factor,
              psi = c(-0.594, 1),
              ncmp = num_comp,
              init.method = my_init,
              optim.method = my_opt
)

assign(paste("fit", 5, sep=""), fit)

results <- data.frame(
  mae = calc_mae(fit$pred$yhat, fit$pred$y),
  mse = calc_mse(fit$pred$yhat, fit$pred$y),
  rsq = calc_rsq(fit$pred$yhat, fit$pred$y),
  ll = -fit$gof$ll,
  aic = fit$gof$aic,
  bic = fit$gof$bic
)


#####################
# Predictions
#####################

all_predictions <- data.frame(
  Disutility = glm$y,
  #Disutility_scaled = data_factor$Disutility_scaled,
  GLM_pred = glm$fitted.values,
  OLS_pred = ols$fitted.values,
  #TwoPart_GLM_pred = predict(twopm_glm, data_factor),
  #TwoPart_OLS_pred = predict(twopm_ols, data_factor),
  CLAD_pred = predict(clad, data_factor),
  ALDVMM_2_pred = predict(fit_2_1, data_factor),
  ALDVMM_3_pred = predict(fit_3_1, data_factor),
  ALDVMM_4_pred = predict(fit_4_1, data_factor)
)

write.csv(all_predictions, paste(RESULTS_PATH, "predictions",  "all_predictions.csv", sep="/"), row.names=FALSE, na='')

# Fix all covariates to their baselines values
ISI_Score_mean = 12.12
ISI_Score_SD = 5.31
mean_untreated_dataset = data.frame(
  EQ5DINDEX = rep(100, 29),
  ISI_Score = (c(0:28) - ISI_Score_mean) / ISI_Score_SD,
  deage = rep(mean(data_factor$deage), 29),
  female = rep(mean(data$female), 29),
  married = rep(mean(data$married), 29),
  degree_4years = rep(mean(data$degree_4years), 29),
  employed = rep(mean(data$employed), 29),
  retired = rep(mean(data$retired), 29),
  current_smoker  = rep(mean(data$current_smoker), 29),
  former_smoker = rep(mean(data$former_smoker), 29),
  heavy_drinker = rep(mean(data$heavy_drinker), 29),
  low_moderate_drinker = rep(mean(data$low_moderate_drinker), 29),
  BMI_R = rep(mean(data_factor$BMI_R), 29),
  CCI = rep(mean(data_factor$CCI), 29),
  Combined_DP_PTS_AX = rep(mean(data$Combined_DP_PTS_AX), 29),
  pain = rep(mean(data$pain), 29),
  treated = 0,
  diagnosed = rep(mean(data$diagnosed), 29),
  treated_ISI_Score = 0,
  Not_Country_UK = 0
)


y_pred2_mean <- predict(fit2, newdata=mean_untreated_dataset)
y_pred2_mean <- as.double(y_pred2_mean$yhat)
y_pred3_mean <- predict(fit3, newdata=mean_untreated_dataset)
y_pred3_mean <- as.double(y_pred3_mean$yhat)

untreated_mapping = data.frame(
  ISI_score = c(0:28),
  EQ_5D_3L_3 = y_pred3_mean,
  EQ_5D_3L_2 = y_pred2_mean
)

write.csv(untreated_mapping, paste(RESULTS_PATH, "predictions", "ALDVMM_untreated_23.csv", sep="/"), row.names=FALSE, na='')

# Fix all covariates to their baselines values
ISI_Score_mean = 12.12
ISI_Score_SD = 5.31
mean_dataset = data.frame(
  EQ5DINDEX = rep(100, 29),
  ISI_Score = (c(0:28) - ISI_Score_mean) / ISI_Score_SD,
  deage = rep(mean(data$deage), 29),
  female = rep(mean(data$female), 29),
  married = rep(mean(data$married), 29),
  degree_4years = rep(mean(data$degree_4years), 29),
  employed = rep(mean(data$employed), 29),
  retired = rep(mean(data$retired), 29),
  current_smoker  = rep(mean(data$current_smoker), 29),
  former_smoker = rep(mean(data$former_smoker), 29),
  heavy_drinker = rep(mean(data$heavy_drinker), 29),
  low_moderate_drinker = rep(mean(data$low_moderate_drinker), 29),
  BMI_R = rep(mean(data$BMI_R), 29),
  CCI = rep(mean(data$CCI), 29),
  Combined_DP_PTS_AX = rep(mean(data$Combined_DP_PTS_AX), 29),
  pain = rep(mean(data$pain), 29),
  treated = rep(mean(data$treated), 29),
  diagnosed = rep(mean(data$diagnosed), 29),
  treated_ISI_Score = rep(mean(data$treated_ISI_Score), 29),
  Not_Country_UK = rep(mean(data$Not_Country_UK), 29)
)

y_pred3_mean <- predict(fit3, newdata=mean_dataset)
y_pred3_mean <- as.double(y_pred3_mean$yhat)

glm_mean <- predict(glm, newdata=mean_dataset, type="response")

all_mapping = data.frame(
  ISI_score = c(0:28),
  EQ5D3L_ALDVMM_pred = y_pred3_mean,
  Disutility_GLM_pred = glm_mean
)
all_mapping$EQ5D3L_GLM_pred <- 1- all_mapping$Disutility_GLM_pred
write.csv(all_mapping, paste(RESULTS_PATH, "predictions", "mean_of_covariates.csv", sep="/"), row.names=FALSE, na='')