#!/usr/bin/env Rscript
#########################################################
# TASK: Corr plots 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")
source("my_pairs_panel.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
#=======
# can't combine by cowplot because not ggplots
#corr_plot_combined = "corr_combined.svg"
#plot_corr_plot_combined  =  paste0(plotdir,"/", corr_plot_combined)

# PS
corr_ps_adjusted = "corr_PS_adjusted.svg"
plot_corr_ps_adjusted =  paste0(plotdir,"/", corr_ps)

# LIG
corr_lig_adjusted = "corr_LIG_adjusted.svg"
plot_corr_lig_adjusted =  paste0(plotdir,"/", corr_lig)

####################################################################
#               end of loading libraries and functions                 #
########################################################################

#%%%%%%%%%%%%%%%%%%%%%%%%%
df_ps = merged_df3_comp 
df_lig = merged_df3_comp_lig
#%%%%%%%%%%%%%%%%%%%%%%%%%

rm( merged_df2, merged_df2_comp, merged_df2_lig, merged_df2_comp_lig, my_df_u, my_df_u_lig)

########################################################################
#               end of data extraction and cleaning for plots          #
########################################################################

#===========================
# Data for Correlation plots:PS
#===========================
table(df_ps$duet_outcome)


#===========================
# Data for Correlation plots:foldx
#===========================
#============================
# adding foldx scaled values
# scale data b/w -1 and 1
#============================
n = which(colnames(df_ps) == "ddg"); n 

my_min = min(df_ps[,n]); my_min 
my_max = max(df_ps[,n]); my_max 

df_ps$foldx_scaled = ifelse(df_ps[,n] < 0
                            , df_ps[,n]/abs(my_min)
                            , df_ps[,n]/my_max) 
# sanity check
my_min = min(df_ps$foldx_scaled); my_min 
my_max = max(df_ps$foldx_scaled); my_max

if (my_min == -1 && my_max == 1){
  cat("PASS: foldx ddg successfully scaled b/w -1 and 1"
      , "\nProceeding with assigning foldx outcome category")
}else{
  cat("FAIL: could not scale foldx ddg values"
      , "Aborting!")
}


#================================
# adding foldx outcome category
# ddg<0 = "Stabilising" (-ve)
#=================================

c1 = table(df_ps$ddg < 0)
df_ps$foldx_outcome = ifelse(df_ps$ddg < 0, "Stabilising", "Destabilising")
c2 = table(df_ps$ddg < 0)

if ( all(c1 == c2) ){
  cat("PASS: foldx outcome successfully created")
}else{
  cat("FAIL: foldx outcome could not be created. Aborting!")
  exit()
}

table(df_ps$foldx_outcome)


#======================
# adding log cols 
#======================
 
df_ps$log10_or_mychisq = log10(df_ps$or_mychisq)
df_ps$neglog_pval_fisher = -log10(df_ps$pval_fisher)

df_ps$log10_or_kin = log10(df_ps$or_kin)
df_ps$neglog_pwald_kin = -log10(df_ps$pwald_kin)

# subset data to generate pairwise correlations
cols_to_select =  c("duet_scaled"
                    
                    , "foldx_scaled"
                    
                    #, "log10_or_mychisq"
                    #, "neglog_pval_fisher"
                    
                    , "or_kin"
                    , "neglog_pwald_kin"
                    
                    , "af"
                    
                    , "asa"
                    , "rsa"
                    , "kd_values"
                    , "rd_values"
                    
                    , "duet_outcome"
                    , drug)

corr_data_ps = df_ps[, cols_to_select]

dim(corr_data_ps)

#p_italic = substitute(paste("-Log(", italic('P'), ")"));p_italic 
#p_adjusted_italic = substitute(paste("-Log(", italic('P adjusted'), ")"));p_adjusted_italic

# assign nice colnames (for display)
my_corr_colnames = c("DUET"
                     
                     , "Foldx"
                     #, "Log(OR)"
                     #, "-Log(P)"
                     
                     , "OR adjusted"
                     , "-Log(P wald)"
                     
                     , "AF"
                     
                     , "ASA"
                     , "RSA"
                     , "KD"
                     , "RD"
                     
                     , "duet_outcome"
                     , drug)

length(my_corr_colnames)

colnames(corr_data_ps)
colnames(corr_data_ps) <- my_corr_colnames
colnames(corr_data_ps)

#-----------------
# generate corr PS plot
#-----------------
start = 1
end = which(colnames(corr_data_ps) == drug); end # should be the last column
offset = 1

my_corr_ps = corr_data_ps[start:(end-offset)]
head(my_corr_ps)

#my_cols = c("#f8766d", "#00bfc4")
# deep blue :#007d85
# deep red: #ae301e

cat("Corr plot PS:", plot_corr_ps_adjusted)
svg(plot_corr_ps_adjusted, width = 15, height = 15)

OutPlot1 = pairs.panels(my_corr_ps[1:(length(my_corr_ps)-1)]
             , method = "spearman" # correlation method
             , hist.col = "grey" ##00AFBB
             , density = TRUE  # show density plots
             , ellipses = F # show correlation ellipses
             , stars = T
             , rug = F
             , breaks = "Sturges"
             , show.points = T
             , bg = c("#f8766d", "#00bfc4")[unclass(factor(my_corr_ps$duet_outcome))]
             , pch = 21
             , jitter = T
             #, alpha = .05
             #, points(pch = 19, col = c("#f8766d", "#00bfc4"))
             , cex = 2
             , cex.axis = 1.5
             , cex.labels = 1.5
             , cex.cor = 1
             , smooth = F
)

print(OutPlot1)
dev.off()

#===========================
# Data for Correlation plots: LIG
#===========================
table(df_lig$ligand_outcome)

df_lig$log10_or_mychisq = log10(df_lig$or_mychisq)
df_lig$neglog_pval_fisher = -log10(df_lig$pval_fisher)


df_lig$log10_or_kin = log10(df_lig$or_kin)
df_lig$neglog_pwald_kin = -log10(df_lig$pwald_kin)


# subset data to generate pairwise correlations
cols_to_select =  c("affinity_scaled"
                    
                    , "log10_or_mychisq"
                    , "neglog_pval_fisher"
                    
                    #, "or_kin"
                    #, "neglog_pwald_kin"
                    
                    , "af"
                    
                    , "ligand_outcome"
                    , drug)

corr_data_lig = df_lig[, cols_to_select]


dim(corr_data_lig)

# assign nice colnames (for display)
my_corr_colnames = c("Ligand Affinity"
                     
                     , "Log(OR)"
                     , "-Log(P)"
                     
                     #, "OR adjusted"
                     #, "-Log(P wald)"
                     
                     , "AF"
                     
                     , "ligand_outcome"
                     , drug)

length(my_corr_colnames)

colnames(corr_data_lig)
colnames(corr_data_lig) <- my_corr_colnames
colnames(corr_data_lig)

#-----------------
# generate corr LIG plot
#-----------------

start = 1
end = which(colnames(corr_data_lig) == drug); end # should be the last column
offset = 1

my_corr_lig = corr_data_lig[start:(end-offset)]
head(my_corr_lig)

cat("Corr LIG plot:", plot_corr_lig_adjusted)
svg(plot_corr_lig_adjusted, width = 15, height = 15)

OutPlot2  = pairs.panels(my_corr_lig[1:(length(my_corr_lig)-1)]
                  , method = "spearman" # correlation method
                  , hist.col = "grey" ##00AFBB
                  , density = TRUE  # show density plots
                  , ellipses = F # show correlation ellipses
                  , stars = T
                  , rug = F
                  , breaks = "Sturges"
                  , show.points = T
                  , bg = c("#f8766d", "#00bfc4")[unclass(factor(my_corr_lig$ligand_outcome))]
                  , pch = 21
                  , jitter = T
                  #, alpha = .05
                  #, points(pch = 19, col = c("#f8766d", "#00bfc4"))
                  , cex = 3
                  , cex.axis = 2.5
                  , cex.labels = 2.1
                  , cex.cor = 1
                  , smooth = F
)

print(OutPlot2)
dev.off()
#######################################################


