#!/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("../functions/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
#=======

# PS
corr_ps = "corr_PS.svg"
plot_corr_ps =  paste0(plotdir,"/", corr_ps)

# LIG
corr_lig = "corr_LIG.svg"
plot_corr_lig =  paste0(plotdir,"/", corr_lig)

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

#%%%%%%%%%%%%%%%%%%%%%%%%%
df_ps = merged_df3 
df_lig = merged_df3_lig
#%%%%%%%%%%%%%%%%%%%%%%%%%
rm( merged_df2, merged_df2_comp, merged_df2_lig, merged_df2_comp_lig, my_df_u, my_df_u_lig)
########################################################################

################################
# Data for Correlation plots: PS
#################################

#======================
# 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)

#===============================
# Data for Correlation plots:PS
#===============================
# subset data to generate pairwise correlations
cols_to_select =  c("duet_scaled"
                    
                    , "foldx_scaled"
                    
                    , "log10_or_mychisq"
                    , "neglog_pval_fisher"
                    
                    , "af"
                    
                    , "duet_outcome"
                    , drug)

corr_data_ps = df_ps[cols_to_select]

dim(corr_data_ps)

# assign nice colnames (for display)
my_corr_colnames = c("DUET"
                     
                     , "Foldx"
                     
                     , "Log (OR)"
                     , "-Log (P)"

                     , "MAF"
                     
                     , "duet_outcome"
                     , drug)

length(my_corr_colnames)

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

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

#---------------------------------------
# generate corr PS plot: both panels
#---------------------------------------
cat("Corr plot PS DUET with coloured dots:", plot_corr_ps)
svg(plot_corr_ps, width = 15, height = 15)

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))] # foldx colours are reveresed
             , pch = 21 # for bg
             , jitter = T
             , alpha = 1
             , cex = 1.8
             , cex.axis = 2
             , cex.labels = 4
             , cex.cor = 1
             , smooth = F
            
)

dev.off()

corr_ps_rho = corr.test(my_corr_ps[1:5], method = "spearman")$r
corr_ps_p = corr.test(my_corr_ps[1:5], method = "spearman")$p

################################################################################################
###################################
# 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"
                    
                    , "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)"
                     
                     , "MAF"
                     
                     , "ligand_outcome"
                     , drug)

length(my_corr_colnames)

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

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)

#---------------------------------------
# generate corr LIG plot: both panels
#---------------------------------------
cat("Corr LIG plot:", plot_corr_lig)
svg(plot_corr_lig, width = 15, height = 15)

# uncomment as necessary
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 # for bg
                         , jitter = T
                         , cex = 2
                         , cex.axis = 2
                         , cex.labels = 4
                         , cex.cor = 1
                         , smooth = F
)

dev.off()
corr_lig_rho = corr.test(my_corr_lig[1:4], method = "spearman")$r
corr_lig_p = corr.test(my_corr_lig[1:4], method = "spearman")$p
#######################################################