#!/usr/bin/env Rscript
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting")
getwd()
source("~/git/LSHTM_analysis/scripts/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_multiple_muts = "logo_multiple_muts.svg"
plot_logo_multiple_muts = paste0(plotdir,"/", logo_multiple_muts)

#*********************
# Plot 1: mutant logo
#*********************
p0 = ggseqlogo(tab_mt
               , method = 'custom'
               , 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 = 17, angle = 90)
                        , axis.text.y = element_blank())
#mut_logo_p
cat('\nDone: p0+mut_logo_p')

#***********************
# Plot 2: wild_type logo
#***********************
p2 = ggseqlogo(tab_wt
               , method = 'custom'
               , seq_type = 'aa'
               #, col_scheme = "taylor"
               #, col_scheme = chemistry2
) + 
  #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 done')

# 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 = 17, 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')

#------------------------------------
# Now combine using cowplot
# which ensures the plots are aligned
#------------------------------------
#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()