getwd()
setwd('~/git/LSHTM_analysis/scripts/plotting')
getwd()

#########################################################
# TASK:

#########################################################

########################################################################
# 		Installing and loading required packages and functions		   #	
########################################################################

#source('Header_TT.R')
#source('barplot_colour_function.R')

########################################################################
#		 Read file: call script for combining df for PS			   	   #
########################################################################
#?????????????
#
########################################################
#%% variable assignment: input and output paths & filenames
drug = 'pyrazinamide'
gene = 'pncA'
gene_match = paste0(gene,'_p.')
cat(gene_match)

#=============
# directories
#=============
datadir = paste0('~/git/Data')
indir = paste0(datadir, '/', drug, '/input')
outdir = paste0('~/git/Data', '/', drug, '/output')

#======
# input
#======
#in_filename = 'mcsm_complex1_normalised.csv'
in_filename_params = paste0(tolower(gene), '_all_params.csv') 
infile_params = paste0(outdir, '/', in_filename_params)
cat(paste0('Input file:', infile_params) )

#=======
# output
#=======


#%%===============================================================
###########################
# 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
   , dup_muts)