#logo plots

# one for multiple muts
  # --> select/drop down option to filter count of nsSNPs
  # --> select/drop down  option for colour
  # --> should include WT

# Data used

#tab_mt # mutant logo plot
#tab_wt # wt logo plot

# Make it hover over position and then get the corresponding data table!
#%%======================================================================

#==================
# logo data: OR
#==================
# NOTE: my_logo_col

LogoPlotMSA <- function(msaSeq_mut
                        , msaSeq_wt
                        , msa_method = 'bits' # or probability
                     , my_logo_col = "chemistry" 
                     , x_lab = "Wild-type position"
                     , y_lab = "Count"
                     , x_ats = 13 # text size
                     , x_tangle = 90 # text angle
                     , y_ats = 13
                     , y_tangle = 0
                     , x_tts = 13 # title size
                     , y_tts = 13
                     , leg_pos = "top" # can be top, left, right and bottom or c(0.8, 0.9)
                     , leg_dir = "horizontal" #can be vertical or horizontal
                     , leg_ts = 16 # leg text size
                     , leg_tts = 16 # leg title size
                     )

{
   ############################################
   # Data processing for logo plot for nsSNPS
   ############################################

  cat("\nLength of MSA", nrow(msaSeq_mut)
      , "\nlength of WT seq:", nrow(msaSeq_wt))

   ######################################
   # Generating plots for muts and wt
   #####################################
   LogoPlotMSAL <- list()
  
   if (my_logo_col %in% c('clustalx','taylor')) {
    cat("\nSelected colour scheme:", my_logo_col
        , "\nUsing black theme\n")
    
    theme_bgc  = "black"
    xfont_bgc  = "white"
    yfont_bgc  = "white"
    xtt_col    = "white"
    ytt_col    = "white"
  }
  
  if (my_logo_col %in% c('chemistry', 'hydrophobicity')) {
    cat('\nSelected colour scheme:', my_logo_col
        , "\nUsing grey theme")
    
    theme_bgc = "grey"
    xfont_bgc  = "black"
    yfont_bgc  = "black"
    xtt_col    = "black"
    ytt_col    = "black"
  }
  
  #####################################
  # Generating logo plots for nsSNPs
  #####################################
  
  #-------------------
  # Mutant logo plot
  #-------------------
  p0 = ggseqlogo(msaSeq_mut
                 #msaSeq_mut$V1
                 , facet = "grid"
                 , method = msa_method
                 , col_scheme = my_logo_col
                 , seq_type = 'aa') 
    
  # further customisation
    msa_mut_logo_P = p0 + theme(legend.position = leg_pos
                          , legend.direction = leg_dir
                          #, legend.title = element_blank()
                          , legend.title = element_text(size = leg_tts
                                                        , colour = ytt_col)
                          , legend.text = element_text(size = leg_ts)

                          , axis.text.x = element_text(size = x_ats
                                                     , angle = x_tangle
                                                     , hjust = 1
                                                     , vjust = 0.4
                                                     , colour = xfont_bgc)
                          #, axis.text.y = element_blank()
                          , axis.text.y = element_text(size = y_ats
                                                       , angle = y_tangle
                                                       , hjust = 1
                                                       , vjust = -1.0
                                                       , colour = yfont_bgc)
                          , axis.title.x = element_text(size = x_tts
                                                        , colour = xtt_col)
                          , axis.title.y = element_text(size = y_tts
                                                        , colour = ytt_col)
                          
                          , plot.background = element_rect(fill = theme_bgc))
  
  cat('\nDone: msa_mut_logo_P')
  #return(msa_mut_logoP)
  LogoPlotMSAL[['msa_mut_logoP']] <- msa_mut_logo_P

  #---------------------------------
  # Wild-type MSA: gene_fasta file
  #---------------------------------
  p1 = ggseqlogo(msaSeq_wt
                 #msaSeq_wt$V1
                 , facet = "grid"
                 , method = msa_method
                 , col_scheme = my_logo_col
                 , seq_type = 'aa')
  
  # further customisation
    msa_wt_logo_P = p1 + 
    
    theme(legend.position = "none"
          , legend.direction = leg_dir
          #, legend.title = element_blank()
          , legend.title = element_text(size = leg_tts
                                        , colour = ytt_col)
          , legend.text = element_text(size = leg_ts)
          
          , axis.text.x = element_text(size = x_ats
                                       , angle = x_tangle
                                       , hjust = 1
                                       , vjust = 0.4
                                       , colour = xfont_bgc)
          , axis.text.y = element_blank()
          
          , axis.title.x = element_text(size = x_tts
                                        , colour = xtt_col)
          , axis.title.y = element_text(size = y_tts
                                        , colour = ytt_col)
          
          , plot.background = element_rect(fill = theme_bgc)) +
      ylab("")
  
  cat('\nDone: msa_wt_logo_P')
  #return(msa_wt_logoP)
  LogoPlotMSAL[['msa_wt_logoP']] <- msa_wt_logo_P
  
  #=========================================
  # Output
  # Combined plot: logo_MSA
  #=========================================

  cat('\nDone: msa_mut_logoP + msa_wt_logoP')

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

  LogoMSA_comb = cowplot::plot_grid(LogoPlotMSAL[['msa_mut_logoP']]
                                     , LogoPlotMSAL[['msa_wt_logoP']]
                                          , nrow = 2
                                          , align = "v"
                                          , rel_heights = c(3/4, 1/4))

  return(LogoMSA_comb)

}