#!/usr/bin/env Rscript   
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting")
getwd()

#########################################################
# TASK: output barplot by position with each bar coloured by 
# its stability value and active site positions indicated 
# according to colour specified in "subcols_axis_PS.R"
#########################################################

#=======================================================================

############################################################
# 1: Installing and loading required packages and functions
############################################################

#source("~/git/LSHTM_analysis/scripts/Header_TT.R")
library(ggplot2)
library(data.table)
source("../functions/bp_subcolours.R")
#source("subcols_axis.R")
source("subcols_axis.R")

# should return the following dfs, directories and variables
# mut_pos_cols
# mut_pos_cols_lig
# my_df_cols
# my_df_u_cols
# my_df_u_lig_cols
# dup_muts_cols

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(dup_muts_cols, mut_pos_cols, my_df_cols, my_df_u_cols, upos)

#=======================================================================
# !!! very important!!!!
#================
# Inspecting mut_pos_cols
# position numbers and colours and assigning axis colours based on lab_fg
# of the correct df
# open file from desktop ("sample_axis_cols") for cross checking
#================
if ( nrow(mut_pos_cols_lig) == length(unique(my_df_u_cols_lig$position)) ){
  print("PASS: lengths checked, assigning axis colours")
  my_axis_colours = mut_pos_cols_lig$lab_fg
  cat("length of axis colours:", length(my_axis_colours)
      , "\nwhich corresponds to the number of positions on the x-axis of the plot")
}else{
  print("FAIL:lengths mismatch, could not assign axis colours")
  quit()
}

# further sanity checks
table(mut_pos_cols_lig$lab_bg)
check_lab_bg = sum( table(mut_pos_cols_lig$lab_bg) ) == nrow(mut_pos_cols_lig) # should  be True
check_lab_bg    

table(mut_pos_cols_lig$lab_bg2)
check_lab_bg2 = sum( table(mut_pos_cols_lig$lab_bg2) ) ==  nrow(mut_pos_cols_lig) # should  be True
check_lab_bg2

table(mut_pos_cols_lig$lab_fg)
check_lab_fg = sum( table(mut_pos_cols_lig$lab_fg) ) ==  nrow(mut_pos_cols_lig) # should  be True
check_lab_fg

# sanity checks:
if (check_lab_bg && check_lab_bg2 && check_lab_fg) {
  print("PASS: No. of assigned colours match length")
}else{
  print("FAIL: length of assigned colours mismatch")
  quit()
}
#=======================================================================
#=======
# output
#=======
# plot name and location
print(paste0("plot will be in:", plotdir))
bp_aa_subcols_ligand = "barplot_acoloured_LIG.svg"
plot_bp_aa_subcols_ligand = paste0(plotdir, "/", bp_aa_subcols_ligand)

#=======================================================================
#================
# Data for plots
#================
# REASSIGNMENT as necessary
df  = my_df_u_cols_lig

cat("ligand df dim:", dim(df))

# sanity checks
str(df)

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

###########################
# Plot: Ligand affinity
###########################

#==========================
# Barplot with scores (unordered)
# corresponds to ligand_outcome
# Stacked Barplot with colours: ligand_outcome @ position coloured by 
# stability scores. This is a barplot where each bar corresponds 
# to a SNP and is coloured by its corresponding ligand stability value.
# Normalised values (range between -1 and 1 ) to aid visualisation
# NOTE: since barplot plots discrete values, colour = score, so number of
# colours will be equal to the no. of unique normalised scores 
# rather than a continuous scale
# will require generating the colour scale separately.
#============================
# sanity checks
upos = unique(df$position)

table(df$ligand_outcome)
table(df$ligand_outcome)

# add frequency of positions (from lib data.table)
setDT(df)[, pos_count := .N, by = .(position)]

# this is cummulative
table(df$pos_count)

# use group by on this
library(dplyr)
snpsBYpos_df <- df %>%
  group_by(position) %>%
  summarize(snpsBYpos = mean(pos_count))

table(snpsBYpos_df$snpsBYpos)

snp_count = sort(unique(snpsBYpos_df$snpsBYpos))

# sanity checks
# should be a factor
if (is.factor(df$ligand_outcome)){
  print("ligand_outcome is factor")
}else{
  print("converting ligand_outcome to factor")
  df$ligand_outcome = as.factor(df$ligand_outcome)
}

is.factor(df$ligand_outcome)

table(df$ligand_outcome)

# may not be -1 and 1 since these are filtered within 10A
min(df$affinity_scaled)
max(df$affinity_scaled)

# sanity checks
# very important!!!!
tapply(df$affinity_scaled, df$ligand_outcome, min)

tapply(df$affinity_scaled, df$ligand_outcome, max)

# My colour FUNCTION: based on group and subgroup
# in my case;
# df = df
# group = ligand_outcome
# subgroup = normalised score i.e affinity_scaled

# check unique values in normalised data
u = unique(df$affinity_scaled) 
cat("No. of unique values in normalised data:", length(u))

# Define group
# Create an extra column called group which contains the "gp name and score" 
# so colours can be generated for each unique values in this column
my_grp = df$affinity_scaled 
df$group <- paste0(df$ligand_outcome, "_", my_grp, sep = "") 

# Call the function to create the palette based on the group defined above
colours <- ColourPalleteMulti(df, "ligand_outcome", "my_grp")
print(paste0("Colour palette generated for: ", length(colours), " colours"))
my_title = "Ligand affinity"
cat("No. of axis colours: ", length(my_axis_colours))

#========================
# plot with axis colours
#========================
class(df$lab_bg)

# define cartesian coord
my_xlim = length(unique(df$position)); my_xlim

# axis label size
my_xals = 18
my_yals = 18

# axes text size
my_xats = 14
my_yats = 18

#******************
# generate plot: with axis colours
#******************
print(paste0("plot name:", plot_bp_aa_subcols_ligand))

svg(plot_bp_aa_subcols_ligand, width = 26, height = 4)

g = ggplot(df, aes(factor(position, ordered = T)))

OutPlot_aa_LIG = g + 
  coord_cartesian(xlim = c(1, my_xlim)
                  #, ylim = c(0, 6)
                  , ylim = c(0, max(snp_count))
                  , clip = "off") +
  geom_bar(aes(fill = group), colour = "grey") +
  scale_fill_manual(values = colours
                     , guide = "none") +
  geom_tile(aes(,-0.8, width = 0.95, height = 0.85)
            , fill = df$lab_bg) +
  geom_tile(aes(,-1.2, width = 0.95, height = -0.2)
            , fill = df$lab_bg2) +
  
# Here it"s important to specify that your axis goes from 1 to max number of levels
  theme(axis.text.x = element_text(size = my_xats
                                    , angle = 90
                                    , hjust = 1
                                    , vjust = 0.4
                                    , colour = my_axis_colours)
        
         , axis.text.y = element_text(size = my_yats 
                                      , angle = 0
                                      , hjust = 1
                                      , vjust = 0)
         , axis.title.x = element_text(size = my_xals
                                       #, hjust = 1
                                       , vjust = -0.4)
         , axis.title.y = element_text(size = my_yals ) 
         , axis.ticks.x = element_blank()) +
  labs(title = ""
       #title = my_title
       , x = "Position"
       , y = "Frequency")

print(OutPlot_aa_LIG)
dev.off()

#!!!!!!!!!!!!!!!!
#Warning message:
#  Vectorized input to `element_text()` is not officially supported.
#Results may be unexpected or may change in future versions of ggplot2. 
#!!!!!!!!!!!!!!!!!

# for sanity and good practice
#rm(df)
