#########################################################
# TASK: Adding colours to dfs so they can be used for plotting
# add  cols to each of the my_df* dfs
#########################################################
#=======================================================================
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

cat(paste0("Directories imported:"
           , "\ndatadir:", datadir
           , "\nindir:", indir
           , "\noutdir:", outdir
           , "\nplotdir:", plotdir))

cat(paste0("Variables imported:"
           , "\ndrug:", drug
           , "\ngene:", gene
           , "\ngene_match:", gene_match
           , "\nLength of upos:", length(upos)
    , "\nAngstrom symbol:", angstroms_symbol))       

# clear excess variable
rm(upos, dup_muts, my_df_u, my_df_u_lig)

# This is because we want to assign the colours to my_df
# and then resubset accordingly for our plots to avoid multiple merges

#=======================================================================
# df to use: my_df
# NOTE: my_df contains duplicate muts but its ok as you are only adding 
# colours to positions

# sanity checks: ensure my_df is ordered by position: it should be 
my_df$position; my_df$mutationinformation

my_df_o = my_df[order(my_df$position),]
my_df_o$position; my_df_o$mutationinformation

head(my_df_o$position) == head(my_df$position)
head(my_df_o$mutationinformation) == head(my_df$mutationinformation)
tail(my_df_o$position) == tail(my_df$position)
tail(my_df_o$mutationinformation) == tail(my_df$mutationinformation)

my_df = my_df_o

# create a new df with unique position numbers and cols
position = unique(my_df$position) 
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_cols = merge(df0, df1
              , by = "position"
              , all.x = TRUE)

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

###########################
# extract unique mutation entries
###########################

# check for duplicate mutations
if ( length(unique(my_df_cols$mutationinformation)) != length(my_df_cols$mutationinformation)){
   cat(paste0("\nCAUTION:", " Duplicate mutations identified"
              , "\nExtracting these..."))
   dup_muts_cols = my_df_cols[duplicated(my_df_cols$mutationinformation),]
   dup_muts_cols_nu = length(unique(dup_muts_cols$mutationinformation))
   cat(paste0("\nDim of duplicate mutation df:", nrow(dup_muts_cols)
              , "\nNo. of unique duplicate mutations:", dup_muts_cols_nu
              , "\n\nExtracting df with unique mutations only"))
   my_df_u_cols = my_df_cols[!duplicated(my_df_cols$mutationinformation),]
}else{
   cat(paste0("\nNo duplicate mutations detected"))
   my_df_u_cols = my_df_cols
}

upos = unique(my_df_u_cols$position)
cat("\nDim of clean df:"); cat(dim(my_df_u_cols))
cat("\nNo. of unique mutational positions:"); cat(length(upos), "\n")


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

###########################
# extract mutations <10Angstroms
###########################
table(my_df_u_cols$ligand_distance<10)

my_df_u_cols_lig = my_df_u_cols[my_df_u_cols$ligand_distance <10,]
angstroms_symbol = "\u212b"
cat(paste0("There are ", nrow(my_df_u_cols_lig), " sites lying within 10", angstroms_symbol, " of the ligand"))

#=================
# very important!
#=================
#my_axis_colours = mut_pos_cols$lab_fg # doesn't work if positions numbers are subsetted as in ligand
# need the equivalent of the mut_pos_cols for ligand

# get position numbers for ligand
lig_pos = my_df_u_cols_lig$position

# subset mut_pos_cols for ligand positions
mut_pos_cols_lig = mut_pos_cols[mut_pos_cols$position %in% lig_pos,]
#my_axis_colours = mut_pos_cols_lig$lab_fg

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

#######################################################################
# end of script
#######################################################################