workDir = '/home/nick/notebook/SIPSim/dev/bac_genome1210/'
buildDir = os.path.join(workDir, 'percTaxaIncorpUnif')
genomeDir = '/home/nick/notebook/SIPSim/dev/bac_genome1210/genomes/'
R_dir = '/home/nick/notebook/SIPSim/lib/R/'

"""
Explanation: Goal
Question: how does the percentage of taxa incorporating isotope affect the accuracy of DESeq2


Using genome dataset created in the "dataset" notebook


Simulates isotope dilution or short incubations

Method  
% of taxa w/ incorporation: 0, 5, 10, 25, 50, 75, 100
% isotope incorporated per taxon: 0, 20, 40, 60, 80, 100%
no variance on incorporation %



User variables
End of explanation
"""


import glob
import itertools
from os.path import abspath
import nestly

%load_ext rpy2.ipython

%%R
library(ggplot2)
library(dplyr)
library(tidyr)
library(gridExtra)

if not os.path.isdir(buildDir):
    os.makedirs(buildDir)

"""
Explanation: Init
End of explanation
"""


# 8 hour delay
!sleep 28800

"""
Explanation: Creating input files (eg., fragments & communities)
End of explanation
"""


!cd $buildDir; \
    SIPSim fragments \
    $genomeDir/genome_index.txt \
    --fp $genomeDir \
    --fr ../../515F-806R.fna \
    --fld skewed-normal,9000,2500,-5 \
    --flr None,None \
    --nf 10000 \
    --np 24 \
    2> ampFrags.log \
    > ampFrags.pkl        

"""
Explanation: Simulating fragments
End of explanation
"""


!cd $buildDir; \
    SIPSim fragment_kde \
    ampFrags.pkl \
    > ampFrags_kde.pkl

"""
Explanation: Converting to kde object
End of explanation
"""


!cd $buildDir; \
    SIPSim diffusion \
    ampFrags_kde.pkl \
    --np 24 \
    > ampFrags_kde_dif.pkl    

"""
Explanation: Adding diffusion
End of explanation
"""


!cd $buildDir; \
    SIPSim communities \
    $genomeDir/genome_index.txt \
    --n_comm 2 \
    > comm.txt

"""
Explanation: Creating a community file
End of explanation
"""


# building tree structure
nest = nestly.Nest()

## varying params
nest.add('percTaxaIncorp', [5, 10, 25, 50, 75, 100])
nest.add('percIncorp', range(0,101,25))

## set params
nest.add('np_many', [24], create_dir=False)
nest.add('np_few', [8], create_dir=False)
nest.add('abs', ['1e10'], create_dir=False)
#nest.add('subsample', [20000], create_dir=False) 
nest.add('subsample_mean', [30000], create_dir=False) 
nest.add('subsample_scale', [5000], create_dir=False) 
nest.add('BD_min', [1.71], create_dir=False)
nest.add('BD_max', [1.75], create_dir=False)
nest.add('padj', [0.1], create_dir=False)
nest.add('log2', [0.25], create_dir=False)
nest.add('topTaxaToPlot', [100], create_dir=False)

## input/output files
nest.add('buildDir', [buildDir], create_dir=False)
nest.add('frag_file', ['ampFrags_kde_dif'], create_dir=False)
nest.add('comm_file', ['comm.txt'], create_dir=False)
nest.add('genome_index', [os.path.join(genomeDir, 'genome_index.txt')], create_dir=False)
nest.add('genome_dir', [genomeDir], create_dir=False)
nest.add('primers', [os.path.join(workDir, '../', '515F-806R.fna')], create_dir=False)
nest.add('R_dir', [R_dir], create_dir=False)

# building directory tree
nest.build(buildDir)

bashFile = os.path.join(buildDir, 'SIPSimRun.sh')

%%writefile $bashFile
#!/bin/bash

# symlinking input files
ln -s {buildDir}/{frag_file}.pkl {frag_file}.pkl
ln -s {buildDir}/{comm_file} {comm_file}


# making incorp file
SIPSim incorpConfigExample \
  --percTaxa {percTaxaIncorp} \
  --percIncorpUnif {percIncorp} \
  > {percTaxaIncorp}_{percIncorp}.config

 
# adding isotope incorporation to BD distribution
SIPSim isotope_incorp \
    {frag_file}.pkl \
    {percTaxaIncorp}_{percIncorp}.config \
    --comm {comm_file} \
    --np {np_many} \
    > {frag_file}_incorp.pkl
     
# calculating BD shift from isotope incorporation
SIPSim BD_shift \
    {frag_file}.pkl \
    {frag_file}_incorp.pkl \
    --np {np_few} \
    > {frag_file}_incorp_BD-shift.txt
           
    
# simulating gradient fractions
SIPSim gradient_fractions \
    {comm_file} \
    > fracs.txt

# simulating an OTU table
SIPSim OTU_table \
    {frag_file}_incorp.pkl \
    {comm_file} \
    fracs.txt \
    --abs {abs} \
    --np {np_few} \
    > OTU_abs{abs}.txt
        
    
# subsampling from the OTU table (simulating sequencing of the DNA pool)
SIPSim OTU_subsample \
    --dist normal \
    --dist_params loc:{subsample_mean},scale:{subsample_scale} \
    OTU_abs{abs}.txt \
    > OTU_n2_abs{abs}_sub-norm.txt

# making a wide table
SIPSim OTU_wideLong -w \
    OTU_n2_abs{abs}_sub-norm.txt \
    > OTU_n2_abs{abs}_sub-norm_w.txt
    
# making metadata (phyloseq: sample_data)
SIPSim OTU_sampleData \
    OTU_n2_abs{abs}_sub-norm.txt \
    > OTU_n2_abs{abs}_sub-norm_meta.txt

!chmod 775 $bashFile

!cd $workDir; \
    nestrun -j 1 --template-file $bashFile -d percTaxaIncorpUnif --log-file log.txt

"""
Explanation: Using nestly: different incorporation percentages
End of explanation
"""


buildDir_DESeq = os.path.join(buildDir, 'DESeq2')
buildDir_DESeq

# building tree structure
nest = nestly.Nest()

## varying params
nest.add('percTaxaIncorp', [5, 10, 25, 50, 75, 100])
nest.add('percIncorp', range(0,101,25))
### BD range
BD_min = np.arange(1.67, 1.77, 0.02).tolist()
BD_max = [x + 0.04 for x in BD_min]
f = lambda x: {'BD_range': str(x[0]) + '-' + str(x[1]), 
               'BD_min':x[0], 
               'BD_max':x[1]}
BD_range = [f(x) for x in itertools.product(BD_min, BD_max) 
            if x[0] < x[1]]
#### contains BD_min & BD_max
nest.add('BD_range', BD_range, update=True)

## set params
nest.add('np_many', [24], create_dir=False)
nest.add('np_few', [8], create_dir=False)
nest.add('abs', ['1e10'], create_dir=False)
#nest.add('subsample', [20000], create_dir=False) 
nest.add('subsample_mean', [30000], create_dir=False) 
nest.add('subsample_scale', [5000], create_dir=False) 
nest.add('padj', [0.1], create_dir=False)
nest.add('log2', [0.25], create_dir=False)
nest.add('topTaxaToPlot', [100], create_dir=False)

## input/output files
nest.add('buildDir', [buildDir], create_dir=False)
nest.add('frag_file', ['ampFrags_kde_dif'], create_dir=False)
nest.add('comm_file', ['comm.txt'], create_dir=False)
nest.add('genome_index', [os.path.join(genomeDir, 'genome_index.txt')], create_dir=False)
nest.add('genome_dir', [genomeDir], create_dir=False)
nest.add('primers', [os.path.join(workDir, '../', '515F-806R.fna')], create_dir=False)
nest.add('R_dir', [R_dir], create_dir=False)

# building directory tree
nest.build(buildDir_DESeq)

%%writefile $bashFile
#!/bin/bash

# copying all input files
cp {buildDir}/{percTaxaIncorp}/{percIncorp}/OTU_n2_abs{abs}_sub-norm.txt OTU_n2_abs{abs}_sub-norm.txt
cp {buildDir}/{percTaxaIncorp}/{percIncorp}/OTU_n2_abs{abs}_sub-norm_w.txt OTU_n2_abs{abs}_sub-norm_w.txt
cp {buildDir}/{percTaxaIncorp}/{percIncorp}/OTU_n2_abs{abs}_sub-norm_meta.txt OTU_n2_abs{abs}_sub-norm_meta.txt
cp {buildDir}/{percTaxaIncorp}/{percIncorp}/{frag_file}_incorp_BD-shift.txt {frag_file}_incorp_BD-shift.txt


#-- R analysis --#
export PATH={R_dir}:$PATH
# plotting taxon abundances
OTU_taxonAbund.r \
    OTU_n2_abs{abs}_sub-norm.txt \
    -r {topTaxaToPlot} \
    -o OTU_n2_abs{abs}_subsub-norm
# running DeSeq2 and making confusion matrix on predicting incorporators
## making phyloseq object from OTU table
phyloseq_make.r \
    OTU_n2_abs{abs}_sub-norm_w.txt \
    -s OTU_n2_abs{abs}_sub-norm_meta.txt \
    > OTU_n2_abs{abs}_sub-norm.physeq
## filtering phyloseq object to just taxa/samples of interest
phyloseq_edit.r \
    OTU_n2_abs{abs}_sub-norm.physeq \
    --BD_min {BD_min} \
    --BD_max {BD_max} \
    > OTU_n2_abs{abs}_sub-norm_filt.physeq
## making ordination
#phyloseq_ordination.r \
#    OTU_n2_abs{abs}_sub-norm_filt.physeq \
#    OTU_n2_abs{abs}_sub-norm_bray-NMDS.pdf
## DESeq2
phyloseq_DESeq2.r \
    OTU_n2_abs{abs}_sub-norm_filt.physeq \
    --log2 {log2} \
    --hypo greater \
    > OTU_n2_abs{abs}_sub-norm_DESeq2
## Confusion matrix
DESeq2_confuseMtx.r \
    {frag_file}_incorp_BD-shift.txt \
    OTU_n2_abs{abs}_sub-norm_DESeq2 \
    --padj {padj} 

%%R -i buildDir
# debug
setwd(buildDir)

physeq = readRDS('50/75/1.69-1.77/OTU_n2_abs1e10_sub-norm.physeq')



!chmod 775 $bashFile

!cd $workDir; \
    nestrun -j 30 --template-file $bashFile -d percTaxaIncorpUnif/DESeq2 --log-file log.txt

# Creating NA tables for any failures
sDir = os.path.join(buildDir, 'DESeq2', '*', '*', '*')
dirs = glob.glob(sDir)

tbl = """
,byClass
Sensitivity,NA
Specificity,NA
Pos Pred Value,NA
Neg Pred Value,NA
Prevalence,NA
Detection Rate,NA
Detection Prevalence,NA
Balanced Accuracy,NA
"""

for x in dirs:
    f = os.path.join(x, 'DESeq2-cMtx_byClass.csv')
    if not os.path.isfile(f):
        with open(f, 'wb') as inFH:
            inFH.write(tbl)
        

# aggregating confusion matrix data
## table
# !cd $buildDir; \
#     nestagg delim \
#         -d DESeq2 \
#         -k percIncorp,percTaxaIncorp,BD_min,BD_max \
#         -o ./DESeq2/cMtx_table.csv \
#         DESeq2-cMtx_table.csv

# ## overall
# !cd $buildDir; \
#     nestagg delim \
#         -d DESeq2 \
#         -k percIncorp,percTaxaIncorp,BD_min,BD_max \
#         -o ./DESeq2/cMtx_overall.csv \
#         DESeq2-cMtx_overall.csv
    
## byClass
!cd $buildDir; \
    nestagg delim \
        -d DESeq2 \
        -k percIncorp,percTaxaIncorp,BD_min,BD_max \
        -o ./DESeq2/cMtx_byClass.csv \
        DESeq2-cMtx_byClass.csv

"""
Explanation: DESeq2 Analysis pipeline
End of explanation
"""


%%R -i buildDir -w 600 -h 600
setwd(buildDir)

byClass = read.csv('./DESeq2/cMtx_byClass.csv')

byClass$byClass[is.na(byClass$byClass)] = 0

byClass %>% head

%%R -w 800 -h 750

col2keep = c('Balanced Accuracy', 'Sensitivity','Specificity')
byClass.f = byClass %>%
    filter(X %in% col2keep) %>%
    mutate(BD_min = as.numeric(BD_min),
           BD_max = as.numeric(BD_max),
           byClass = as.numeric(byClass),
           byClass.inv = 1 - byClass) 

byClass.fs = byClass.f %>%
    group_by(X, percIncorp, percTaxaIncorp) %>%
    summarize(byClass.max = max(byClass))

just.true = function(x){
    if(x == TRUE){
        return(1)
        } else{
        return(NA)
        }
    }

byClass.j = inner_join(byClass.f, byClass.fs, c('X' = 'X',
                                                'percIncorp' = 'percIncorp',
                                                'percTaxaIncorp' = 'percTaxaIncorp')) %>%
    mutate(max_val = as.numeric(byClass == byClass.max),
           byClass.txt = round(byClass, 2))
byClass.jf = byClass.j %>%
    filter(max_val == 1)

byClass.j %>% head

%%R

x.breaks = unique(byClass.j$BD_min)

taxaIncorp = unique(byClass.j$percTaxaIncorp)

for (x in taxaIncorp){
    tbl = byClass.j %>% filter(percTaxaIncorp == x)
    
    p = ggplot(tbl, aes(BD_min, BD_max, fill=byClass.inv)) +
    geom_tile() +
    geom_text(aes(label=byClass.txt), color=c('white'), size=4) +
    geom_text(data=byClass.jf, aes(label=byClass.txt), color=c('red'), size=4) +
    scale_x_continuous(breaks=x.breaks) +
    labs(x='Minimum Buoyant Density', y='Maximum Buoyant Density') +
    facet_grid(percIncorp ~ X) +
    theme_bw() +
    theme(
        text=element_text(size=16),
        axis.text.x=element_text(angle=90, hjust=1, vjust=0.5)
        )
    
    print(p)
    }


%%R -i workDir -w 600 -h 600
setwd(workDir)

byClass = read.csv('./percIncorpUnif/DESeq2-cMtx_byClass.csv')

byClass$byClass[is.na(byClass$byClass)] = 0

byClass$percIncorp = factor(byClass$percIncorp, levels=as.character(unique(sort(byClass$percIncorp))))
p = ggplot(byClass, aes(X, byClass, color=percIncorp)) +
    geom_point() +
    labs(y='value') +
    theme_bw() +
    theme(
        text=element_text(size=16),
        axis.text.x=element_text(angle=90, hjust=1, vjust=0.5),
        axis.title.x=element_blank()
        )
p

%%R -w 600 -h 400

p + scale_y_continuous(limits=c(0,0.15))

"""
Explanation: Plotting results
End of explanation
"""


%%R -i workDir

setwd(workDir)

tbl.c = read.csv('percIncorpUnif/50/DESeq2-cMtx_data.csv')

tbl.c.f = tbl.c %>% 
    filter(incorp.known==TRUE & incorp.pred==FALSE) %>%
    select(taxon)

write.table(tbl.c.f, file='percIncorpUnif/50/DESeq2-cMtx_FN', row.names=F, col.names=F, quote=F)

!cd $workDir'percIncorpUnif/50/'; \
    $R_dir/OTU_taxonAbund.r OTU_n2_abs1e8_sub30000.txt -t DESeq2-cMtx_FN -l -o DESeq2-cMtx_FN --width 14

%%R -i workDir

setwd(workDir)

tbl.c = read.csv('percIncorpUnif/50/DESeq2-cMtx_data.csv')

tbl.c.f = tbl.c %>% 
    filter(incorp.known==FALSE & incorp.pred==TRUE) %>%
    select(taxon)

write.table(tbl.c.f, file='percIncorpUnif/50/DESeq2-cMtx_FP', row.names=F, col.names=F, quote=F)

!cd $workDir'percIncorpUnif/50/'; \
    $R_dir/OTU_taxonAbund.r OTU_n2_abs1e8_sub30000.txt -t DESeq2-cMtx_FP -l -o DESeq2-cMtx_FP --width 14

%%R
tbl.c %>% filter(taxon %in% tbl.c.f$taxon)

"""
Explanation: Plotting false negative taxa (actual incorporators)
End of explanation
"""


%%R -i workDir

setwd(workDir)

tbl.c = read.csv('percIncorpUnif/100/DESeq2-cMtx_data.csv')
tbl.otu = read.delim('percIncorpUnif/100/OTU_n2_abs1e10_sub20000.txt', sep='\t')

%%R
# OTU total counts
tbl.otu.sum = tbl.otu %>%
    group_by(library, taxon) %>%
    summarize(total_count = sum(count))
tbl.otu.sum %>% head

%%R
# 
label.tp.fn = function(known, pred){
    if(is.na(known) | is.na(pred)){
        return(NA)
        } else
    if(known==TRUE & pred==TRUE){
        return('TP')
        } else
    if(known==TRUE & pred==FALSE){
        return('FN')
        } else {
        return(NA)
        }
    }

tbl.c.tp.fn = tbl.c %>% 
    mutate(tp.fn = mapply(label.tp.fn, incorp.known, incorp.pred)) %>%
    filter(! is.na(tp.fn))

tbl.tp.fn = inner_join(tbl.c.tp.fn, tbl.otu.sum, c('taxon' = 'taxon')) 
tbl.tp.fn %>% head

%%R
# how many TP & FN?
tbl.tp.fn %>% 
    group_by(library, tp.fn) %>%
    summarize(n = n())

%%R -h 350
tbl.tp.fn$library = as.character(tbl.tp.fn$library)
ggplot(tbl.tp.fn, aes(library, total_count, color=tp.fn)) +
    geom_boxplot() +
    labs(y='Total count') +
    theme(
        text = element_text(size=18)
        )  

%%R -h 700 -w 900
tbl.tp.fn$library = as.character(tbl.tp.fn$library)
ggplot(tbl.tp.fn, aes(taxon, total_count, group=taxon, color=library)) +
    geom_point(size=3) +
    geom_line() +
    scale_y_log10() +
    labs(y='Total count') +
    facet_grid(. ~ tp.fn, scales='free_x') +
    theme(
        text = element_text(size=18),
        axis.text.x = element_text(angle=90, hjust=1)
        )  

"""
Explanation: Enrichment of TP for abundant incorporators?

What is the abundance distribution of TP and FP?
Are more abundant incorporators being detected more than low abundant taxa
End of explanation
"""


%%R
# OTU total counts

heavy.cut = 1.71

tbl.otu.sum = tbl.otu %>%
    filter(! grepl('inf', fraction)) %>%
    separate(fraction, into = c('BD_min','BD_max'), sep='-', convert=TRUE) %>%
    filter(BD_min >= heavy.cut & BD_max <= 2) %>%
    group_by(library, taxon) %>%
    summarize(total_count = sum(count))
tbl.otu.sum %>% head

%%R

tbl.tp.fn = inner_join(tbl.c.tp.fn, tbl.otu.sum, c('taxon' = 'taxon')) 
tbl.tp.fn %>% head

%%R -h 350
tbl.tp.fn$library = as.character(tbl.tp.fn$library)
ggplot(tbl.tp.fn, aes(library, total_count, color=tp.fn)) +
    geom_boxplot() +
    labs(y='Total count') +
    theme(
        text = element_text(size=18)
        )  

%%R -h 700 -w 900
tbl.tp.fn$library = as.character(tbl.tp.fn$library)
ggplot(tbl.tp.fn, aes(taxon, total_count, group=taxon, color=library)) +
    geom_point(size=3) +
    geom_line() +
    scale_y_log10() +
    labs(y='Total count') +
    facet_grid(. ~ tp.fn, scales='free_x') +
    theme(
        text = element_text(size=18),
        axis.text.x = element_text(angle=90, hjust=1)
        )  

%%R -i workDir -w 1000 -h 450
setwd(workDir)

tbl.ds = read.csv('DESeq2-cMtx_data.csv')

# loading file
tbl.otu = read.delim('OTU_abs1e10.txt', sep='\t')

tbl.otu = tbl.otu %>% 
    filter(!grepl('inf', fraction, ignore.case=T)) %>%
    separate(fraction, into = c('BD_min','BD_max'), sep='-', convert=TRUE) %>%

tbl.j = inner_join(tbl.otu, tbl.ds, c('taxon' = 'taxon'))

"""
Explanation: Notes:

FNs can be abundant
TP can have a lower abundance in the 'treatment' library

Just looking at the 'heavy' fractions
End of explanation
"""


# building tree structure
from os.path import abspath
nest = nestly.Nest()

##  values
vals = [str(x) for x in range(1,5)]
nest.add('vals', vals)


## input files
nest.add('--np', [1], create_dir=False)

buildDir = '/home/nick/t/nestly/' #os.path.join(workDir, 'vals')
nest.build(buildDir)

%%writefile /home/nick/t/example.sh
#!/bin/bash

export TIME='elapsed,maxmem,exitstatus\n%e,%M,%x'

echo {--np} > {--np}_test.txt

!cd /home/nick/t/; \
    chmod 777 example.sh 

!cd /home/nick/t/; \
    nestrun -j 2 --template-file example.sh -d nestly

%%writefile $bashFile
#!/bin/bash


#-- R analysis --#
export PATH={R_dir}:$PATH
# plotting taxon abundances
OTU_taxonAbund.r \
    OTU_n2_abs{abs}_sub{subsample}.txt \
    -r {topTaxaToPlot} \
    -o OTU_n2_abs{abs}_sub{subsample} 
# running DeSeq2 and making confusion matrix on predicting incorporators
## making phyloseq object from OTU table
phyloseq_make.r \
    OTU_n2_abs{abs}_sub{subsample}_w.txt \
    -s OTU_n2_abs{abs}_sub{subsample}_meta.txt \
    > OTU_n2_abs{abs}_sub{subsample}.physeq
## filtering phyloseq object to just taxa/samples of interest
phyloseq_edit.r \
    OTU_n2_abs{abs}_sub{subsample}.physeq \
    --BD_min {BD_min} --BD_max {BD_max} \
    > OTU_n2_abs{abs}_sub{subsample}_filt.physeq
## making ordination
phyloseq_ordination.r \
    OTU_n2_abs{abs}_sub{subsample}_filt.physeq \
    OTU_n2_abs{abs}_sub{subsample}_bray-NMDS.pdf
## DESeq2
phyloseq_DESeq2.r \
    OTU_n2_abs{abs}_sub{subsample}_filt.physeq \
    > OTU_n2_abs{abs}_sub{subsample}_DESeq2
## Confusion matrix
DESeq2_confuseMtx.r \
    {fileName}_kde_dif_incorp_BD-shift.txt \
    OTU_n2_abs{abs}_sub{subsample}_DESeq2 \
    --padj {padj} --log2 {log2} 
    
    

"""
Explanation: Notes:

The TPs (for the most part) are dramatically different in abundance between control and treatment


Sandbox
End of explanation
"""


