workDir = '../../t/SIPSim_example/'
nprocs = 3

"""
Explanation: Description

Time to make a simple SIP data simulation with the dataset that you alreadly created


Make sure you have created the dataset before trying to run this notebook

Setting variables

"workDir" is the path to the working directory for this analysis (where the files will be download to)
NOTE: MAKE SURE to modify this path to the directory where YOU want to run the example.
"nprocs" is the number of processors to use (3 by default, since only 3 genomes). Change this if needed.
End of explanation
"""


import os

# Note: you will need to install `rpy2.ipython` and the necessary R packages (see next cell)
%load_ext rpy2.ipython

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

workDir = os.path.abspath(workDir)
if not os.path.isdir(workDir):
    os.makedirs(workDir)
%cd $workDir    

genomeDir = os.path.join(workDir, 'genomes_rn')

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


%%bash
source activate SIPSim

# creating example config
SIPSim incorp_config_example \
  --percTaxa 34 \
  --percIncorpUnif 50 \
  --n_reps 1 \
  > incorp.config

!cat incorp.config

"""
Explanation: Experimental design

How many gradients?
Which are labeled treatments & which are controls?
For this tutorial, we'll keep things simple and just simulate one control & one treatment
For the labeled treatment, 34% of the taxa (1 of 3) will incorporate 50% isotope

The script below ("SIPSim incorpConfigExample") is helpful for making simple experimental designs
End of explanation
"""


%%bash
source activate SIPSim

SIPSim communities \
    --config incorp.config \
    ./genomes_rn/genome_index.txt \
    > comm.txt  

!cat comm.txt

"""
Explanation: Pre-fractionation communities

What is the relative abundance of taxa in the pre-fractionation samples?
End of explanation
"""


%%bash 
source activate SIPSim

SIPSim gradient_fractions \
    --BD_min 1.67323 \
    --BD_max 1.7744 \
    comm.txt \
    > fracs.txt    

!head -n 6 fracs.txt

"""
Explanation: Note: "library" = gradient
Simulating gradient fractions

BD size ranges for each fraction (& start/end of the total BD range)
End of explanation
"""


# primers = """>515F
# GTGCCAGCMGCCGCGGTAA
# >806R
# GGACTACHVGGGTWTCTAAT
# """

# F = os.path.join(workDir, '515F-806R.fna')
# with open(F, 'wb') as oFH:
#     oFH.write(primers)
    
# print 'File written: {}'.format(F)

"""
Explanation: Simulating fragments

Simulating shotgun-fragments
Fragment length distribution: skewed-normal

Primer sequences (wait... what?)

If you were to simulate amplicons, instead of shotgun fragments, you can use something like the following:
End of explanation
"""


%%bash -s $genomeDir
source activate SIPSim 

# skewed-normal
SIPSim fragments \
    $1/genome_index.txt \
    --fp $1 \
    --fld skewed-normal,9000,2500,-5 \
    --flr None,None \
    --nf 1000 \
    --debug \
    --tbl \
    > shotFrags.txt       

!head -n 5 shotFrags.txt
!tail -n 5 shotFrags.txt

"""
Explanation: Simulation
End of explanation
"""


%%R -w 700 -h 350

df = read.delim('shotFrags.txt')

p = ggplot(df, aes(fragGC, fragLength, color=taxon_name)) +
    geom_density2d() +
    scale_color_discrete('Taxon') +
    labs(x='Fragment G+C', y='Fragment length (bp)') +
    theme_bw() +
    theme(
        text = element_text(size=16)
    )
plot(p)

"""
Explanation: Plotting fragments
End of explanation
"""


%%bash 
source activate SIPSim

SIPSim fragment_KDE \
    shotFrags.txt \
    > shotFrags_kde.pkl

!ls -thlc shotFrags_kde.pkl

"""
Explanation: Note: for information on what's going on in this config file, use the command: SIPSim isotope_incorp -h
Converting fragments to a 2d-KDE

Estimating the joint-probabilty for fragment G+C & length
End of explanation
"""


%%bash 
source activate SIPSim

SIPSim diffusion \
    shotFrags_kde.pkl \
    --np 3 \
    > shotFrags_kde_dif.pkl        

!ls -thlc shotFrags_kde_dif.pkl 

"""
Explanation: Note: The generated list of KDEs (1 per taxon per gradient) are in a binary file format
To get a table of length/G+C values, use the command: SIPSim KDE_sample

Adding diffusion

Simulating the BD distribution of fragments as Gaussian distributions. 
One Gaussian distribution per homogeneous set of DNA molecules (same G+C and length)


See the README if you get MKL errors with the next step and re-run the fragment KDE generation step
End of explanation
"""


n = 100000

%%bash -s $n
source activate SIPSim

SIPSim KDE_sample -n $1 shotFrags_kde.pkl > shotFrags_kde.txt
SIPSim KDE_sample -n $1 shotFrags_kde_dif.pkl > shotFrags_kde_dif.txt

ls -thlc shotFrags_kde*.txt

"""
Explanation: Plotting fragment distribution w/ and w/out diffusion
Making a table of fragment values from KDEs
End of explanation
"""


%%R
df1 = read.delim('shotFrags_kde.txt', sep='\t')
df2 = read.delim('shotFrags_kde_dif.txt', sep='\t')

df1$data = 'no diffusion'
df2$data = 'diffusion'
df = rbind(df1, df2) %>%
    gather(Taxon, BD, Clostridium_ljungdahlii_DSM_13528, 
           Escherichia_coli_1303, Streptomyces_pratensis_ATCC_33331) %>%
    mutate(Taxon = gsub('_(ATCC|DSM)', '\n\\1', Taxon))

df %>% head(n=3)

%%R -w 800 -h 300

p = ggplot(df, aes(BD, fill=data)) +
    geom_density(alpha=0.25) +
    facet_wrap( ~ Taxon) +    
    scale_fill_discrete('') +
    theme_bw() +
    theme(
        text=element_text(size=16),
        axis.title.y = element_text(vjust=1),
        axis.text.x = element_text(angle=50, hjust=1)
        )

plot(p)

"""
Explanation: Plotting

plotting KDE with or without diffusion added
End of explanation
"""


%%bash 
source activate SIPSim

SIPSim DBL \
    shotFrags_kde_dif.pkl \
    --np 3 \
    > shotFrags_kde_dif_DBL.pkl

# viewing DBL logs
!ls -thlc *pkl

"""
Explanation: Adding diffusive boundary layer (DBL) effects

'smearing' effects
End of explanation
"""


%%bash
source activate SIPSim

SIPSim isotope_incorp \
    --comm comm.txt \
    --np 3 \
    shotFrags_kde_dif_DBL.pkl \
    incorp.config \
    > shotFrags_KDE_dif_DBL_inc.pkl

!ls -thlc *.pkl

"""
Explanation: Adding isotope incorporation

Using the config file produced in the Experimental Design section
End of explanation
"""


%%R
df = read.delim('BD-shift_stats.txt', sep='\t')
df 

"""
Explanation: Note: statistics on how much isotope was incorporated by each taxon are listed in "BD-shift_stats.txt"
End of explanation
"""


%%bash
source activate SIPSim

SIPSim OTU_table \
    --abs 1e7 \
    --np 3 \
    shotFrags_KDE_dif_DBL_inc.pkl \
    comm.txt \
    fracs.txt \
    > OTU.txt

!head -n 7 OTU.txt

"""
Explanation: Making an OTU table

Number of amplicon-fragment in each fraction in each gradient
Assuming a total pre-fractionation community size of 1e7
End of explanation
"""


%%R -h 350 -w 750

df = read.delim('OTU.txt', sep='\t')

p = ggplot(df, aes(BD_mid, count, fill=taxon)) +
    geom_area(stat='identity', position='dodge', alpha=0.5) +
    scale_x_continuous(expand=c(0,0)) +
    labs(x='Buoyant density') +
    labs(y='Shotgun fragment counts') +
    facet_grid(library ~ .) +
    theme_bw() +
    theme( 
        text = element_text(size=16),
        axis.title.y = element_text(vjust=1),        
        axis.title.x = element_blank()
    )
plot(p)

"""
Explanation: Plotting fragment count distributions
End of explanation
"""


%%R -h 350 -w 750

p = ggplot(df, aes(BD_mid, count, fill=taxon)) +
    geom_area(stat='identity', position='fill') +
    scale_x_continuous(expand=c(0,0)) +
    scale_y_continuous(expand=c(0,0)) +
    labs(x='Buoyant density') +
    labs(y='Shotgun fragment counts') +
    facet_grid(library ~ .) +
    theme_bw() +
    theme( 
        text = element_text(size=16),
        axis.title.y = element_text(vjust=1),        
        axis.title.x = element_blank()
    )
plot(p)

"""
Explanation: Notes: 

This plot represents the theoretical number of amplicon-fragments at each BD across each gradient. 
Derived from subsampling the fragment BD proability distributions generated in earlier steps.
The fragment BD distribution of one of the 3 taxa should have shifted in Gradient 2 (the treatment gradient).
The fragment BD distributions of the other 2 taxa should be approx. the same between the two gradients.

Viewing fragment counts as relative quantities
End of explanation
"""


%%bash
source activate SIPSim

SIPSim OTU_PCR OTU.txt > OTU_PCR.txt    

!head -n 5 OTU_PCR.txt
!tail -n 5 OTU_PCR.txt

"""
Explanation: Adding effects of PCR

This will alter the fragment counts based on the PCR kinetic model of:
Suzuki MT, Giovannoni SJ. (1996). Bias caused by template annealing in the
  amplification of mixtures of 16S rRNA genes by PCR. Appl Environ Microbiol
  62:625-630.
End of explanation
"""


%%bash
source activate SIPSim

SIPSim OTU_subsample OTU_PCR.txt > OTU_PCR_sub.txt

!head -n 5 OTU_PCR_sub.txt

"""
Explanation: Notes

The table is in the same format as with the original OTU table, but the counts and relative abundances should be altered.

Simulating sequencing

Sampling from the OTU table
End of explanation
"""


%%R -h 350 -w 750

df = read.delim('OTU_PCR_sub.txt', sep='\t')


p = ggplot(df, aes(BD_mid, rel_abund, fill=taxon)) +
    geom_area(stat='identity', position='fill') +
    scale_x_continuous(expand=c(0,0)) +
    scale_y_continuous(expand=c(0,0)) +
    labs(x='Buoyant density') +
    labs(y='Taxon relative abundances') +
    facet_grid(library ~ .) +
    theme_bw() +
    theme( 
        text = element_text(size=16),
        axis.title.y = element_text(vjust=1),        
        axis.title.x = element_blank()
    )
plot(p)

"""
Explanation: Notes

The table is in the same format as with the original OTU table, but the counts and relative abundances should be altered.

Plotting
End of explanation
"""


%%bash
source activate SIPSim

SIPSim OTU_wide_long -w \
    OTU_PCR_sub.txt \
    > OTU_PCR_sub_wide.txt

!head -n 4 OTU_PCR_sub_wide.txt

"""
Explanation: Misc
A 'wide' OTU table

If you want to reformat the OTU table to a more standard 'wide' format (as used in Mothur or QIIME):
End of explanation
"""


%%bash
source activate SIPSim

SIPSim OTU_sample_data \
    OTU_PCR_sub.txt \
    > OTU_PCR_sub_meta.txt

!head OTU_PCR_sub_meta.txt

"""
Explanation: SIP metadata

If you want to make a table of SIP sample metadata
End of explanation
"""


%%bash
source activate SIPSim

SIPSim -l

"""
Explanation: Other SIPSim commands
SIPSim -l will list all available SIPSim commands
End of explanation
"""


