#!/usr/bin/env Rscript
#########################################################
# TASK: To combine struct params and meta data for plotting
# Input csv files:
# 1) <gene>_all_params.csv
# 2) <gene>_meta_data.csv

# Output: 
# 1) muts with opposite effects on stability
# 2) large combined df including NAs for AF, OR,etc
# 		Dim: same no. of rows as gene associated meta_data_with_AFandOR
# 3) small combined df including NAs for AF, OR, etc.
# 		Dim: same as mcsm data
# 4) large combined df excluding NAs 
# 		Dim: dim(#1) - na_count_df2
# 5) small combined df excluding NAs
# 		Dim: dim(#2) - na_count_df3

# This script is sourced from other .R scripts for plotting
#########################################################
#=======================================================================
# working dir and loading libraries
getwd()
setwd("~/git/LSHTM_analysis/scripts/plotting/")
getwd()

require("getopt", quietly = TRUE) # cmd parse arguments

# load functions
source("~/git/LSHTM_analysis/scripts/Header_TT.R")
source("../functions/plotting_globals.R")
source("../functions/plotting_data.R")

#############################################################
# command line args
#********************
# !!!FUTURE TODO!!!
# Can pass additional params of output/plot dir by user. 
# Not strictly required for my workflow  since it is optimised
# to have a streamlined input/output flow without filename worries.
#********************
spec = matrix(c(
  "drug"   ,"d", 1, "character",
  "gene"   ,"g", 1, "character",
  "data"   ,"f", 2, "character" 
), byrow = TRUE, ncol = 4)

opt = getopt(spec)

#FIXME: detect if script running from cmd, then set these
drug   = opt$drug
gene   = opt$gene
infile = opt$data

# hardcoding when not using cmd
#drug = "streptomycin"
#gene = "gid"

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

#########################################################
# call functions with relevant args
#***********************************
# import_dirs(): returns
  # datadir
  # indir
  # outdir
  # plotdir
  # dr_muts_col
  # other_muts_col
  # resistance_col
#***********************************
import_dirs(drug, gene)
#***********************************
# plotting_data(): returns
  # my_df
  # my_df_u
  # my_df_u_lig
  # dup_muts
#***********************************
#infile = "/home/tanu/git/Data/streptomycin/output/gid_comb_stab_struc_params.csv"

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

# Get the DFs out of plotting_data()
pd_df = plotting_data(infile)
my_df       = pd_df[[1]]
my_df_u     = pd_df[[2]]
my_df_u_lig = pd_df[[3]]
dup_muts    = pd_df[[4]]

cat(paste0("Directories imported:"
           , "\ndatadir:" , datadir
           , "\nindir:"   , indir
           , "\noutdir:"  , outdir
           , "\nplotdir:" , plotdir))

cat(paste0("\nVariables imported:"
           , "\ndrug:"       , drug
           , "\ngene:"       , gene
           , "\ngene match:" , gene_match
           , "\n"))
#========================================================
#===========
# input
#===========
#in_file1: output of plotting_data.R
# my_df_u

# infile 2: gene associated meta data
#in_filename_gene_metadata = paste0(tolower(gene),  "_meta_data_with_AFandOR.csv")
in_filename_gene_metadata = paste0(tolower(gene),  "_metadata.csv")
infile_gene_metadata = paste0(outdir, "/", in_filename_gene_metadata)
cat(paste0("Input infile 2:", infile_gene_metadata))

#===========
# output
#===========
# other variables that you can write
# primarily called by other scripts for plotting

# PS combined: 
# 1) merged_df2
# 2) merged_df2_comp
# 3) merged_df3
# 4) merged_df3_comp

# LIG combined: 
# 5) merged_df2_lig
# 6) merged_df2_comp_lig
# 7) merged_df3_lig
# 8) merged_df3_comp_lig

#%%===============================================================

###########################
# 2: Read file: <gene>_meta data.csv
###########################
cat("Reading meta data file:", infile_gene_metadata)

gene_metadata <- read.csv(infile_gene_metadata
                      , stringsAsFactors = F
                      , header = T)
cat("Dim:", dim(gene_metadata))

table(gene_metadata$mutation_info)

# counting NAs in AF, OR cols
# or_mychisq
if (identical(sum(is.na(my_df_u$or_mychisq))
              , sum(is.na(my_df_u$pval_fisher))
              , sum(is.na(my_df_u$af)))){
  cat("\nPASS: NA count match for OR, pvalue and AF\n")
  na_count = sum(is.na(my_df_u$af))
  cat("\nNo. of NAs: ", sum(is.na(my_df_u$or_mychisq)))
} else{
  cat("\nFAIL: NA count mismatch"
      , "\nNA in OR: ", sum(is.na(my_df_u$or_mychisq))
      , "\nNA in pvalue: ", sum(is.na(my_df_u$pval_fisher))
      , "\nNA in AF:", sum(is.na(my_df_u$af)))
}

# or kin
if (identical(sum(is.na(my_df_u$or_kin))
              , sum(is.na(my_df_u$pwald_kin))
              , sum(is.na(my_df_u$af_kin)))){
  cat("\nPASS: NA count match for OR, pvalue and AF\n from Kinship matrix calculations")
  na_count = sum(is.na(my_df_u$af_kin))
  cat("\nNo. of NAs: ", sum(is.na(my_df_u$or_kin)))
} else{
  cat("\nFAIL: NA count mismatch"
      , "\nNA in OR: ", sum(is.na(my_df_u$or_kin))
      , "\nNA in pvalue: ", sum(is.na(my_df_u$pwald_kin))
      , "\nNA in AF:", sum(is.na(my_df_u$af_kin)))
}

str(gene_metadata)

###################################################################
#                           combining: PS
###################################################################
# sort by position (same as my_df)
head(gene_metadata$position)
gene_metadata = gene_metadata[order(gene_metadata$position),]
head(gene_metadata$position)

#=========================
# Merge 1: merged_df2
# dfs with NAs in ORs
#=========================
head(my_df_u$mutationinformation)
head(gene_metadata$mutationinformation)

# Find common columns b/w two df
merging_cols = intersect(colnames(my_df_u), colnames(gene_metadata))

cat(paste0("Merging dfs with NAs: big df (1-many relationship b/w id & mut)"
           , "\nNo. of merging cols:", length(merging_cols)
           , "\nMerging columns identified:"))

print(merging_cols)

# using all common cols create confusion, so pick one!
# merging_cols = merging_cols[[1]]
merging_cols = 'mutationinformation'

cat("\nLinking column being used: mutationinformation")

# important checks!
table(nchar(my_df_u$mutationinformation))
table(nchar(my_df_u$wild_type))
table(nchar(my_df_u$mutant_type))
table(nchar(my_df_u$position))

# all.y  because x might contain non-structural positions!
merged_df2 = merge(x = gene_metadata
                  , y = my_df_u
                  , by = merging_cols
                  , all.y = T)

cat("Dim of merged_df2: ", dim(merged_df2))

# Remove duplicate columns 
dup_cols =  names(merged_df2)[grepl("\\.x$|\\.y$", names(merged_df2))]
cat("\nNo. of duplicate cols:", length(dup_cols))
check_df_cols = merged_df2[dup_cols]

identical(check_df_cols$wild_type.x, check_df_cols$wild_type.y)
identical(check_df_cols$position.x, check_df_cols$position.y)
identical(check_df_cols$mutant_type.x, check_df_cols$mutant_type.y)
# False: because some of the ones with OR don't have mutation
identical(check_df_cols$mutation.x, check_df_cols$mutation.y)

cols_to_drop = names(merged_df2)[grepl("\\.y",names(merged_df2))]
cat("\nNo. of cols to drop:", length(cols_to_drop))

# Drop duplicate columns
merged_df2 = merged_df2[,!(names(merged_df2)%in%cols_to_drop)]

# Drop the '.x' suffix in the colnames
names(merged_df2)[grepl("\\.x$|\\.y$", names(merged_df2))]
colnames(merged_df2) <- gsub("\\.x$", "", colnames(merged_df2))
names(merged_df2)[grepl("\\.x$|\\.y$", names(merged_df2))]

head(merged_df2$position)

# sanity check 
cat("Checking nrows in merged_df2")
if(nrow(gene_metadata) == nrow(merged_df2)){
  cat("PASS: nrow(merged_df2) = nrow (gene associated gene_metadata)"
      ,"\nExpected no. of rows: ",nrow(gene_metadata) 
      ,"\nGot no. of rows: ", nrow(merged_df2))
} else{
  cat("FAIL: nrow(merged_df2)!= nrow(gene associated gene_metadata)"
      , "\nExpected no. of rows after merge: ", nrow(gene_metadata)
      , "\nGot no. of rows: ", nrow(merged_df2)
      , "\nFinding discrepancy")
  merged_muts_u = unique(merged_df2$mutationinformation)
  meta_muts_u = unique(gene_metadata$mutationinformation)
  # find the index where it differs
  unique(meta_muts_u[! meta_muts_u %in% merged_muts_u])
  quit()
}

#=========================
# Merge 2: merged_df3
# dfs with NAs in ORs
#
# Cannot trust lineage, country from this df as the same mutation
# can have many different lineages
# but this should be good for the numerical corr plots
#=========================
# remove duplicated mutations
cat("Merging dfs without NAs: small df (removing muts with no AF|OR associated)"
    ,"\nCannot trust lineage info from this"
    ,"\nlinking col: mutationinforamtion"
    ,"\nfilename: merged_df3")

merged_df3 = merged_df2[!duplicated(merged_df2$mutationinformation),] 
head(merged_df3$position); tail(merged_df3$position) # should be sorted

# sanity check
cat("Checking nrows in merged_df3")
if(nrow(my_df_u) == nrow(merged_df3)){
  cat("PASS: No. of rows match with my_df"
      ,"\nExpected no. of rows: ", nrow(my_df_u)
      ,"\nGot no. of rows: ", nrow(merged_df3))
} else {
  cat("FAIL: No. of rows mismatch"
      , "\nNo. of rows my_df: ", nrow(my_df_u)
      , "\nNo. of rows merged_df3: ", nrow(merged_df3))
  quit()
}

# counting NAs in AF, OR cols in merged_df3
# this is because mcsm has no AF, OR cols,
# so you cannot count NAs
if (identical(sum(is.na(merged_df3$or_kin))
              , sum(is.na(merged_df3$pwald_kin))
              , sum(is.na(merged_df3$af_kin)))){
  cat("PASS: NA count match for OR, pvalue and AF\n")
  na_count_df3 = sum(is.na(merged_df3$af_kin))
  cat("No. of NAs: ", sum(is.na(merged_df3$or_kin)))
} else{
  cat("FAIL: NA count mismatch"
      , "\nNA in OR: ", sum(is.na(merged_df3$or_kin))
      , "\nNA in pvalue: ", sum(is.na(merged_df3$pwald_kin))
      , "\nNA in AF:", sum(is.na(merged_df3$af_kin)))
}

#=========================
# Merge3: merged_df2_comp
# same as merge 1 but excluding NAs from ORs, etc.
#=========================
cat("Merging dfs without any NAs: big df (1-many relationship b/w id & mut)"
    ,"\nfilename: merged_df2_comp")

na_count_df2 = sum(is.na(merged_df2$af))
merged_df2_comp = merged_df2[!is.na(merged_df2$af),] 

# sanity check: no +-1 gymnastics
cat("Checking nrows in merged_df2_comp")
if(nrow(merged_df2_comp) == (nrow(merged_df2) - na_count_df2)){
  cat("\nPASS: No. of rows match"
      ,"\nDim of merged_df2_comp: "
      ,"\nExpected no. of rows: ", nrow(merged_df2) - na_count_df2
      , "\nNo. of rows: ", nrow(merged_df2_comp)
      , "\nNo. of cols: ", ncol(merged_df2_comp))
}else{
  cat("FAIL: No. of rows mismatch"
      ,"\nExpected no. of rows: ", nrow(merged_df2) - na_count_df2
      ,"\nGot no. of rows: ", nrow(merged_df2_comp))
}

#=========================
# Merge4: merged_df3_comp
# same as merge 2 but excluding NAs from ORs, etc or 
# remove duplicate mutation information
#=========================
na_count_df3 = sum(is.na(merged_df3$af))
#merged_df3_comp = merged_df3_comp[!duplicated(merged_df3_comp$mutationinformation),] # a way

merged_df3_comp = merged_df3[!is.na(merged_df3$af),] # another way
cat("Checking nrows in merged_df3_comp")

if(nrow(merged_df3_comp) == (nrow(merged_df3) - na_count_df3)){
  cat("\nPASS: No. of rows match"
      ,"\nDim of merged_df3_comp: "
      ,"\nExpected no. of rows: ", nrow(merged_df3) - na_count_df3
      , "\nNo. of rows: ", nrow(merged_df3_comp)
      , "\nNo. of cols: ", ncol(merged_df3_comp))
}else{
  cat("FAIL: No. of rows mismatch"
      ,"\nExpected no. of rows: ", nrow(merged_df3) - na_count_df3
      ,"\nGot no. of rows: ", nrow(merged_df3_comp))
}

# alternate way of deriving merged_df3_comp
foo = merged_df3[!is.na(merged_df3$af),]
bar = merged_df3_comp[!duplicated(merged_df3_comp$mutationinformation),]
# compare dfs: foo and merged_df3_com
all.equal(foo, bar)
#summary(comparedf(foo, bar))

#==============================================================

#####################################################################
#                       Combining: LIG
#####################################################################

#=========================
# Merges 5-8
#=========================

merged_df2_lig = merged_df2[merged_df2$ligand_distance<10,]
merged_df2_comp_lig = merged_df2_comp[merged_df2_comp$ligand_distance<10,]

merged_df3_lig = merged_df3[merged_df3$ligand_distance<10,]
merged_df3_comp_lig = merged_df3_comp[merged_df3_comp$ligand_distance<10,]

# sanity check
if (nrow(merged_df3_lig) == nrow(my_df_u_lig)){
  print("PASS: verified merged_df3_lig")
}else{
  cat(paste0("FAIL: nrow mismatch for merged_df3_lig"
             , "\nExpected:", nrow(my_df_u_lig)
             , "\nGot:", nrow(merged_df3_lig)))
}

#==============================================================

#################
# OPTIONAL: write output files in one go
#################
#outvars = c(#"merged_df2",
            #"merged_df2_comp",
            #"merged_df2_lig",
            #"merged_df2_comp_lig",
            
            #"meregd_df3_comp"
            #"merged_df3_comp_lig",
            #"merged_df3",
            #"merged_df3_lig")

#cat("Writing output files: "
    #, "\nPath:", outdir)

#for (i in outvars){
  #out_filename = paste0(i, ".csv")
  #outfile = paste0(outdir, "/", out_filename)
  #cat("Writing output file:"
  #    ,"\nFilename: ", out_filename,"\n")
  #write.csv(get(i), outfile, row.names = FALSE)
  #cat("Finished writing: ", outfile
  #    , "\nNo. of rows: ", nrow(get(i))
  #    , "\nNo. of cols: ", ncol(get(i)), "\n")
#}

# clear variables
rm(foo, bar, gene_metadata
   , in_filename_params, infile_params, merging_cols
   , in_filename_gene_metadata, infile_gene_metadata)

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