# takes a dataframe and returns the same dataframe with two extra columns for colours and position
library('viridis')

generate_distance_colour_map = function(plot_df,
                                        xvar_colname = "position",
                                        lig_dist_colname = "ligand_distance",
                                        debug = TRUE
)
{
  if (debug) {
    cat("\nAnnotating x-axis ~", lig_dist_colname, "requested...")
  }
  
  plot_df['lig_distR'] = round(plot_df[[lig_dist_colname]], digits = 0)
  
  lig_min  = min(round(plot_df[[lig_dist_colname]]), na.rm = T); lig_min
  lig_max  = max(round(plot_df[[lig_dist_colname]]), na.rm = T); lig_max 
  lig_mean = round(mean(round(plot_df[[lig_dist_colname]]), na.rm = T)); lig_mean
  n_colours  = length(sort(unique(round(plot_df[[lig_dist_colname]], digits = 0))))
  lig_cols = magma(n_colours, direction=-1)
  ligD_valsR = sort(unique(round(plot_df[[lig_dist_colname]], digits = 0)))
  
  if (debug) {
    length(ligD_valsR)
    if (n_colours == length(ligD_valsR)) {
      cat("\nStarting: mapping b/w"
          , lig_dist_colname
          , "and colours")
    }else{
      cat("\nCannot start mapping b/w", lig_dist_colname, "and colours..."
          , "\nLength mismatch:"
          , "No. of colours: ", n_colours
          , "\nValues to map:", length(ligD_valsR))
    }
  }
  
  ligDcolKey <- data.frame(ligD_colours = lig_cols
                           , lig_distR = ligD_valsR); ligDcolKey
  if (debug) {
    names(ligDcolKey)
    cat("\nSuccessful: Mapping b/w", lig_dist_colname, "and colours")
  }
  # merge colour key with plot_df
  plot_df = merge(plot_df, ligDcolKey, by = 'lig_distR')
  
  return(plot_df)
}

generate_distance_legend = function(plot_df,
                                    xvar_colname = 'position',
                                    lig_dist_colname = "ligand_distance",
                                    legend_title = "Ligand\nDistance"
)
{
  # build legend for ligand distance "heat bar"
  lig_min  = min(round(plot_df[[lig_dist_colname]]), na.rm = T); lig_min
  lig_max  = max(round(plot_df[[lig_dist_colname]]), na.rm = T); lig_max 
  lig_mean = round(mean(round(plot_df[[lig_dist_colname]]), na.rm = T)); lig_mean
  
  labels  = seq(lig_min, lig_max, len = 5); labels
  labelsD = round(labels, digits = 0); labelsD
  
  get_legend(
    ggplot(plot_df, aes_string(x = sprintf("factor(%s)", xvar_colname), y=0)) +
      
      geom_tile(aes(fill = .data[[lig_dist_colname]])
                , colour = "white") +
      scale_fill_gradient2(midpoint = lig_mean
                           , low  = magma(3, direction=-1)[1]
                           , mid  = magma(3, direction=-1)[2]
                           , high = magma(3, direction=-1)[3]
                           , breaks = labels
                           , limits = c(lig_min, lig_max)
                           , labels = labelsD
                           , name   = legend_title)
  )
}