#!/usr/bin/env Rscript
#########################################################
# TASK: producing logo-type plot showing
  ## OR
  ## Log OR
# multiple muts (>1 per position) coloured by aa property
  ## mutant_type
  ## wild_type

# Note: Replaced scripts:
  ## logo_plot.R
  ## logo_multiple_muts.R
  ## logo_combined.R
#########################################################
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting")
getwd()
source("Header_TT.R")

spec = matrix(c(
  "drug"       , "d",  1, "character",
  "gene"       , "g",  1, "character",
  "data_file1" , "fa", 2, "character",
  "data_file2" , "fb", 2, "character" 
), byrow = TRUE, ncol = 4)

opt = getopt(spec)

drug            = opt$drug
gene            = opt$gene
infile_params   = opt$data_file1
infile_metadata = opt$data_file2

if(is.null(drug)|is.null(gene)) {
  stop("Missing arguments: --drug and --gene must both be specified (case-sensitive)")
}

#===========
# Input
#===========
source("get_plotting_dfs.R")

#===========
# output
#===========
logo_or_plotname = "logo_or_plot.svg"
plot_logo_or = paste0(plotdir,"/", logo_or_plotname)

logo_logOR_plotname = "logo_logOR_plot.svg"
plot_logo_logOR = paste0(plotdir,"/", logo_logOR_plotname)

logo_multiple_muts = "logo_multiple_muts.svg"
plot_logo_multiple_muts = paste0(plotdir,"/", logo_multiple_muts)

logo_combined_labelled = "logo_combined_labelled.svg"
plot_logo_combined_labelled  = paste0(plotdir,"/", logo_combined_labelled)

#########################################################
#logo_or_mult_p + theme_dark()
#logo_or_mult_p + theme(plot.background = element_rect(fill = "black"))

#==================================
# Output
# Logo plot OR: custom height (OR)   
#==================================
cat("Logo plot with OR as y axis:", plot_logo_or)
svg(plot_logo_or, width = 30 , height = 6)

logo_or = ggseqlogo(wide_df_or
                    , method = "custom"
                    , seq_type = "aa") + ylab("my custom height") +
  theme( axis.text.x = element_text(size = x_ats
                                    , angle = x_tangle
                                    , hjust = 1
                                    , vjust = 0.4)
         , axis.text.y = element_text(size = y_ats
                                      , angle = y_tangle
                                      , hjust = 1
                                      , vjust = 0)
         , axis.title.y = element_text(size = x_tts)
         , axis.title.x = element_text(size = y_tts)
         #, legend.position = "bottom") + 
         , legend.position = leg_pos)+
  #, legend.text = element_text(size = leg_ts)
  #, legend.title = element_text(size = leg_tts))+
  scale_x_discrete(x_lab
                   #, breaks
                   , labels = position_or
                   , limits = factor(1:length(position_or))) +
  ylab(y_lab)

print(logo_or)
#dev.off()

#=============================================
# Output
# Logo plot LOG OR: custom height (Log10 OR) 
#=============================================
cat("Logo plot with log10 OR as y axis:", plot_logo_logOR)
svg(plot_logo_logOR, width = 30 , height = 6)

logo_logOR = ggseqlogo(wide_df_logor
                       , method = "custom"
                       , seq_type="aa") + ylab("my custom height") +
  theme(legend.position = "bottom"
        , axis.text.x = element_text(size = 13
                                     , angle = 90
                                     , hjust = 1
                                     , vjust = 0.4)
        , axis.text.y = element_text(size = 20
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0)
        , axis.title.y = element_text(size = 25)
        , axis.title.x = element_text(size = 20))+
  scale_x_discrete("Position"
                   #, breaks
                   , labels = position_logor
                   , limits = factor(1:length(position_logor)))+
  ylab("Log (Odds Ratio)") + 
  scale_y_continuous(limits = c(0, 9))

print(logo_logOR)
#dev.off()

#*****************************
# Mutant logo plot: >1 nsSNP
#******************************
aa_col_choices = c('chemistry', 'hydrophobicity', 'clustalx', 'taylor')
my_logo_col = aa_col_choices[[1]]

if (my_logo_col == 'clustalx || taylor'){
  cat("\nSelected colour scheme:", my_logo_col
      , "\nUsing black theme\n")
  theme_bgc = "black"
  font_bgc  = "white"
} if (my_logo_col == 'chemistry || hydrophobicity') {
  cat('\nSelected colour scheme:', my_logo_col
      , "\nUsing grey theme")
  theme_bgc = "grey"
  font_bgc  = "black"
}

p0 = ggseqlogo(tab_mt
               , method = 'custom'
               , col_scheme = my_logo_col
               , seq_type = 'aa') + 
  #ylab('my custom height') +
  theme(axis.text.x = element_blank()) +
  theme(text=element_text(family="FreeSans"))+
  theme_logo()+ 
  scale_x_continuous(breaks = 1:ncol(tab_mt)
                     , labels = colnames(tab_mt))+
  scale_y_continuous( breaks = 1:max_mult_mut
                      , limits = c(0, max_mult_mut))

#p0
cat('\nDone: p0')

# further customisation
mut_logo_p = p0 + theme(legend.position = "none"
                , legend.title = element_blank()
                , legend.text = element_text(size = 20)
                , axis.text.x = element_text(size = 14
                                             , angle = 90
                                             , colour = font_bgc)
                , axis.text.y = element_blank()
                , plot.background = element_rect(fill = theme_bgc))
mut_logo_p
cat('\nDone: p0+mut_logo_p')

#*************************
# Wild logo plot: >1 nsSNP
#*************************
p2 = ggseqlogo(tab_wt
               , method = 'custom'
               , seq_type = 'aa'
               , col_scheme = my_logo_col) + 
  #ylab('my custom height') +
  theme(text=element_text(family="FreeSans"))+
  theme(axis.text.x = element_blank()
        , axis.text.y = element_blank()) +
  theme_logo() +
  scale_x_continuous(breaks = 1:ncol(tab_wt)
                     , labels = colnames(tab_wt)) 
#p2
cat('\nDone: p2')

# further customise
wt_logo_p = p2 +
  theme(legend.position = "bottom"
        #, legend.title = element_blank()
        , legend.title = element_text("Amino acid properties"
                                      , size = 20)
        , legend.text = element_text(size = 20)
        , axis.text.x = element_text(size = 14, angle = 90)
        , axis.text.y = element_blank()
        , axis.title.x = element_text(size = 22))+
  
  labs(x= "Position")

#wt_logo_p
cat('\nDone: wt_logo_p')

#***********************
# Logo OR >1 nsSNP
#***********************
logo_or_mult_p = ggseqlogo(wide_df_or_mult
                           , method = "custom"
                           , col_scheme = my_logo_col
                           , seq_type="aa") + 
  ylab("my custom height") +
  theme(axis.text.x = element_text(size = 14
                                   , angle = 90
                                   , hjust = 1
                                   , vjust = 0.4)
        , axis.text.y = element_text(size = 18
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0)
        , axis.title.y = element_text(size = 18)
        , axis.title.x = element_blank()
        , legend.position = "none")+
  scale_x_discrete( labels = position_or_mult
                    , limits = factor(1:length(position_or_mult))) +
  scale_y_discrete(breaks = c(50, 150, 250, 350)
                   , labels = c(50, 150, 250, 350)
                   , limits = c(50, 150, 250, 350)
  ) +
  xlab("Position") +
  ylab("Odds Ratio")

#logo_or_mult_p
cat('\nDone: logo_or_mult_p')

#=========================================
# Output
# Combined plot: logo_mutliple_muts.svg
#=========================================
#suppressMessages( require(cowplot) )
cat('\nDone: wt_logo_p')

#plot_grid(p1, p3, ncol = 1, align = 'v') 
cat('\nDone: mut_logo_p + wt_logo_p')

# colour scheme: https://rdrr.io/cran/ggseqlogo/src/R/col_schemes.r
cat("\nOutput plot:", plot_logo_multiple_muts, "\n")
svg(plot_logo_multiple_muts, width = 32, height = 10)

mult_muts_combined = cowplot::plot_grid(mut_logo_p, wt_logo_p
                              , nrow = 2
                              , align = "v"
                              , rel_heights = c(3/4, 1/4))

print(mult_muts_combined)
#dev.off()

#======================================================
# Output
# Combined ALL logo plots: logOR, mut_logo and mut_logo
#=====================================================
cat("Output plot:", plot_logo_combined_labelled)
svg(plot_logo_combined_labelled, width = 25, height = 10)

all_logo_plots = cowplot::plot_grid(logo_or_mult_p
                                    , mut_logo_p
                                    , wt_logo_p
                                    , nrow = 3
                                    , align = "hv"
                                    #, labels = c("(a)","(b)", "(c)")
                                    , labels = "AUTO"
                                    , rel_heights = c(3/8, 3/8, 1.5/8)
                                    , rel_widths = c(0.85, 1, 1)
                                    , label_size = 25)

print(all_logo_plots)
#dev.off()