#########################################################
# 1b: Define function: coloured barplot by subgroup
# LINK: https://stackoverflow.com/questions/49818271/stacked-barplot-with-colour-gradients-for-each-bar
#########################################################
#source("~/git/LSHTM_analysis/scripts/functions/generate_distance_colour_map.R")

ColourPalleteMulti = function(df, group, subgroup){
  
  # Find how many colour categories to create and the number of colours in each
  categories <- aggregate(as.formula(paste(subgroup, group, sep="~" ))
                          , df
                          , function(x) length(unique(x)))
  #  return(categories) }
  
  category.start <- (scales::hue_pal(l = 100)(nrow(categories))) # Set the top of the colour pallete
  
  category.end  <- (scales::hue_pal(l = 40)(nrow(categories))) # set the bottom
  
  #return(category.start); return(category.end)}
  
  # Build Colour pallette
  colours <- unlist(lapply(1:nrow(categories),
                           function(i){
                             colorRampPalette(colors = c(category.start[i]
                                                         , category.end[i]))(categories[i,2])}))
  return(colours)
}
#########################################################################

########################
# Generate bp with
# colour palette derived
# from the data using 
# above function
#########################

bp_stability_hmap <- function(plot_df = merged_df3
                              , xvar_colname = "position"
                              , yvar_colname = 'avg_stability_scaled' # Only here so that you can do function(df)
                              #, bar_col_colname = "group"
                              , stability_colname = "avg_stability_scaled"  # Only here so that you can do function(df)
                              , stability_outcome_colname = "avg_stability_outcome"  # Only here so that you can do function(df)
                              , p_title = "DUMMY TITLE",  # Only here so that you can do function(df)
                              my_xaxls = 6, # x-axis label size
                              my_yaxls = 6, # y-axis label size
                              my_xaxts = 9, # x-axis text size
                              my_yaxts = 10, # y-axis text size
                              my_pts  = 10  # plot-title size
                              , my_xlab = "Position"
                              , my_ylab = "No. of nsSNPs"
                              
                              # Custom 2: x-axis: geom tiles ~ lig distance
                              #, A_xvar_lig = T
                              , lig_dist_colname = LigDist_colname # from globals
                              , tpos0 = 0 # 0 is a magic number that does my sensible default
                              , tW0 = 1
                              , tH0 = 0.2,
                              y_max_override = 1, # an override for tidily plotting multiple different-ranged plots together
                              reorder_position = FALSE, # enable to reorder according to plot_df$pos_count
                              ...
                              
                              
                              
                              
)
{
  # Custom 2: x-axis geom tiles ~ lig distance
  
  # order the df by position and ensure it is a factor
  plot_df = plot_df[order(plot_df[[xvar_colname]]), ]
  plot_df[[xvar_colname]] = factor(plot_df[[xvar_colname]])
  
  #cat("\nSneak peak:\n")
  head(data.frame( plot_df[[xvar_colname]], plot_df[[stability_colname]] ) )
  
  # stability values isolated to help with generating column called: 'group'
  my_grp = plot_df[[stability_colname]]
  # cat( "\nLength of nsSNPs:", length(my_grp)
  #      , "\nLength of unique values for nsSNPs:", length(unique(my_grp)) )
  # 
  # Add col: 'group'
  plot_df$group = paste0(plot_df[[stability_outcome_colname]], "_", my_grp, sep = "")
  plot_df=plot_df %>% dplyr::add_count(position)
  plot_df$pos_count=plot_df$n
  plot_df$n=NULL
  
  # define a "max Y" in case the user didn't supply one
  if(reorder_position) {
    y_max = max(plot_df$pos_count)
  } 
  else{
    y_max = 1 # boring default
  }
  y_axis_limit = round_any(y_max, y_max_override, ceiling)
  
  # Call the function to create the palette based on the group defined above
  #subcols_ps
  subcols_bp_hmap = ColourPalleteMulti(plot_df, stability_outcome_colname, stability_colname)
  
  cat("\nNo. of sub colours generated:", length(subcols_bp_hmap))
  anno_bar=position_annotation(plot_df,
                               reorder_position=reorder_position,
                               aa_pos_drug=aa_pos_drug,
                               active_aa_pos=active_aa_pos,
                               aa_pos_lig1=aa_pos_lig1,
                               aa_pos_lig2=aa_pos_lig2,
                               aa_pos_lig3=aa_pos_lig3,
                               ...
  )
  
  subcols_plot = ggplot(plot_df) +
    scale_fill_manual( values = subcols_bp_hmap
                       , guide = "none") +
    # scale_x_discrete("Position", labels=factor(plot_df$position)) +
    scale_y_continuous(limits=c(0,y_axis_limit)) +
    theme(
      panel.grid = element_line(color="lightgrey", size=0.125)
      , axis.text.x = element_text(size = my_xaxls
                                   , angle = 90
                                   , hjust = 1
                                   , vjust = 0.4)
      , axis.text.y = element_text(size = my_yaxls
                                   , angle = 0
                                   , hjust = 1
                                   , vjust = 0)
      , axis.title.x = element_blank()
      , axis.ticks = element_blank()
      #, axis.title.x = element_text(size = my_xaxts)
      , axis.title.y = element_text(size = my_yaxts )
      , plot.title = element_text(size = my_pts
                                  , hjust = 0.5)
      # , panel.grid = element_blank()
      , panel.background = element_rect(fill = "transparent", colour=NA)
    ) +
    labs(title = p_title
         , x = my_xlab
         , y = my_ylab) +
    if(reorder_position) {
      geom_bar(aes(x=reorder(position,-pos_count), fill = group),
               colour = "grey",
               size=0.125
      )
      
    }else{
      geom_bar(aes(x=position, fill = group),
               colour = "grey",
               size=0.125
      )
    }
  
  
  # Generate the subcols barplot
  cowplot::plot_grid(
    subcols_plot,
    NULL,
    anno_bar,
    ncol = 1,
    align = "v",
    rel_heights = c(6,-0.1,1)
  )
  
}
# bp_stability_hmap(small_df3)