cox.disc.filename <- '../../output/all-cv-survreg-boot-try5-surv-model.rds'
caliber.missing.coefficients.filename <-
  '../../output/caliber-replicate-with-missing-survreg-6-linear-age-coeffs-3.csv'
rf.filename <- '../../output/rfsrc-cv-nsplit-try3-var-effects.csv'

source('../lib/shared.R')
requirePlus('cowplot')

# Amount of padding at the right-hand side to make space for missing values
missing.padding <- 0.05

continuous.vars <-
  c(
    'age', 'total_chol_6mo', 'hdl_6mo', 'pulse_6mo', 'crea_6mo',
    'total_wbc_6mo', 'haemoglobin_6mo'
  )

# Load in the discretised Cox model for plotting
surv.model.fit.boot <- readRDS(cox.disc.filename)

# Pull coefficients from model with missing data
caliber.missing.coeffs <- read.csv(caliber.missing.coefficients.filename)
# Log them to get them on the same scale as discrete model
caliber.missing.coeffs$our_value <- -log(caliber.missing.coeffs$our_value)
caliber.missing.coeffs$our_lower <- -log(caliber.missing.coeffs$our_lower)
caliber.missing.coeffs$our_upper <- -log(caliber.missing.coeffs$our_upper)

# Load the data
COHORT.use <- data.frame(fread(data.filename))

# Open the calibration to find the best binning scheme
calibration.filename <- '../../output/survreg-crossvalidation-try5.csv'
cv.performance <- read.csv(calibration.filename)

# Find the best calibration...
# First, average performance across cross-validation folds
cv.performance.average <-
  aggregate(
    c.index.val ~ calibration,
    data = cv.performance,
    mean
  )
# Find the highest value
best.calibration <-
  cv.performance.average$calibration[
    which.max(cv.performance.average$c.index.val)
    ]
# And finally, find the first row of that calibration to get the n.bins values
best.calibration.row1 <-
  min(which(cv.performance$calibration == best.calibration))

# Get its parameters
n.bins <-
  t(
    cv.performance[best.calibration.row1, continuous.vars]
  )

# Prepare the data with those settings...

# Reset process settings with the base setings
process.settings <-
  list(
    var        = c('anonpatid', 'time_death', 'imd_score', 'exclude'),
    method     = c(NA, NA, NA, NA),
    settings   = list(NA, NA, NA, NA)
  )
for(j in 1:length(continuous.vars)) {
  process.settings$var <- c(process.settings$var, continuous.vars[j])
  process.settings$method <- c(process.settings$method, 'binByQuantile')
  process.settings$settings <-
    c(
      process.settings$settings,
      list(
        seq(
          # Quantiles are obviously between 0 and 1
          0, 1,
          # Choose a random number of bins (and for n bins, you need n + 1 breaks)
          length.out = n.bins[j]
        )
      )
    )
}

# prep the data given the variables provided
COHORT.optimised <-
  prepData(
    # Data for cross-validation excludes test set
    COHORT.use,
    cols.keep,
    process.settings,
    surv.time, surv.event,
    surv.event.yes,
    extra.fun = caliberExtraPrep
  )

# Unpack variable and level names
cph.coeffs <- cphCoeffs(
  bootStats(surv.model.fit.boot, uncertainty = '95ci', transform = `-`),
  COHORT.optimised, surv.predict, model.type = 'boot.survreg'
)

# We'll need the CALIBER scaling functions for plotting
source('../cox-ph/caliber-scale.R')

# set up list to store the plots
cox.discrete.plots <- list()
# Add dummy columns for x-position of missing values
caliber.missing.coeffs$missing.x.pos.cont <- NA
cph.coeffs$missing.x.pos.disc <- NA

for(variable in unique(cph.coeffs$var)) {
  # If it's a continuous variable, get the real centres of the bins
  if(variable %in% process.settings$var) {
    process.i <- which(variable == process.settings$var)
    
    if(process.settings$method[[process.i]] == 'binByQuantile') {
      
      variable.quantiles <-
        getQuantiles(
          COHORT.use[, variable],
          process.settings$settings[[process.i]]
        )
      # For those rows which relate to this variable, and whose level isn't
      # missing, put in the appropriate quantile boundaries for plotting
      cph.coeffs$bin.min[cph.coeffs$var == variable & 
                           cph.coeffs$level != 'missing'] <-
        variable.quantiles[1:(length(variable.quantiles) - 1)]
      cph.coeffs$bin.max[cph.coeffs$var == variable & 
                           cph.coeffs$level != 'missing'] <-
        variable.quantiles[2:length(variable.quantiles)]
      # Make the final bin the 99th percentile
      cph.coeffs$bin.max[cph.coeffs$var == variable & 
                           cph.coeffs$level != 'missing'][
                             length(variable.quantiles) - 1] <-
        quantile(COHORT.use[, variable], 0.99, na.rm = TRUE)
      
      # Add a fake data point at the highest value to finish the graph
      cph.coeffs <-
        rbind(
          cph.coeffs,
          cph.coeffs[cph.coeffs$var == variable & 
                       cph.coeffs$level != 'missing', ][
                         length(variable.quantiles) - 1, ]
        )
      # Change it so that bin.min is bin.max from the old one
      cph.coeffs$bin.min[nrow(cph.coeffs)] <-
        cph.coeffs$bin.max[cph.coeffs$var == variable & 
                             cph.coeffs$level != 'missing'][
                               length(variable.quantiles) - 1]
      
      # Work out data range by taking the 1st and 99th percentiles
      # Use the max to provide a max value for the final bin
      # Also use for x-axis limits, unless there are missing values to
      # accommodate on the right-hand edge.
      x.data.range <-
        quantile(COHORT.use[, variable], c(0.01, 0.99), na.rm = TRUE)
      x.axis.limits <- x.data.range
      
      
      # Finally, we need to scale this such that the baseline value is equal
      # to the value for the equivalent place in the Cox model, to make the
      # risks comparable...
      
      # First, we need to find the average value of this variable in the lowest
      # bin (which is always the baseline here)
      baseline.bin <- variable.quantiles[1:2]
      baseline.bin.avg <- 
        mean(
          # Take only those values of the variable which are in the range
          COHORT.use[
            inRange(COHORT.use[, variable], baseline.bin, na.false = TRUE),
            variable
            ]
        )
      # Then, scale it with the caliber scaling
      baseline.bin.val <-
        caliberScaleUnits(baseline.bin.avg, variable) * 
        caliber.missing.coeffs$our_value[
          caliber.missing.coeffs$quantity == variable
          ]
      
      # And now, add all the discretised values to that value to make them
      # comparable...
      cph.coeffs[cph.coeffs$var == variable, c('val', 'lower', 'upper')] <-
        cph.coeffs[cph.coeffs$var == variable, c('val', 'lower', 'upper')] -
        baseline.bin.val
      
      # Now, plot this variable as a stepped line plot using those quantile
      # boundaries
      cox.discrete.plot <-
        ggplot(
          subset(cph.coeffs, var == variable),
          aes(x = bin.min, y = val)
        ) +   
        geom_step() +
        geom_step(aes(y = lower), colour = 'grey') +
        geom_step(aes(y = upper), colour = 'grey') +
        labs(x = variable, y = 'Bx')
      
      # If there's a missing value risk, add it
      if(any(cph.coeffs$var == variable & cph.coeffs$level == 'missing')) {
        # Expand the x-axis to squeeze the missing values in
        x.axis.limits[2] <- 
          x.axis.limits[2] + diff(x.data.range) * missing.padding
        # Put this missing value a third of the way into the missing area
        cph.coeffs$missing.x.pos.disc[
          cph.coeffs$var == variable &
          cph.coeffs$level == 'missing'] <-
          x.axis.limits[2] + diff(x.data.range) * missing.padding / 3

        # Add the point to the graph (we'll set axis limits later)
        cox.discrete.plot <-
          cox.discrete.plot +
          geom_pointrange(
            data = cph.coeffs[cph.coeffs$var == variable & 
                                cph.coeffs$level == 'missing', ],
            aes(
              x = missing.x.pos.disc,
              y = val, ymin = lower,
              ymax = upper
            ),
            colour = 'red'
          )
      }
      
      # Now, let's add the line from the continuous Cox model. We only need two
      # points because the lines are straight!
      continuous.cox <-
        data.frame(
          var.x.values = x.data.range
        )
      # Scale the x-values
      continuous.cox$var.x.scaled <-
        caliberScaleUnits(continuous.cox$var.x.values, variable)
      # Use the risks to calculate risk per x for central estimate and errors
      continuous.cox$y <-
        -caliber.missing.coeffs$our_value[
          caliber.missing.coeffs$quantity == variable
          ] * continuous.cox$var.x.scaled
      continuous.cox$upper <-
        -caliber.missing.coeffs$our_upper[
          caliber.missing.coeffs$quantity == variable
          ] * continuous.cox$var.x.scaled
      continuous.cox$lower <-
        -caliber.missing.coeffs$our_lower[
          caliber.missing.coeffs$quantity == variable
          ] * continuous.cox$var.x.scaled
      
      cox.discrete.plot <-
        cox.discrete.plot +
        geom_line(
          data = continuous.cox,
          aes(x = var.x.values, y = y),
          colour = 'blue'
        ) +
        geom_line(
          data = continuous.cox,
          aes(x = var.x.values, y = upper),
          colour = 'lightblue'
        ) +
        geom_line(
          data = continuous.cox,
          aes(x = var.x.values, y = lower),
          colour = 'lightblue'
        )
      
      # If there is one, add missing value risk from the continuous model
      if(any(caliber.missing.coeffs$quantity == paste0(variable, '_missing') &
             caliber.missing.coeffs$unit == 'missing')) {
        
        # Put this missing value 2/3rds of the way into the missing area
        caliber.missing.coeffs$missing.x.pos.cont[
          caliber.missing.coeffs$quantity == paste0(variable, '_missing') &
            caliber.missing.coeffs$unit == 'missing'] <-
                x.axis.limits[2] + 2 * diff(x.data.range) * missing.padding / 3
        
        cox.discrete.plot <-
          cox.discrete.plot +
          geom_pointrange(
            data = caliber.missing.coeffs[
              caliber.missing.coeffs$quantity == paste0(variable, '_missing') &
                caliber.missing.coeffs$unit == 'missing',
              ],
            aes(
              x = missing.x.pos.cont,
              y = our_value, ymin = our_lower, ymax = our_upper
            ),
            colour = 'blue'
          )
      }
      
      # Finally, set the x-axis limits; will just be the data range, or data
      # range plus a bit if there are missing values to squeeze in
      cox.discrete.plot <-
        cox.discrete.plot +
        coord_cartesian(xlim = x.axis.limits) +
        theme(axis.title.y = element_blank()) +
        theme(plot.margin = unit(c(0.2, 0.1, 0.2, 0.1), "cm"))
      
      cox.discrete.plots[[variable]] <- cox.discrete.plot
    }
  }
}

# Load the random forest variable effects file
risk.by.variables <- read.csv(rf.filename)
rf.vareff.plots <- list()

for(variable in unique(risk.by.variables$var)) {
  # Get the mean of the normalised risk for every value of the variable
  risk.aggregated <-
    aggregate(
      as.formula(paste0('risk.normalised ~ val')),
      subset(risk.by.variables, var == variable), median
    )
  
  # work out the limits on the axes by taking the 1st and 99th percentiles
  x.axis.limits <-
    quantile(COHORT.use[, variable], c(0.01, 0.99), na.rm = TRUE)
  y.axis.limits <-
    quantile(subset(risk.by.variables, var == variable)$risk.normalised, c(0.05, 0.95), na.rm = TRUE)
  
  # If there's a missing value risk in the graph above, expand the axes so they
  # match
  if(any(cph.coeffs$var == variable & cph.coeffs$level == 'missing')) {
    x.axis.limits[2] <- 
      x.axis.limits[2] + diff(x.data.range) * missing.padding
  }
  
  rf.vareff.plots[[variable]] <-
    ggplot(
      subset(risk.by.variables, var == variable), 
      aes(x = val, y = log(risk.normalised))
    ) +
    geom_line(alpha=0.003, aes(group = id)) +
    geom_line(data = risk.aggregated, colour = 'blue') +
    coord_cartesian(xlim = x.axis.limits, ylim = log(y.axis.limits)) +
    labs(x = variable) +
    theme(
      plot.margin = unit(c(0.2, 0.1, 0.2, 0.1), "cm"),
      axis.title.y = element_blank()
    )
}


plot_grid(
  cox.discrete.plots[['age']],
  cox.discrete.plots[['haemoglobin_6mo']],
  cox.discrete.plots[['total_wbc_6mo']],
  cox.discrete.plots[['crea_6mo']],
  rf.vareff.plots[['age']],
  rf.vareff.plots[['haemoglobin_6mo']],
  rf.vareff.plots[['total_wbc_6mo']],
  rf.vareff.plots[['crea_6mo']],
  labels = c('A', rep('', 3), 'B', rep('', 3)),
  align = "v", ncol = 4
)

ggsave(
  '../../output/variable-effects.pdf',
  width = 16,
  height = 10,
  units = 'cm',
  useDingbats = FALSE
)