#!/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)

#===========
# input
#===========
source("combining_dfs_plotting.R")

#===========
# output
#===========

logo_plot = "logo_plot.svg"
plot_logo_plot = paste0(plotdir,"/", logo_plot)

#==========================
# This will return:

# df with NA for pyrazinamide:
# merged_df2
# merged_df3

# df without NA for pyrazinamide:
# merged_df2_comp
# merged_df3_comp
#===========================

###########################
# Data for plots
# you need merged_df2 or merged_df2_comp
# since this is one-many relationship 
# i.e the same SNP can belong to multiple lineages
# using the _comp dataset means
# we lose some muts and at this level, we should use
# as much info as available, hence use df with NA

# This will the first plotting df
# Then subset this to extract dr muts only (second plottig df)
###########################

#%%%%%%%%%%%%%%%%%%%%%%%%%
# uncomment as necessary
# REASSIGNMENT
#my_data = merged_df2
#my_data =  merged_df2_comp
#my_data = merged_df3
my_data = merged_df3_comp
#%%%%%%%%%%%%%%%%%%%%%%%%%%

# delete variables not required
rm(merged_df2, merged_df2_comp, merged_df3, merged_df3_comp)

# quick checks
colnames(my_data)
str(my_data)

c1 = unique(my_data$position) 
nrow(my_data)
cat("No. of rows in my_data:", nrow(my_data)
    , "\nDistinct positions corresponding to snps:", length(c1)
    , "\n===========================================================")

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# FIXME: Think and decide what you want to remove
# mut_pos_occurence < 1 or sample_pos_occurrence <1

# get freq count of positions so you can subset freq<1
require(data.table)
#setDT(my_data)[, mut_pos_occurrence := .N, by = .(position)] #265, 14

#extract freq_pos>1
#my_data_snp = my_data[my_data$occurrence!=1,]

#u = unique(my_data_snp$position) #73
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# REASSIGNMENT to prevent changing code
my_data_snp = my_data
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#=======================================================================
#%% logo plots from dataframe

#############
# PLOTS: ggseqlogo with custom height
# https://omarwagih.github.io/ggseqlogo/
#############
#require(ggplot2)
#require(tidyverse)
library(ggseqlogo)

foo = my_data_snp[, c("position", "mutant_type","duet_scaled", "or_mychisq"
                      , "mut_prop_polarity", "mut_prop_water") ] 

my_data_snp$log10or = log10(my_data_snp$or_mychisq)
bar = my_data_snp[, c("position", "mutant_type", "or_mychisq", "log10or")]

bar_or = my_data_snp[, c("position", "mutant_type", "or_mychisq")]
wide_df_or <- bar_or %>% spread(position, or_mychisq, fill = 0)
wide_df_or = as.matrix(wide_df_or)
rownames(wide_df_or) = wide_df_or[,1]
wide_df_or = wide_df_or[,-1]

# custom height (OR) logo plot: yayy works
ggseqlogo(wide_df_or, method="custom", seq_type="aa") + ylab("my custom height") +
  theme(legend.position = "bottom"
        , axis.text.x = element_text(size = 11
                                     , angle = 90
                                     , hjust = 1
                                     , vjust = 0.4)
        , axis.text.y = element_text(size = 15
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0))+
  labs(title = "AA logo plot"
       , x = "Wild-type position"
       , y = "OR")
#%% end of logo plot with OR as height
#=======================================================================
# extracting data with log10R
bar_logor = my_data_snp[, c("position", "mutant_type", "log10or")]
wide_df_logor <- bar_logor %>% spread(position, log10or, fill = 0)

wide_df_logor = as.matrix(wide_df_logor)
rownames(wide_df_logor) = wide_df_logor[,1]
wide_df_logor = wide_df_logor[,-1]

#  custom height (log10OR) logo plot: yayy works
ggseqlogo(wide_df_logor, method="custom", seq_type="aa") + ylab("my custom height") +
  theme(legend.position = "bottom"
        , axis.text.x = element_text(size = 11
                                     , angle = 90
                                     , hjust = 1
                                     , vjust = 0.4)
        , axis.text.y = element_text(size = 15
                                     , angle = 0
                                     , hjust = 1
                                     , vjust = 0))+
  labs(title = "AA logo plot"
       , x = "Wild-type position"
       , y = "Log10(OR)")

#=======================================================================
#%% logo plot from sequence

#################
# Plot: LOGOLAS (ED plots)
# link: https://github.com/kkdey/Logolas
# on all pncA samples: output of mutate.py
################
library(Logolas)

# data was pnca_msa.txt
#FIXME: generate this file 

seqs = read.csv("~/git//Data/pyrazinamide/snp_seqsfile.txt"
                , header = FALSE
                , stringsAsFactors = FALSE)$V1


# my_data: useful!
logomaker(seqs, type = "EDLogo", color_type = "per_symbol"
          , return_heights = TRUE)
logomaker(seqs, type = "Logo", color_type = "per_symbol")

#%% end of script
#=======================================================================