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

"""
Explanation: Goal
Accuracy as a function of community evenness and isotope incorporation
Variable parameters:

Community rank-abundance:
lognormal abundance distribution
very uneven
moderately uneven
slightly uneven


% incorp
range = 0:100:5
n-reps (stocastic: taxon abundances & which incorporate)
20

Set parameters:

% taxa that incorporate
10%?
Total community abundance:
1e9
'heavy' BD range:
1.71 - 1.75

User variables
End of explanation
"""


import glob
from os.path import abspath
import nestly
from IPython.display import Image, display

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


!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: Creating input files (eg., fragments & communities)
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 \
    --abund_dist_p=mean:10,sigma:1 \
    --n_comm 1 \
    > comm_m10-s1.txt

!cd $buildDir; \
    SIPSim communities \
    $genomeDir/genome_index.txt \
    --abund_dist_p=mean:10,sigma:2 \
    --n_comm 1 \
    > comm_m10-s2.txt

!cd $buildDir; \
    SIPSim communities \
    $genomeDir/genome_index.txt \
    --abund_dist_p=mean:10,sigma:3 \
    --n_comm 1 \
    > comm_m10-s3.txt

%%R -i buildDir
# loading comm files
setwd(buildDir)

inFiles = c('comm_m10-s1.txt', 'comm_m10-s2.txt', 'comm_m10-s3.txt')

tbl.l = list()
for (x in inFiles){
    tbl.l[[x]] = read.delim(x, sep='\t')
    }
tbl = do.call(rbind, tbl.l)

%%R -w 800

tbl$distribution = gsub('[0-9]+$', '', rownames(tbl))

ggplot(tbl, aes(rank, rel_abund_perc, color=distribution)) +
    geom_point() +
    scale_y_log10() +
    labs(x='Rank', y='Percent relative abundance') +
    theme_bw() +
    theme(
        text=element_text(size=16)
        )

%%R
# min-max relative abundance 

tbl %>% 
    group_by(distribution) %>%
    summarize(
        min_abund = min(rel_abund_perc),
        max_abund = max(rel_abund_perc),
        min_max_ratio = max_abund / min_abund
        )

"""
Explanation: Making example plots of community size distribution
End of explanation
"""


# building tree structure
nest = nestly.Nest()

## varying params
#-- full run --#
nest.add('abund_dist_p', ['mean:10,sigma:1','mean:10,sigma:2','mean:10,sigma:3'])
nest.add('percIncorp', [0,5,10,15,20,30,40,50,75,100])
nest.add('rep', range(1,21))

## set params
nest.add('percTaxa', [10], create_dir=False)
nest.add('np', [6], create_dir=False)
nest.add('abs', ['1e9'], 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('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

# Creating a community file
SIPSim communities \
    {genome_index} \
    --abund_dist_p {abund_dist_p} \
    --n_comm 2 \
    > comm.txt    

# simulating gradient fractions
SIPSim gradient_fractions \
    {comm_file} \
    > fracs.txt
    
# making incorp file
SIPSim incorpConfigExample \
  --percTaxa {percTaxa} \
  --percIncorpUnif {percIncorp} \
  > {percTaxa}_{percIncorp}.config

# adding isotope incorporation to BD distribution
SIPSim isotope_incorp \
    {frag_file}.pkl \
    {percTaxa}_{percIncorp}.config \
    --comm {comm_file} \
    --np {np} \
    > {frag_file}_incorp.pkl
 
# calculating BD shift from isotope incorporation
SIPSim BD_shift \
    {frag_file}.pkl \
    {frag_file}_incorp.pkl \
    --np {np} \
    > {frag_file}_incorp_BD-shift.txt

# simulating an OTU table
SIPSim OTU_table \
    {frag_file}_incorp.pkl \
    {comm_file} \
    fracs.txt \
    --abs {abs} \
    --np {np} \
    > OTU_n2_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_n2_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 5 --template-file $bashFile -d atomIncorp_evenness --log-file log.txt

"""
Explanation: Running nestly
End of explanation
"""


%%writefile $bashFile
#!/bin/bash

#-- R analysis --#
export PATH={R_dir}:$PATH
    
# plotting taxon abundances
OTU_taxonAbund.r \
    OTU_n2_abs{abs}.txt \
    -r {topTaxaToPlot} \
    -o OTU_n2_abs{abs}
# 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} 

!chmod 775 $bashFile

!cd $workDir; \
    nestrun -j 30 --template-file $bashFile -d atomIncorp_evenness --log-file logR.txt

# aggregating confusion matrix data
## table
!cd $workDir; \
    nestagg delim \
        -d atomIncorp_evenness \
        -k abund_dist_p,percIncorp,rep \
        -o ./atomIncorp_evenness/DESeq2-cMtx_table.csv \
        DESeq2-cMtx_table.csv

## overall
!cd $workDir; \
    nestagg delim \
        -d atomIncorp_evenness \
        -k abund_dist_p,percIncorp,rep \
        -o ./atomIncorp_evenness/DESeq2-cMtx_overall.csv \
        DESeq2-cMtx_overall.csv
    
## byClass
!cd $workDir; \
    nestagg delim \
        -d atomIncorp_evenness \
        -k abund_dist_p,percIncorp,rep \
        -o ./atomIncorp_evenness/DESeq2-cMtx_byClass.csv \
        DESeq2-cMtx_byClass.csv

"""
Explanation: R analysis
End of explanation
"""


%%R -i workDir
setwd(workDir)

byClass = read.csv('./atomIncorp_evenness/DESeq2-cMtx_byClass.csv')
#byClass$byClass[is.na(byClass$byClass)] = 0
byClass %>% head

%%R -w 700
to.keep = c('Sensitivity', 'Specificity', 'Balanced Accuracy')

byClass.f = byClass %>%
    filter(X %in% to.keep) %>%
    mutate(percIncorp = as.character(percIncorp))

byClass.f$percIncorp = factor(byClass.f$percIncorp, levels=sort(unique(as.numeric(byClass.f$percIncorp))))

ggplot(byClass.f, aes(percIncorp, byClass, fill=abund_dist_p)) +
    geom_boxplot() +
    labs(y='Value', x='Atom percent 13C') +
    facet_grid(X ~ .) +
        theme(
        text = element_text(size=16)
        )

%%R -w 700 -h 400

byClass.f.s = byClass.f %>%
    group_by(X, abund_dist_p, percIncorp) %>%
    summarize(mean.val = mean(byClass),
              sd.val = sd(byClass)) %>%
    ungroup() %>%
    mutate(percIncorp = as.numeric(as.character(percIncorp)))

#byClass.f.s %>% head

ggplot(byClass.f.s, aes(percIncorp, mean.val, group=abund_dist_p, fill=abund_dist_p)) +
    geom_smooth(method='loess') +
    labs(y='Value', x='Atom percent 13C') +
    facet_grid(X ~ .) +
        theme(
        text = element_text(size=16)
        )

%%R -w 800 -h 450

byClass.f.s = byClass.f %>%
    group_by(X, abund_dist_p, percIncorp) %>%
    summarize(mean.val = mean(byClass),
              sd.val = sd(byClass)) %>%
    ungroup() %>%
    mutate(percIncorp = as.numeric(as.character(percIncorp)))


ggplot(byClass.f.s, aes(percIncorp, mean.val, group=abund_dist_p, color=abund_dist_p,
                       ymin=mean.val - sd.val, ymax=mean.val + sd.val)) +
    geom_pointrange(alpha=0.6, size=0.7) +
    geom_line() +
    labs(y='Value', x='Atom percent 13C') +
    theme_bw() +
    facet_grid(X ~ .) +
        theme(
        text = element_text(size=16)
        )

"""
Explanation: Analyzing the data

figure: 
pointrange (percent_incorp ~ accuracy_value; color by community evenness)
End of explanation
"""


