#!/usr/bin/env Rscript  
#########################################################
# TASK: producing boxplots for dr and other muts

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

#source("Header_TT.R")
library(ggplot2)
library(data.table)
library(dplyr)
source("combining_dfs_plotting.R")

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


cols_to_select = c("mutation", "mutationinformation"
                   , "wild_type", "position", "mutant_type"
                   , "mutation_info")

merged_df3_short = merged_df3[, cols_to_select]

# write merged_df3 to generate structural figure
write.csv(merged_df3_short, "merged_df3_short.csv")

#========================================================================
#%%%%%%%%%%%%%%%%%%%
# REASSIGNMENT: PS
#%%%%%%%%%%%%%%%%%%%%
df_ps = merged_df3

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

# name tidying
df_ps$mutation_info = as.factor(df_ps$mutation_info)
df_ps$duet_outcome = as.factor(df_ps$duet_outcome)
df_ps$foldx_outcome  = as.factor(df_ps$foldx_outcome)
df_ps$ligand_outcome  = as.factor(df_ps$ligand_outcome)

# check
table(df_ps$mutation_info)

 # further checks to make sure dr and other muts are indeed unique
dr_muts = df_ps[df_ps$mutation_info == dr_muts_col,]
dr_muts_names = unique(dr_muts$mutation)

other_muts = df_ps[df_ps$mutation_info == other_muts_col,]
other_muts_names = unique(other_muts$mutation)

if ( table(dr_muts_names%in%other_muts_names)[[1]] == length(dr_muts_names) &&
  table(other_muts_names%in%dr_muts_names)[[1]] == length(other_muts_names) ){
  cat("PASS: dr and other muts are indeed unique")
}else{
  cat("FAIL: dr adn others muts are NOT unique!")
  quit()
}


#%%%%%%%%%%%%%%%%%%%
# REASSIGNMENT: LIG
#%%%%%%%%%%%%%%%%%%%%

df_lig = merged_df3_lig

# name tidying
df_lig$mutation_info = as.factor(df_lig$mutation_info)
df_lig$duet_outcome = as.factor(df_lig$duet_outcome)
#df_lig$ligand_outcome  = as.factor(df_lig$ligand_outcome)

# check
table(df_lig$mutation_info)

#========================================================================
#===========
# Data: ps
#===========
# keep similar dtypes cols together
cols_to_select_ps = c("mutationinformation", "mutation", "position", "mutation_info"
                   , "duet_outcome"

                   , "duet_scaled"
                   , "ligand_distance"
                   , "asa"
                   , "rsa"
                   , "rd_values"
                   , "kd_values")

df_wf_ps = df_ps[, cols_to_select_ps]

pivot_cols_ps = cols_to_select_ps[1:5]; pivot_cols_ps

expected_rows_lf_ps = nrow(df_wf_ps) * (length(df_wf_ps) - length(pivot_cols_ps))
expected_rows_lf_ps

# LF data: duet
df_lf_ps = gather(df_wf_ps, param_type, param_value, duet_scaled:kd_values, factor_key=TRUE)

if (nrow(df_lf_ps) == expected_rows_lf_ps){
  cat("PASS: long format data created for duet")
}else{
  cat("FAIL: long format data could not be created for duet")
  exit()
}

str(df_wf_ps)
str(df_lf_ps)

# assign pretty labels: param_type
levels(df_lf_ps$param_type); table(df_lf_ps$param_type)

ligand_dist_colname = paste0("Distance to ligand (", angstroms_symbol, ")")
ligand_dist_colname

duet_stability_name = paste0(delta_symbol, delta_symbol, "G")
duet_stability_name
  
#levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="duet_scaled"] <- "Stability"
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="duet_scaled"] <- duet_stability_name
#levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="ligand_distance"] <- "Ligand Distance"
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="ligand_distance"] <- ligand_dist_colname
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="asa"] <- "ASA"
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="rsa"] <- "RSA"
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="rd_values"] <- "RD"
levels(df_lf_ps$param_type)[levels(df_lf_ps$param_type)=="kd_values"] <- "KD"
# check
levels(df_lf_ps$param_type); table(df_lf_ps$param_type)

# assign pretty labels: mutation_info
levels(df_lf_ps$mutation_info); table(df_lf_ps$mutation_info)
sum(table(df_lf_ps$mutation_info)) == nrow(df_lf_ps)

levels(df_lf_ps$mutation_info)[levels(df_lf_ps$mutation_info)==dr_muts_col] <- "DM"
levels(df_lf_ps$mutation_info)[levels(df_lf_ps$mutation_info)==other_muts_col] <- "OM"
# check
levels(df_lf_ps$mutation_info); table(df_lf_ps$mutation_info)

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

#===========
# LF data: LIG
#===========
# keep similar dtypes cols together
cols_to_select_lig = c("mutationinformation", "mutation", "position", "mutation_info"
                       , "ligand_outcome"
                       
                       , "affinity_scaled"
                       #, "ligand_distance"
                       , "asa"
                       , "rsa"
                       , "rd_values"
                       , "kd_values")

df_wf_lig = df_lig[, cols_to_select_lig]

pivot_cols_lig = cols_to_select_lig[1:5]; pivot_cols_lig

expected_rows_lf_lig = nrow(df_wf_lig) * (length(df_wf_lig) - length(pivot_cols_lig))
expected_rows_lf_lig

# LF data: foldx
df_lf_lig = gather(df_wf_lig, param_type, param_value, affinity_scaled:kd_values, factor_key=TRUE)

if (nrow(df_lf_lig) == expected_rows_lf_lig){
  cat("PASS: long format data created for foldx")
}else{
  cat("FAIL: long format data could not be created for foldx")
  exit()
}

# assign pretty labels: param_type
levels(df_lf_lig$param_type); table(df_lf_lig$param_type)

levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="affinity_scaled"] <- "Ligand Affinity"
#levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="ligand_distance"] <- "Ligand Distance"
levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="asa"] <- "ASA"
levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="rsa"] <- "RSA"
levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="rd_values"] <- "RD"
levels(df_lf_lig$param_type)[levels(df_lf_lig$param_type)=="kd_values"] <- "KD"
#check
levels(df_lf_lig$param_type); table(df_lf_lig$param_type)

# assign pretty labels: mutation_info
levels(df_lf_lig$mutation_info); table(df_lf_lig$mutation_info)
sum(table(df_lf_lig$mutation_info)) == nrow(df_lf_lig)

levels(df_lf_lig$mutation_info)[levels(df_lf_lig$mutation_info)==dr_muts_col] <- "DM"
levels(df_lf_lig$mutation_info)[levels(df_lf_lig$mutation_info)==other_muts_col] <- "OM"
# check
levels(df_lf_lig$mutation_info); table(df_lf_lig$mutation_info)

#############################################################################
#===========
# Data: foldx
#===========
# keep similar dtypes cols together
cols_to_select_foldx = c("mutationinformation", "mutation", "position", "mutation_info"
                      , "foldx_outcome"
                      
                      , "foldx_scaled")
                      #, "ligand_distance"
                      #, "asa"
                      #, "rsa"
                      #, "rd_values"
                      #, "kd_values")


df_wf_foldx = df_ps[, cols_to_select_foldx]

pivot_cols_foldx = cols_to_select_foldx[1:5]; pivot_cols_foldx
  
expected_rows_lf_foldx = nrow(df_wf_foldx) * (length(df_wf_foldx) - length(pivot_cols_foldx))
expected_rows_lf_foldx

# LF data: foldx
df_lf_foldx = gather(df_wf_foldx, param_type, param_value, foldx_scaled, factor_key=TRUE)

if (nrow(df_lf_foldx) == expected_rows_lf_foldx){
  cat("PASS: long format data created for foldx")
}else{
  cat("FAIL: long format data could not be created for foldx")
  exit()
}

foldx_stability_name = paste0(delta_symbol, delta_symbol, "G")
foldx_stability_name

# assign pretty labels: param type
levels(df_lf_foldx$param_type); table(df_lf_foldx$param_type)

#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="foldx_scaled"] <- "Stability"
levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="foldx_scaled"] <- foldx_stability_name
#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="ligand_distance"] <- "Ligand Distance"
#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="asa"] <- "ASA"
#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="rsa"] <- "RSA"
#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="rd_values"] <- "RD"
#levels(df_lf_foldx$param_type)[levels(df_lf_foldx$param_type)=="kd_values"] <- "KD"
# check
levels(df_lf_foldx$param_type); table(df_lf_foldx$param_type)

# assign pretty labels: mutation_info
levels(df_lf_foldx$mutation_info); table(df_lf_foldx$mutation_info)
sum(table(df_lf_foldx$mutation_info)) == nrow(df_lf_foldx)

levels(df_lf_foldx$mutation_info)[levels(df_lf_foldx$mutation_info)==dr_muts_col] <- "DM"
levels(df_lf_foldx$mutation_info)[levels(df_lf_foldx$mutation_info)==other_muts_col] <- "OM"
# check
levels(df_lf_foldx$mutation_info); table(df_lf_foldx$mutation_info)

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

# clear excess variables
rm(cols_to_select_ps, cols_to_select_foldx, cols_to_select_lig
   , pivot_cols_ps, pivot_cols_foldx, pivot_cols_lig
  , expected_rows_lf_ps, expected_rows_lf_foldx, expected_rows_lf_lig
  , my_max, my_min, na_count, na_count_df2, na_count_df3, dup_muts_nu
  , c1, c2, n)