#!/usr/bin/env Rscript 

###########################################################
# TASK: To combine mcsm combined file and meta data.
# This script is sourced by other .R scripts for plotting.
###########################################################
# load libraries and functions

#source("Header_TT.R")

#==========================================================
# combining_dfs_plotting(): 

# input args

## df1_mcsm_comb: <gene>_meta_data.csv
## df2_gene_metadata: <gene>_all_params.csv
## lig_dist_cutoff = 10, cut off distance

# Output: returns
# 1) large combined df including NAs for AF, OR,etc
# 		Dim: same no. of rows as gene associated meta_data_with_AFandOR
# 2) small combined df including NAs for AF, OR, etc.
# 		Dim: same as mcsm data
# 3) large combined df excluding NAs 
# 		Dim: dim(#1) - na_count_df2
# 4) small combined df excluding NAs
# 		Dim: dim(#2) - na_count_df3
# 5) LIGAND large combined df including NAs for AF, OR,etc
# 		Dim: dim()
# 6) LIGAND small combined df excluding NAs
# 		Dim: dim()
#==========================================================
combining_dfs_plotting <- function(  my_df_u
                                   , gene_metadata
                                   , lig_dist_colname = 'ligand_distance'
                                   , lig_dist_cutoff = 10){

  # 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("\nMerging 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("\nDim 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("\nChecking nrows in merged_df2")
  if(nrow(gene_metadata) == nrow(merged_df2)){
    cat("\nPASS: nrow(merged_df2) = nrow (gene associated gene_metadata)"
        ,"\nExpected no. of rows: ",nrow(gene_metadata) 
        ,"\nGot no. of rows: ", nrow(merged_df2))
  } else{
    cat("\nFAIL: 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()
  }

  # Quick formatting: pretty labels
  #-----------------------
  # mutation_info_labels
  #-----------------------
  merged_df2$mutation_info_labels = ifelse(merged_df2$mutation_info == dr_muts_col
                                              , "DM", "OM")
  merged_df2$mutation_info_labels = factor(merged_df2$mutation_info_labels)
  #-----------------------
  # lineage labels
  #-----------------------
  merged_df2$lineage_labels =  gsub("lineage", "L", merged_df2$lineage)
  
  merged_df2$lineage_labels =  factor(merged_df2$lineage_labels, c("L1"
                                                                   , "L2"
                                                                   , "L3"
                                                                   , "L4"
                                                                   , "L5"
                                                                   , "L6"
                                                                   , "L7"
                                                                   , "LBOV"
                                                                   , "L1;L2"
                                                                   , "L1;L3"
                                                                   , "L1;L4"
                                                                   , "L2;L3"
                                                                   , "L2;L3;L4"
                                                                   , "L2;L4"  
                                                                   , "L2;L6"
                                                                   , "L2;LBOV"  
                                                                   , "L3;L4" 
                                                                   , "L4;L6"
                                                                   , "L4;L7"
                                                                   , ""))
  
  
  #=================================================================
  # 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("\nMerging 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("\nChecking nrows in merged_df3")
  if(nrow(my_df_u) == nrow(merged_df3)){
    cat("\nPASS: No. of rows match with my_df"
        ,"\nExpected no. of rows: ", nrow(my_df_u)
        ,"\nGot no. of rows: ", nrow(merged_df3))
  } else {
    cat("\nFAIL: 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("\nPASS: NA count match for OR, pvalue and AF\n")
    na_count_df3 = sum(is.na(merged_df3$af_kin))
    cat("\nNo. of NAs: ", sum(is.na(merged_df3$or_kin)))
  } else{
    cat("\nFAIL: 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("\nMerging 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("\nChecking 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("\nFAIL: 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("\nChecking 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("\nFAIL: 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))
  cat("\n------------------------"
      , "\nSummary of created dfs:"
      , "\n------------------------"
      , "\n1) Dim of merged_df2: " , nrow(merged_df2), "," , ncol(merged_df2)
      , "\n2) Dim of merged_df2_comp: " , nrow(merged_df2_comp), "," , ncol(merged_df2_comp)
      , "\n3) Dim of merged_df3: " , nrow(merged_df3), "," , ncol(merged_df3)
      , "\n4) Dim of merged_df3_comp: " , nrow(merged_df3_comp), "," , ncol(merged_df3_comp))
  
  #####################################################################
  #                       Combining: LIG
  #####################################################################
  
  #============
  # Merges 5-8
  #============
  cat("\n=========================================="
      , "\nStarting filtering for mcsm ligand df"
      , "\n===========================================")
  
  if (lig_dist_colname%in%names(my_df_u)){
    cat("\nFiltering column: ", lig_dist_colname
        , "\nCut off criteria: ", lig_dist_cutoff, "Angstroms")
    df_lig = my_df_u[my_df_u[[lig_dist_colname]] < lig_dist_cutoff,]
  
    #merged_df2_lig = merged_df2[merged_df2$ligand_distance<lig_dist_cutoff,]
    merged_df2_lig = merged_df2[merged_df2[[lig_dist_colname]] < lig_dist_cutoff,]
    dim(merged_df2_lig)
    
    merged_df2_comp_lig = merged_df2_comp[merged_df2_comp[[lig_dist_colname]] < lig_dist_cutoff,]
    
    merged_df3_lig = merged_df3[merged_df3[[lig_dist_colname]] < lig_dist_cutoff,]
    merged_df3_comp_lig = merged_df3_comp[merged_df3_comp[[lig_dist_colname]] < lig_dist_cutoff,]
    
    cat("\n------------------------"
        , "\nSummary of created ligand dfs:"
        , "\n------------------------"
        , "\n1) Dim of merged_df2_lig: " , nrow(merged_df2_lig), "," , ncol(merged_df2_lig)
        , "\n2) Dim of merged_df2_comp_lig: " , nrow(merged_df2_comp_lig), "," , ncol(merged_df2_comp_lig)
        , "\n3) Dim of merged_df3_lig: " , nrow(merged_df3_lig), "," , ncol(merged_df3_lig)
        , "\n4) Dim of merged_df3_comp_lig: " , nrow(merged_df3_comp_lig), "," , ncol(merged_df3_comp_lig))
} else {
    cat("\nFiltering column: ", lig_dist_colname, " not found\n")
  }
  #quit()
  
  # sanity check
  if (nrow(merged_df3_lig) == nrow(df_lig)){
    print("\nPASS: verified merged_df3_lig")
  }else{
    cat(paste0("\nFAIL: nrow mismatch for merged_df3_lig"
               , "\nExpected:", nrow(df_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")
  #}
  
  return(list(  merged_df2
              , merged_df3
              , merged_df2_comp
              , merged_df3_comp
              , merged_df2_lig
              , merged_df3_lig
              , merged_df2_comp_lig
              , merged_df3_comp_lig))
  
}