# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# + [markdown] code_folding=[]
# # Goal
#
# * Simulating fullCyc Day1 control gradients
#   * Not simulating incorporation (all 0% isotope incorp.)
#     * Don't know how much true incorporatation for emperical data
# * Using parameters inferred from emperical data (fullCyc Day1 seq data), or if not available, default SIPSim parameters
# * Determining whether simulated taxa show similar distribution to the emperical  data
# * Simulating higher levels of richness
# -

# ## Init

import os
import glob
import re
import nestly
# %load_ext rpy2.ipython
# %load_ext pushnote

# + language="R"
# library(ggplot2)
# library(dplyr)
# library(tidyr)
# library(gridExtra)
# library(phyloseq)
#
# ## BD for G+C of 0 or 100
# BD.GCp0 = 0 * 0.098 + 1.66
# BD.GCp100 = 1 * 0.098 + 1.66
# -

# # Nestly
#
# * assuming fragments already simulated

# +
workDir = '/home/nick/notebook/SIPSim/dev/fullCyc/n1147_frag_norm_9_2.5_n5/'
buildDir = os.path.join(workDir, 'Day1_xRich')
R_dir = '/home/nick/notebook/SIPSim/lib/R/'

fragFile= '/home/nick/notebook/SIPSim/dev/bac_genome1147/validation/ampFrags.pkl'
targetFile = '/home/nick/notebook/SIPSim/dev/fullCyc/CD-HIT/target_taxa.txt'

physeqDir = '/var/seq_data/fullCyc/MiSeq_16SrRNA/515f-806r/lib1-7/phyloseq/'
physeq_bulkCore = 'bulk-core'
physeq_SIP_core = 'SIP-core_unk'


prefrac_comm_abundance = ['1e9']
richness = [2503, 4000, 8000, 120000]   # 2503=  chao1 estimate for bulk Day 1

seq_per_fraction = ['lognormal', 9.432, 0.5, 10000, 30000] # dist, mean, scale, min, max
bulk_days = [1]
nprocs = 24

# +
# building tree structure
nest = nestly.Nest()

## varying params
nest.add('richness', richness)

## set params
nest.add('abs', prefrac_comm_abundance, create_dir=False)
nest.add('bulk_day', bulk_days, create_dir=False)
nest.add('percIncorp', [0], create_dir=False)
nest.add('percTaxa', [0], create_dir=False)
nest.add('np', [nprocs], create_dir=False)
nest.add('subsample_dist', [seq_per_fraction[0]], create_dir=False)
nest.add('subsample_mean', [seq_per_fraction[1]], create_dir=False)
nest.add('subsample_scale', [seq_per_fraction[2]], create_dir=False)
nest.add('subsample_min', [seq_per_fraction[3]], create_dir=False)
nest.add('subsample_max', [seq_per_fraction[4]], create_dir=False)

### input/output files
nest.add('buildDir', [buildDir], create_dir=False)
nest.add('R_dir', [R_dir], create_dir=False)
nest.add('fragFile', [fragFile], create_dir=False)
nest.add('targetFile', [targetFile], create_dir=False)
nest.add('physeqDir', [physeqDir], create_dir=False)
nest.add('physeq_bulkCore', [physeq_bulkCore], create_dir=False)


# building directory tree
nest.build(buildDir)

# bash file to run
bashFile = os.path.join(buildDir, 'SIPSimRun.sh')

# +
# %%writefile $bashFile
# #!/bin/bash

export PATH={R_dir}:$PATH

#-- making DNA pool similar to gradient of interest
# echo '# Creating comm file from phyloseq'
phyloseq2comm.r {physeqDir}{physeq_bulkCore} -s 12C-Con -d {bulk_day} > {physeq_bulkCore}_comm.txt
printf 'Number of lines: '; wc -l {physeq_bulkCore}_comm.txt

# echo '## Adding target taxa to comm file'
comm_add_target.r {physeq_bulkCore}_comm.txt {targetFile} > {physeq_bulkCore}_comm_target.txt
printf 'Number of lines: '; wc -l {physeq_bulkCore}_comm_target.txt

# echo '# Adding extra richness to community file'
printf "1\t{richness}\n" > richness_needed.txt
comm_add_richness.r -s {physeq_bulkCore}_comm_target.txt richness_needed.txt > {physeq_bulkCore}_comm_all.txt
### renaming comm file for downstream pipeline
# cat {physeq_bulkCore}_comm_all.txt > {physeq_bulkCore}_comm_target.txt
# rm -f {physeq_bulkCore}_comm_all.txt 

# echo '## parsing out genome fragments to make simulated DNA pool resembling the gradient of interest'
## all OTUs without an associated reference genome will be assigned a random reference (of the reference genome pool)
### this is done through --NA-random
SIPSim fragment_KDE_parse {fragFile} {physeq_bulkCore}_comm_target.txt \
    --rename taxon_name --NA-random > fragsParsed.pkl


# echo '#-- SIPSim pipeline --#'
# echo '# converting fragments to KDE'
SIPSim fragment_KDE \
    fragsParsed.pkl \
    > fragsParsed_KDE.pkl
    
# echo '# adding diffusion'    
SIPSim diffusion \
    fragsParsed_KDE.pkl \
    --np {np} \
    > fragsParsed_KDE_dif.pkl    

# echo '# adding DBL contamination'
SIPSim DBL \
    fragsParsed_KDE_dif.pkl \
    --np {np} \
    > fragsParsed_KDE_dif_DBL.pkl      
    
# echo '# making incorp file'
SIPSim incorpConfigExample \
  --percTaxa {percTaxa} \
  --percIncorpUnif {percIncorp} \
  > {percTaxa}_{percIncorp}.config

# echo '# adding isotope incorporation to BD distribution'
SIPSim isotope_incorp \
    fragsParsed_KDE_dif_DBL.pkl \
    {percTaxa}_{percIncorp}.config \
    --comm {physeq_bulkCore}_comm_target.txt \
    --np {np} \
    > fragsParsed_KDE_dif_DBL_inc.pkl
    
# #echo '# calculating BD shift from isotope incorporation'
#SIPSim BD_shift \
#    fragsParsed_KDE_dif_DBL.pkl \
#    fragsParsed_KDE_dif_DBL_inc.pkl \
#    --np {np} \
#    > fragsParsed_KDE_dif_DBL_inc_BD-shift.txt

# echo '# simulating gradient fractions'
SIPSim gradient_fractions \
    {physeq_bulkCore}_comm_target.txt \
    > fracs.txt 

# echo '# simulating an OTU table'
SIPSim OTU_table \
    fragsParsed_KDE_dif_DBL_inc.pkl \
    {physeq_bulkCore}_comm_target.txt \
    fracs.txt \
    --abs {abs} \
    --np {np} \
    > OTU_abs{abs}.txt
    
# #echo '# simulating PCR'
SIPSim OTU_PCR \
    OTU_abs{abs}.txt \
    > OTU_abs{abs}_PCR.txt    
    
# echo '# subsampling from the OTU table (simulating sequencing of the DNA pool)'
SIPSim OTU_subsample \
    --dist {subsample_dist} \
    --dist_params mean:{subsample_mean},sigma:{subsample_scale} \
    --min_size {subsample_min} \
    --max_size {subsample_max} \
    OTU_abs{abs}_PCR.txt \
    > OTU_abs{abs}_PCR_sub.txt
        
# echo '# making a wide-formatted table'
SIPSim OTU_wideLong -w \
    OTU_abs{abs}_PCR_sub.txt \
    > OTU_abs{abs}_PCR_sub_w.txt
    
# echo '# making metadata (phyloseq: sample_data)'
SIPSim OTU_sampleData \
    OTU_abs{abs}_PCR_sub.txt \
    > OTU_abs{abs}_PCR_sub_meta.txt
# -

# !chmod 777 $bashFile
# !cd $workDir; \
#     nestrun  --template-file $bashFile -d Day1_xRich --log-file log.txt -j 1

# #### Notes
#
# * richness of 8000 & 12000 failed due to memory errors

# # BD min/max
#
# * what is the min/max BD that we care about?

# + language="R"
# ## min G+C cutoff
# min_GC = 13.5
# ## max G+C cutoff
# max_GC = 80
# ## max G+C shift
# max_13C_shift_in_BD = 0.036
#
#
# min_BD = min_GC/100.0 * 0.098 + 1.66    
# max_BD = max_GC/100.0 * 0.098 + 1.66    
#
# max_BD = max_BD + max_13C_shift_in_BD
#
# cat('Min BD:', min_BD, '\n')
# cat('Max BD:', max_BD, '\n')
# -

# # Loading simulated OTU tables

# + magic_args="-i OTU_files -i buildDir" language="R"
# # loading files
#
# OTU_files = c('/home/nick/notebook/SIPSim/dev/fullCyc/n1147_frag_norm_9_2.5_n5/Day1_xRich/2503/OTU_abs1e9_PCR_sub.txt',
#               '/home/nick/notebook/SIPSim/dev/fullCyc/n1147_frag_norm_9_2.5_n5/Day1_xRich/4000/OTU_abs1e9_PCR_sub.txt')
#
# df.SIM = list()
# for (x in OTU_files){
#     richness = gsub(paste0(buildDir, '/'), '', x)
#     richness = gsub('/OTU_abs1e9_PCR_sub.txt', '', richness)
#     df.SIM[[richness]] = read.delim(x, sep='\t') 
#     }
# df.SIM = do.call('rbind', df.SIM)
# df.SIM$richness = gsub('\\.[0-9]+$', '', rownames(df.SIM))
# rownames(df.SIM) = 1:nrow(df.SIM)
# df.SIM %>% head

# + language="R"
# ## edit table
# df.SIM.nt = df.SIM %>%
#     filter(count > 0) %>%
#     group_by(richness, library, BD_mid) %>%
#     summarize(n_taxa = n()) %>%
#     filter(BD_mid >= min_BD, 
#            BD_mid <= max_BD)
# df.SIM.nt %>% head    
# -

# ## Plotting number of taxa in each fraction

# ### Emperical data (fullCyc)

# + language="R"
# # simulated OTU table file
# OTU.table.dir = '/home/nick/notebook/SIPSim/dev/fullCyc/frag_norm_9_2.5_n5/Day1_default_run/1e9/'
# OTU.table.file = 'OTU_abs1e9_PCR_sub.txt'
# #OTU.table.file = 'OTU_abs1e9_sub.txt'
# #OTU.table.file = 'OTU_abs1e9.txt'

# + magic_args="-i physeqDir -i physeq_SIP_core -i bulk_days" language="R"
#
# # bulk core samples
# F = file.path(physeqDir, physeq_SIP_core)
# physeq.SIP.core = readRDS(F) 
# physeq.SIP.core.m = physeq.SIP.core %>% sample_data
#
# physeq.SIP.core = prune_samples(physeq.SIP.core.m$Substrate == '12C-Con' & 
#                                 physeq.SIP.core.m$Day %in% bulk_days, 
#                                 physeq.SIP.core) %>%
#     filter_taxa(function(x) sum(x) > 0, TRUE)
# physeq.SIP.core.m = physeq.SIP.core %>% sample_data        
#
# physeq.SIP.core

# + magic_args="-w 800 -h 300 " language="R"
#
# ## dataframe
# df.EMP = physeq.SIP.core %>% otu_table %>%
#     as.matrix %>% as.data.frame
# df.EMP$OTU = rownames(df.EMP)
# df.EMP = df.EMP %>%    
#     gather(sample, abundance, 1:(ncol(df.EMP)-1)) 
#
# df.EMP = inner_join(df.EMP, physeq.SIP.core.m, c('sample' = 'X.Sample')) 
#
# df.EMP.nt = df.EMP %>%
#     group_by(sample) %>%
#     mutate(n_taxa = sum(abundance > 0)) %>%
#     ungroup() %>%
#     distinct(sample) %>%
#     filter(Buoyant_density >= min_BD, 
#            Buoyant_density <= max_BD) 
#     
#
# ## plotting
# p = ggplot(df.EMP.nt, aes(Buoyant_density, n_taxa)) +
#     geom_point(color='blue') +
#     geom_line(color='blue') +
#     #geom_vline(xintercept=c(BD.GCp0, BD.GCp100), linetype='dashed', alpha=0.5) +
#     labs(x='Buoyant density', y='Number of taxa') +
#     theme_bw() +
#     theme( 
#         text = element_text(size=16),
#         legend.position = 'none'
#     )
# p
# -

# ### w/ simulated data

# + magic_args="-w 800 -h 300" language="R"
# # plotting
# p = ggplot(df.SIM.nt, aes(BD_mid, n_taxa)) +
#     geom_point(aes(color=richness, group=richness)) +
#     geom_line(aes(color=richness, group=richness), alpha=0.5) +
#     #geom_smooth(color='red') +
#     geom_point(data=df.EMP.nt, aes(x=Buoyant_density), color='blue') +
#     geom_line(data=df.EMP.nt, aes(x=Buoyant_density), color='blue') +
#     #geom_vline(xintercept=c(BD.GCp0, BD.GCp100), linetype='dashed', alpha=0.5) +
#     labs(x='Buoyant density', y='Number of taxa') +
#     theme_bw() +
#     theme( 
#         text = element_text(size=16)#,
#     #    legend.position = 'none'
#     )
# p
# -

# ## Total sequence count

# + magic_args="-w 800 -h 300" language="R"
#
# # simulated
# df.SIM.s = df.SIM %>%
#     group_by(richness, library, BD_mid) %>%
#     summarize(total_abund = sum(count)) %>%
#     rename('Day' = library, 'Buoyant_density' = BD_mid) %>%
#     ungroup() %>%
#     mutate(dataset='simulated')
# # emperical
# df.EMP.s = df.EMP %>% 
#     group_by(Day, Buoyant_density) %>%
#     summarize(total_abund = sum(abundance))  %>%
#     ungroup() %>%
#     mutate(dataset='emperical', richness = NA)
#
# # join
# df.j = rbind(df.SIM.s, df.EMP.s) %>%
#     filter(Buoyant_density >= min_BD, 
#            Buoyant_density <= max_BD)
#     
# df.SIM.s = df.EMP.s = ""
#
# # plot
# ggplot(df.j, aes(Buoyant_density, total_abund, color=dataset, group=richness)) +
#     geom_point() +
#     geom_line(alpha=0.3) +
#     geom_line(data=df.j %>% filter(dataset=='emperical')) +
#     scale_color_manual(values=c('blue', 'red')) +
#     labs(x='Buoyant density', y='Total sequences per sample') +
#     theme_bw() +
#     theme( 
#         text = element_text(size=16),
#         legend.position = 'none'
#     )
# -

# ## Plotting Shannon diversity for each

# + language="R"
# shannon_index_long = function(df, abundance_col, ...){
#     # calculating shannon diversity index from a 'long' formated table
#     ## community_col = name of column defining communities
#     ## abundance_col = name of column defining taxon abundances
#     df = df %>% as.data.frame
#     cmd = paste0(abundance_col, '/sum(', abundance_col, ')')
#     df.s = df %>%
#         group_by_(...) %>%
#         mutate_(REL_abundance = cmd) %>%
#         mutate(pi__ln_pi = REL_abundance * log(REL_abundance),
#                shannon = -sum(pi__ln_pi, na.rm=TRUE)) %>%
#         ungroup() %>% 
#         dplyr::select(-REL_abundance, -pi__ln_pi) %>%
#         distinct_(...) 
#     return(df.s)
# }

# + language="R"
# # calculating shannon
# df.SIM.shan = shannon_index_long(df.SIM, 'count', 'richness', 'library', 'fraction') %>%
#     filter(BD_mid >= min_BD, 
#            BD_mid <= max_BD)
#     
# df.EMP.shan = shannon_index_long(df.EMP, 'abundance', 'sample') %>%
#     filter(Buoyant_density >= min_BD, 
#            Buoyant_density <= max_BD)
#     

# + magic_args="-w 800 -h 300" language="R"
# # plotting
# p = ggplot(df.SIM.shan, aes(BD_mid, shannon, group=richness)) +
#     geom_point(aes(color=richness)) +
#     geom_line(aes(color=richness), alpha=0.3) +
#     geom_point(data=df.EMP.shan, aes(x=Buoyant_density), color='blue') +
#     geom_line(data=df.EMP.shan, aes(x=Buoyant_density), color='blue') +
#     scale_y_continuous(limits=c(4, 7.5)) +
#     labs(x='Buoyant density', y='Shannon index') +
#     theme_bw() +
#     theme( 
#         text = element_text(size=16)#,
#         #legend.position = 'none'
#     )
# p
#
# -

# # min/max abundances of taxa

# + code_folding=[] magic_args="-h 300 -w 800" language="R"
#
# # simulated
# df.SIM.s = df.SIM %>% 
#     filter(rel_abund > 0) %>%
#     group_by(richness, BD_mid) %>%
#     summarize(min_abund = min(rel_abund),
#               max_abund = max(rel_abund)) %>%
#     ungroup() %>%
#     rename('Buoyant_density' = BD_mid) %>%
#     mutate(dataset = 'simulated')
#
# # emperical
# df.EMP.s = df.EMP %>%
#     group_by(Buoyant_density) %>%
#     mutate(rel_abund = abundance / sum(abundance)) %>%
#     filter(rel_abund > 0) %>%
#     summarize(min_abund = min(rel_abund),
#               max_abund = max(rel_abund)) %>%
#     ungroup() %>%
#     mutate(dataset = 'emperical',
#            richness = NA)
#
# df.j = rbind(df.SIM.s, df.EMP.s) %>%
#     filter(Buoyant_density >= min_BD, 
#            Buoyant_density <= max_BD)
#     
#
# # plotting
# ggplot(df.j, aes(Buoyant_density, max_abund, color=dataset, group=richness)) +
#     geom_point(aes(color=richness)) +
#     geom_line(aes(color=richness), alpha=0.3) +
#     geom_line(data=df.j %>% filter(dataset=='emperical')) +
#     scale_color_manual(values=c('green', 'red', 'blue')) +
#     labs(x='Buoyant density', y='Maximum relative abundance') +
#     theme_bw() +
#     theme( 
#         text = element_text(size=16),
#         legend.position = 'none'
#     )
# -

# # BD range where an OTU is detected 
#
# * Do the simulated OTU BD distributions span the same BD range of the emperical data?

# ### Simulated

# + language="R"
# comm_files = c('/home/nick/notebook/SIPSim/dev/fullCyc/n1147_frag_norm_9_2.5_n5/Day1_xRich/2503/bulk-core_comm_target.txt',
#               '/home/nick/notebook/SIPSim/dev/fullCyc/n1147_frag_norm_9_2.5_n5/Day1_xRich/4000/bulk-core_comm_target.txt')

# + language="R"
#
# df.comm = list()
# for (f in comm_files){
#     rep = gsub('.+/Day1_xRich/([0-9]+)/.+', '\\1', f)
#     df.comm[[rep]] = read.delim(f, sep='\t') %>%
#         dplyr::select(library, taxon_name, rel_abund_perc) %>%
#         rename('bulk_abund' = rel_abund_perc) %>%
#         mutate(bulk_abund = bulk_abund / 100)
# }
#
# df.comm = do.call('rbind', df.comm)
# df.comm$richness = gsub('\\.[0-9]+$', '', rownames(df.comm))
# rownames(df.comm) = 1:nrow(df.comm)
# df.comm %>% head(n=3)

# + language="R"
#
# ## joining
# df.SIM.j = inner_join(df.SIM, df.comm, c('richness' = 'richness',
#                                          'library' = 'library',
#                                          'taxon' = 'taxon_name')) %>%
#     filter(BD_mid >= min_BD, 
#            BD_mid <= max_BD)
#     
# df.SIM.j %>% head(n=3)
# -

# ### Emperical

# + language="R"
# bulk_days = c(1)

# + language="R"
# physeq.dir = '/var/seq_data/fullCyc/MiSeq_16SrRNA/515f-806r/lib1-7/phyloseq/'
# physeq.bulk = 'bulk-core'
# physeq.file = file.path(physeq.dir, physeq.bulk)
# physeq.bulk = readRDS(physeq.file)
# physeq.bulk.m = physeq.bulk %>% sample_data
# physeq.bulk = prune_samples(physeq.bulk.m$Exp_type == 'microcosm_bulk' &
#                             physeq.bulk.m$Day %in% bulk_days, physeq.bulk)
#
# physeq.bulk.m = physeq.bulk %>% sample_data
# physeq.bulk

# + language="R"
# physeq.bulk.n = transform_sample_counts(physeq.bulk, function(x) x/sum(x))
# physeq.bulk.n

# + code_folding=[] language="R"
# # making long format of each bulk table
# bulk.otu = physeq.bulk.n %>% otu_table %>% as.data.frame
# ncol = ncol(bulk.otu)
# bulk.otu$OTU = rownames(bulk.otu)
# bulk.otu = bulk.otu %>%
#     gather(sample, abundance, 1:ncol) 
#
# bulk.otu = inner_join(physeq.bulk.m, bulk.otu, c('X.Sample' = 'sample')) %>%
#     dplyr::select(OTU, abundance) %>%
#     rename('bulk_abund' = abundance)
# bulk.otu %>% head(n=3)

# + language="R"
# # joining tables
# df.EMP.j = inner_join(df.EMP, bulk.otu, c('OTU' = 'OTU')) %>%
#     filter(Buoyant_density >= min_BD, 
#            Buoyant_density <= max_BD)
#     
# df.EMP.j %>% head(n=3)

# + language="R"
# # filtering & combining emperical w/ simulated data
#
# ## emperical 
# max_BD_range = max(df.EMP.j$Buoyant_density) - min(df.EMP.j$Buoyant_density)
# df.EMP.j.f = df.EMP.j %>%
#     filter(abundance > 0) %>%
#     group_by(OTU) %>%
#     summarize(mean_rel_abund = mean(bulk_abund),
#               min_BD = min(Buoyant_density),
#               max_BD = max(Buoyant_density),
#               BD_range = max_BD - min_BD,
#               BD_range_perc = BD_range / max_BD_range * 100) %>%
#     ungroup() %>%
#     mutate(dataset = 'emperical',
#            richness = NA)
#
# ## simulated
# max_BD_range = max(df.SIM.j$BD_mid) - min(df.SIM.j$BD_mid)
# df.SIM.j.f = df.SIM.j %>%
#     filter(count > 0) %>%
#     group_by(richness, taxon) %>%
#     summarize(mean_rel_abund = mean(bulk_abund),
#               min_BD = min(BD_mid),
#               max_BD = max(BD_mid),
#               BD_range = max_BD - min_BD,
#               BD_range_perc = BD_range / max_BD_range * 100) %>%
#     ungroup() %>%
#     rename('OTU' = taxon) %>%
#     mutate(dataset = 'simulated')
#
# ## join
# df.j = rbind(df.EMP.j.f, df.SIM.j.f) %>%
#     filter(BD_range_perc > 0,
#             mean_rel_abund > 0)
#
# df.j %>% head(n=3)

# + magic_args="-h 400" language="R"
# ## plotting
# ggplot(df.j, aes(mean_rel_abund, BD_range_perc, color=richness)) +
#     geom_point(alpha=0.5, shape='O') +
#     scale_x_log10() +
#     scale_y_continuous() +
#     labs(x='Pre-fractionation abundance', y='% of total BD range') +
#     facet_grid(dataset ~ .) +
#     theme_bw() +
#     theme(
#         text = element_text(size=16),
#         panel.grid = element_blank()#,
#         #legend.position = 'none'
#         )
#
# -

# #### Notes:
#
# * Higher richness may have actually lowered BD span
#   * Would need replicates to tell for sure



StarcoderdataJupyterScriptsDedupFiltered