#########################################################
# TASK: Adding colours to positions labels according to 
# active site residues. This is so these can be seen promptly
# when visualising the barplot.
#########################################################
#=======================================================================
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting")
getwd()

source("plotting_data.R")
   # should return the following dfs and directories
   # my_df
   # my_df_u
   # my_df_u_lig
   # dup_muts

#=======================================================================
###########################
# Read file: struct params
###########################
cat("Reading struct params including mcsm:", in_filename_params)

my_df = read.csv(infile_params
                 #, stringsAsFactors = F
                 , header = T)

cat("Input dimensions:", dim(my_df)) 

# clear variables
rm(in_filename_params, infile_params)

# quick checks
colnames(my_df)
str(my_df)

# check for duplicate mutations
if ( length(unique(my_df$mutationinformation)) != length(my_df$mutationinformation)){
  cat(paste0("CAUTION:", " Duplicate mutations identified"
             , "\nExtracting these..."))
  dup_muts = my_df[duplicated(my_df$mutationinformation),]
  dup_muts_nu = length(unique(dup_muts$mutationinformation))
  cat(paste0("\nDim of duplicate mutation df:", nrow(dup_muts)
             , "\nNo. of unique duplicate mutations:", dup_muts_nu
             , "\n\nExtracting df with unique mutations only"))
  my_df_u = my_df[!duplicated(my_df$mutationinformation),]
}else{
  cat(paste0("No duplicate mutations detected"))
  my_df_u = my_df
}

upos = unique(my_df_u$position)
cat("Dim of clean df:"); cat(dim(my_df_u))
cat("\nNo. of unique mutational positions:"); cat(length(upos))
#=======================================================================
# create a new df with unique position numbers and cols
position = unique(my_df$position) #130
position_cols = as.data.frame(position)

head(position_cols) ; tail(position_cols)

# specify active site residues and bg colour
position = c(49, 51, 57, 71
             , 8, 96, 138
             , 13, 68
             , 103, 137
             , 133, 134) #13

lab_bg = rep(c("purple" 
               , "yellow"
               , "cornflowerblue"
               , "blue"
               , "green"), times = c(4, 3, 2, 2, 2)
)

# second bg colour for active site residues
#lab_bg2 = rep(c("white" 
#                , "green" , "white", "green"
#                , "white"
#                , "white"
#                , "white"), times = c(4
#                                      , 1, 1, 1
#                                      , 2
#                                      , 2
#                                      , 2)
#)

#%%%%%%%%%
# revised: leave the second box coloured as the first one incase there is no second colour
#%%%%%%%%%
lab_bg2 = rep(c("purple" 
                , "green", "yellow", "green"
                , "cornflowerblue"
                , "blue"
                , "green"), times = c(4
                                      , 1, 1, 1
                                      , 2
                                      , 2
                                      , 2))

# fg colour for labels for active site residues
lab_fg = rep(c("white" 
               , "black"
               , "black"
               , "white"
               , "black"), times = c(4, 3, 2, 2, 2))

#%%%%%%%%%
# revised: make the purple ones black
# fg colour for labels for active site residues
#%%%%%%%%%
#lab_fg = rep(c("black" 
#               , "black"
#               , "black"
#               , "white"
#               , "black"), times = c(4, 3, 2, 2, 2))               
               
# combined df with active sites, bg and fg colours
aa_cols_ref = data.frame(position
                         , lab_bg
                         , lab_bg2
                         , lab_fg
                         , stringsAsFactors = F) #13, 4

str(position_cols); class(position_cols)
str(aa_cols_ref); class(aa_cols_ref)

# since position is int and numeric in the two dfs resp, 
# converting numeric to int for consistency
aa_cols_ref$position = as.integer(aa_cols_ref$position)
class(aa_cols_ref$position)

#===========
# Merge 1: merging positions df (position_cols) and
# active site cols (aa_cols_ref)
# linking column: "position"
# This is so you can have colours defined for all positions
#===========
head(position_cols$position); head(aa_cols_ref$position)

mut_pos_cols = merge(position_cols, aa_cols_ref
            , by = "position"
            , all.x = TRUE) 

head(mut_pos_cols)
# replace NA"s 
# :column "lab_bg" with "white"
# : column "lab_fg" with "black"
mut_pos_cols$lab_bg[is.na(mut_pos_cols$lab_bg)] <- "white"
mut_pos_cols$lab_bg2[is.na(mut_pos_cols$lab_bg2)] <- "white"
mut_pos_cols$lab_fg[is.na(mut_pos_cols$lab_fg)] <- "black"
head(mut_pos_cols)

#===========
# Merge 2: Merge mut_pos_cols with mcsm df
# Now combined the positions with aa colours with 
# the mcsm_data
#===========
# dfs to merge
df0 = my_df # my_df_o
df1 = mut_pos_cols

# check the column on which merge will be performed
head(df0$position); tail(df0$position)
head(df1$position); tail(df1$position)

# should now have 3 extra columns
my_df = merge(df0, df1
              , by = "position"
              , all.x = TRUE)

# sanity check
my_df[my_df$position == "49",]
my_df[my_df$position == "13",]

rm(df0, df1)
#===========
# Merge 3: Merge mut_pos_cols with mcsm df_u
# Now combined the positions with aa colours with 
# the mcsm_data
#===========
# dfs to merge
df0 = my_df_u # my_df_u
df1 = mut_pos_cols

# check the column on which merge will be performed
head(df0$position); tail(df0$position)
head(df1$position); tail(df1$position)

# should now have 3 extra columns
my_df_u = merge(df0, df1
              , by = "position"
              , all.x = TRUE)

# sanity check
my_df[my_df$position == "49",]
my_df[my_df$position == "13",]

# clear variables
rm(aa_cols_ref
   , df0
   , df1
   , position_cols
   , lab_bg
   , lab_bg2
   , lab_fg
   , position)