#!/usr/bin/env Rscript  
#########################################################
# TASK: producing logoplot 
# from data and/or from sequence
#########################################################
# working dir and loading libraries
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting")
getwd()

source("Header_TT.R")
source("../functions/plotting_globals.R")
source("../functions/plotting_data.R")
source("../functions/combining_dfs_plotting.R")
###########################################################
# command line args
#********************
drug = 'streptomycin'
gene = 'gid'

#===========
# input
#===========
#---------------------
# call: import_dirs()
#---------------------
import_dirs(drug, gene)

#---------------------------
# call: plotting_data()
#---------------------------
if (!exists("infile_params") && exists("gene")){
  #if (!is.character(infile_params) && exists("gene")){
  #in_filename_params = paste0(tolower(gene), "_all_params.csv") 
  in_filename_params = paste0(tolower(gene), "_comb_afor.csv") # part combined for gid
  infile_params = paste0(outdir, "/", in_filename_params)
  cat("\nInput file for mcsm comb data not specified, assuming filename: ", infile_params, "\n")
}

# Input 1: read <gene>_comb_afor.csv
pd_df = plotting_data(infile_params)
my_df_u       = pd_df[[1]] # this forms one of the input for combining_dfs_plotting()

#--------------------------------
# call: combining_dfs_plotting()
#--------------------------------
if (!exists("infile_metadata") && exists("gene")){
  #if (!is.character(infile_params) && exists("gene")){{
  in_filename_metadata = paste0(tolower(gene), "_metadata.csv") # part combined for gid
  infile_metadata = paste0(outdir, "/", in_filename_metadata)
  cat("\nInput file for gene metadata not specified, assuming filename: ", infile_metadata, "\n")
}

# Input 2: read <gene>_meta data.csv
cat("\nReading meta data file:", infile_metadata)

gene_metadata <- read.csv(infile_metadata
                          , stringsAsFactors = F
                          , header = T)

all_plot_dfs = combining_dfs_plotting(my_df_u
                                      , gene_metadata
                                      , lig_dist_colname = 'ligand_distance'
                                      , lig_dist_cutoff = 10)

merged_df2        = all_plot_dfs[[1]]
merged_df3        = all_plot_dfs[[2]]
#merged_df2_comp   = all_plot_dfs[[3]]
#merged_df3_comp   = all_plot_dfs[[4]]
#merged_df2_lig    = all_plot_dfs[[5]]
#merged_df3_lig    = all_plot_dfs[[6]]

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

###########################
# 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
#%%%%%%%%%%%%%%%%%%%%%%%%%%

# 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/
#############
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)
logo_data = my_data_snp[, c("position"
                            , "mutant_type", "or_mychisq", "log10or")]

logo_data_or = my_data_snp[, c("position", "mutant_type", "or_mychisq")]
wide_df_or <- logo_data_or %>% spread(position, or_mychisq, fill = 0.0)

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

position_or = as.numeric(colnames(wide_df_or))

#===========================================
#custom height (OR) logo plot: CORRECT x-axis labelling
#============================================  
# custom height (OR) logo plot: yayy works
cat("Logo plot with OR as y axis:", plot_logo_plot)
svg(plot_logo_plot, 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 = 12
                                       , angle = 90
                                       , hjust = 1
                                       , vjust = 0.4)
          , axis.text.y = element_text(size = 22
                                       , angle = 0
                                       , hjust = 1
                                       , vjust = 0)
          , axis.title.y = element_text(size = 25)
          , axis.title.x = element_text(size = 20)
          #, legend.position = "bottom") + 
          , legend.position = "none")+
          #, legend.text = element_text(size = 15)
          #, legend.title = element_text(size = 15))+
      scale_x_discrete("Position"
                     #, breaks
                     , labels = position_or
                     , limits = factor(1:length(position_or))) +
  ylab("Odds Ratio")
  
print(logo_or)
dev.off()
#%% end of logo plot with OR as height

#=======================================================================
# extracting data with log10R
logo_data_logor = my_data_snp[, c("position", "mutant_type", "log10or")]
wide_df_logor <- logo_data_logor %>% spread(position, log10or, fill = 0.0)

wide_df_logor = as.matrix(wide_df_logor)

rownames(wide_df_logor) = wide_df_logor[,1]
wide_df_logor = subset(wide_df_logor, select = -c(1) )
colnames(wide_df_logor)
wide_df_logor_m = data.matrix(wide_df_logor)

rownames(wide_df_logor_m)
colnames(wide_df_logor_m)

# FIXME
#my_ylim_up = as.numeric(max(wide_df_logor_m)) * 5
#my_ylim_low = as.numeric(min(wide_df_logor_m))

position_logor = as.numeric(colnames(wide_df_logor_m))

# custom height (log10OR) logo plot: yayy works
ggseqlogo(wide_df_logor_m, 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))+
  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))
#=======================================================================
#%% end of script
#=======================================================================