#!/usr/bin/env Rscript
#########################################################
# TASK: Bubble plot of OR for PS and Lig 

# Output: 1 svg

#=======================================================================
# working dir and loading libraries
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting/")
getwd()


source("~/git/LSHTM_analysis/scripts/Header_TT.R")
require(cowplot)
source("combining_dfs_plotting.R")
# should return the following dfs, directories and variables

# PS combined: 
# 1) merged_df2
# 2) merged_df2_comp
# 3) merged_df3
# 4) merged_df3_comp

# LIG combined: 
# 5) merged_df2_lig
# 6) merged_df2_comp_lig
# 7) merged_df3_lig
# 8) merged_df3_comp_lig

# 9) my_df_u
# 10) my_df_u_lig

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

cat(paste0("Variables imported:"
           , "\ndrug:", drug
           , "\ngene:", gene
           , "\ngene_match:", gene_match
           , "\nAngstrom symbol:", angstroms_symbol
           , "\nNo. of duplicated muts:", dup_muts_nu
           , "\nNA count for ORs:", na_count
           , "\nNA count in df2:", na_count_df2
           , "\nNA count in df3:", na_count_df3))       

#=========================
#=======
# output
#=======
or_combined = "or_combined_PS_LIG.svg"
plot_or_combined  =  paste0(plotdir,"/", or_combined)

or_combined_labelled = "or_combined_PS_LIG_labelled.svg"
plot_or_combined_labelled =  paste0(plotdir,"/", or_combined_labelled)

#or_kin_combined = "or_kin_combined_PS_LIG.svg"
#plot_or_kin_combined  =  paste0(plotdir,"/", or_kin_combined)
#=======================================================================

###########################
# Data for OR and stability plots
# you need merged_df3_comp
# since these are matched 
# to allow pairwise corr
###########################

ps_df  = merged_df3_comp
lig_df = merged_df3_comp_lig

# Ensure correct data type in columns to plot: should be TRUE
is.numeric(ps_df$or_mychisq)
is.numeric(lig_df$or_mychisq)

# delete variables not required
rm(merged_df2, merged_df2_comp, merged_df2_lig, merged_df2_comp_lig, my_df_u, my_df_u_lig)
#%% end of section 1

# sanity check: should be <10
if (max(lig_df$ligand_distance) < 10){
  print ("Sanity check passed: lig data is <10Ang")
}else{
  print ("Error: data should be filtered to be within 10Ang")
}

#############
# Plots: Bubble plot
# x = Position, Y = stability
# size of dots = OR
# col: stability
#############

#-----------------
# Plot 1: DUET vs OR by position as geom_points
#-------------------

my_ats = 20 # axis text size
my_als = 22 # axis label size

# Spelling Correction: made redundant as already corrected at the source
#ps_df$duet_outcome[ps_df$duet_outcome=='Stabilizing'] <- 'Stabilising'
#ps_df$duet_outcome[ps_df$duet_outcome=='Destabilizing'] <- 'Destabilising'

table(ps_df$duet_outcome) ; sum(table(ps_df$duet_outcome))

g1 = ggplot(ps_df, aes(x = factor(position)
                   , y = duet_scaled))

p1 = g1 + 
  geom_point(aes(col = duet_outcome
                 , size = or_mychisq))+
                 #, size = or_kin)) + # not good, almost like log(or)
  theme(axis.text.x = element_text(size = my_ats
                                   , angle = 90
                                   , hjust = 1
                                   , vjust = 0.4)
        , axis.text.y = element_text(size = my_ats
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0)
        , axis.title.x = element_text(size = my_als)
        , axis.title.y = element_text(size = my_als) 
        , legend.text = element_text(size = my_als)
        , legend.title = element_text(size = my_als) ) +
  #, legend.key.size = unit(1, "cm")) +
  labs(title = ""
       , x = "Position"
       , y = "DUET(PS)"
       , size = "Odds Ratio"
       , colour = "DUET Outcome") +
  guides(colour = guide_legend(override.aes = list(size=4))) 

p1 

#=============
#ggplotly
# *so COOL*
fig <- ggplotly(p1)
fig 
#==============
#-------------------
# generate plot 2: Lig vs OR by position as geom_points
#-------------------

# Spelling Correction: made redundant as already corrected at the source

#lig_df$ligand_outcome[lig_df$ligand_outcome=='Stabilizing'] <- 'Stabilising'
#lig_df$ligand_outcome[lig_df$ligand_outcome=='Destabilizing'] <- 'Destabilising'

table(lig_df$ligand_outcome)

g2 = ggplot(lig_df, aes(x = factor(position)
                   , y = affinity_scaled))

p2 = g2 + 
  geom_point(aes(col = ligand_outcome
                 , size = or_mychisq))+
                 #, size = or_kin)) + # not good! almost like log(or)
  theme(axis.text.x = element_text(size = my_ats
                                   , angle = 90
                                   , hjust = 1
                                   , vjust = 0.4)
        , axis.text.y = element_text(size = my_ats
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0)
        , axis.title.x = element_text(size = my_als)
        , axis.title.y = element_text(size = my_als) 
        , legend.text = element_text(size = my_als)
        , legend.title = element_text(size = my_als) ) +
  #, legend.key.size = unit(1, "cm")) +
  labs(title = ""
       , x = "Position"
       , y = "Ligand Affinity"
       , size = "Odds Ratio"
       , colour = "Ligand Outcome"
       ) +
  guides(colour = guide_legend(override.aes = list(size=4))) 

p2

#======================
# combine using cowplot
#======================
#---------------
# no labels
#---------------

svg(plot_or_combined, width = 32, height = 12)
#svg(plot_or_kin_combined, width = 32, height = 12)

theme_set(theme_gray()) # to preserve default theme

Outplot_or_combined = cowplot::plot_grid(plot_grid(p1, p2
                             , ncol = 1
                             , align = 'v'
                             , labels = c("", "")
                             , label_size = my_als+5))
print(Outplot_or_combined)
dev.off()

#---------------
# labels
#---------------
svg(plot_or_combined_labelled, width = 32, height = 12)
#svg(plot_or_kin_combined, width = 32, height = 12)

theme_set(theme_gray()) # to preserve default theme

Outplot_or_combined_labelled = cowplot::plot_grid(plot_grid(p1, p2
                                         , ncol = 1
                                         , align = 'v'
                                         #, labels = c("(a)", "(b)")
                                         , labels = "AUTO"
                                         , label_size = 26))
print(Outplot_or_combined_labelled)
dev.off()