#!/usr/bin/env python
# coding: utf-8

# # <span style="color:gray">ipyrad-analysis toolkit:</span> window_extracter
# 
# The `window_extracter` tool is used to extract sequence data from a selected genomic window, apply filters, concatenate, and write to a phylip or nexus file. This is useful for combining multiple reference mapped RAD loci into a single large contig, for example around a gene of interest, or for applying filtering options to a denovo or refmapped dataset to create new supermatrices for downstream phylogenetic analysis. It is particularly convenient for taking just the single largest RAD-seq assembly of a dataset (e.g., min4 coverage) and applying different filtering options to create alignment matrices for downstream phylogenetic analyses that each include different subsets of taxa and sites.
# 
# Key features:
# 
# 1. Concatenate ref-mapped RAD loci in genomic windows.
# 2. Apply filters to remove sites or taxa by missing data.
# 3. Optionally apply consensus calling to reduce missingness by reducing multiple samples to a single representative.

# <img src="https://eaton-lab.org/slides/data-svg/window-extracter-min4.svg">

# ### Required software

# In[1]:


# conda install ipyrad -c conda-forge -c bioconda


# In[2]:


import ipyrad.analysis as ipa
ipa.__version__


# ### Required input data files
# Your input data should be a `.seqs.hdf` database file produced by ipyrad. This file contains the full sequence alignment for your samples as well as associated meta-data of the genomic positions of RAD loci relative to a reference genome. 

# In[3]:


# path to an HDF5 formatted seqs file
SEQSFILE = "/tmp/oaks.seqs.hdf5"


# In[6]:


# download example seqs file if not already present (~500Mb, takes ~5 minutes)
URL = "https://www.dropbox.com/s/c1u89nwuuv8e6ie/virentes_ref.seqs.hdf5?raw=1"
ipa.download(URL, path=SEQSFILE);


# ### The scaffold table
# 
# The `window_extracter()` tool takes the `.seqs.hdf5` database file from ipyrad as its input file. You select scaffolds by their index (integer) which can be found in the `.scaffold_table`. We can see from the table below that this genome has 12 large scaffolds (chromosome-scale linkage blocks) and many other smaller unplaced scaffolds. If you are working with a high quality reference genome then it will likely look similar to this, whereas many other reference genomes will be composed of many more scaffolds that are mostly smaller in size. Here I will focus just on the large chromosomes. Take note that the scaffolds are ordered the same as in the reference genome file, which may not involve the largest scaffolds coming first, so you may need to sort the table by length.

# In[8]:


# first load the data file with no other arguments to see scaffold table
ext = ipa.window_extracter(SEQSFILE)

# the scaffold table shows scaffold names and lens in order of the ref. genome
ext.scaffold_table.head(15)


# ### Selecting scaffolds
# The `scaffold_idxs` designates the scaffold to extract sequence data from. This is the index (row) of the named scaffold from the scaffold table (e.g., above). The `window_extracter` tool will select all RAD data within this window and exclude any sites that have no data (e.g., the space between RAD markers, or the space between paired reads) to create a clean concise alignment. 
# 
# The `.stats` attribute shows the information content of the selected window before and after filtering. The stats are returned as a dataframe, showing the size, information content, missingness, and number of samples in the alignment. For example, the code cell below selects all RAD data that mapped to the first scaffold (index 0). Very few sites were filtered because we have not yet applied any filtering options (the few that were removed represent sites that are all Ns).

# In[9]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
)

# show stats of the window
ext.stats


# ### Subsetting scaffold windows
# You can use the `start` and `end` arguments to select subsets of scaffolds as smaller window sizes to be extracted. As with the example above the selected window will be filtered to reduce missing data. If there is no data in the selected window the stats will show no sites, and a warning will be printed. An example with no data and with some data are both shown below. 

# In[11]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
    start=0, 
    end=10000,
)

# show stats of the window
ext.stats


# In[12]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
    start=500000, 
    end=800000,
)

# show stats of the window
ext.stats


# ### Filtering missing data with `mincov`
# 
# You can filter sites from the alignment by using `mincov`, which applies a filter to all sites in the alignment. For example, `mincov=0.5` will require that 50% of samples contain a site that is not `N` or `-` for the site to be included in the alignment. This value can be a proportion like 0.5, or it can be a number, like 10. Similarly, the argument `rmincov` filters rows (taxa) from the matrix if they have data for less than some percentage of sites in the matrix.
# 
# <img src="https://eaton-lab.org/slides/data-svg/window-extracter-min8.svg">

# In[15]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
    start=500000, 
    end=800000,
    mincov=0.8,
    rmincov=0.5,
)

# show stats of the window
ext.stats


# ### Filtering missing data with `imap` and `minmap`
# 
# An `imap` dictionary can be used to group samples into populations/species, as in the example below. It takes key,value pairs where the key is the name of the group, and the value is a list of sample names. One way to use an `imap` is to apply a `minmap` filter. This acts just like the global `mincov` filter, but applies to each group separately. Only if a site meets the minimum coverage argument for each group will it be retained in the data set. In this case the `imap` sampling selected 28 samples to include in the dataset and required 75% of data in each group which reduced the number of SNPs from 105 to 76. 

# In[20]:


# assign samples to groups/taxa
imap = {
    "reference": ["reference"],
    "virg": ["TXWV2", "LALC2", "SCCU3", "FLSF33", "FLBA140"],
    "mini": ["FLSF47", "FLMO62", "FLSA185", "FLCK216"],
    "gemi": ["FLCK18", "FLSF54", "FLWO6", "FLAB109"],
    "bran": ["BJSL25", "BJSB3", "BJVL19"],
    "fusi": ["MXED8", "MXGT4", "TXGR3", "TXMD3"],
    "sagr": ["CUVN10", "CUCA4", "CUSV6"],
    "oleo": ["CRL0030", "HNDA09", "BZBB1", "MXSA3017"],
}

# set a simple minmap requiring 1 sample from each group
minmap = {name: 0.75 for name in imap}


# In[21]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
    start=500000, 
    end=800000,
    mincov=0.8,
    imap=imap,
    minmap=minmap,
)

# show stats of the window
ext.stats


# ### Subsample taxa with `imap`
# 
# You can use an imap dictionary to select which samples to include/exclude from an analysis. This is an easy way to remove rogue taxa, hybrids, or technical replicates from phylogenetic analyses. Here I select a subset ot taxa to include in the analyses and keep only sites that have 80% coverage from scaffold 2 (Qrob_Chr03). 

# In[24]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=2,
    mincov=0.8,
    imap={
        "include": [
            "TXWV2", "LALC2", "SCCU3", "FLSF33", "FLBA140",
            "FLSF47", "FLMO62", "FLSA185", "FLCK216",
            "FLCK18", "FLSF54", "FLWO6", "FLAB109",
        ]
    },
)

# show stats of the window
ext.stats


# ### Concatenate multiple scaffolds together
# You can also concatenate multiple scaffolds together using `window_extracter`. This can be useful for creating genome-wide alignments, or smaller subsets of the genome. For example, you may want to combine multiple scaffolds from the same chromosome together, or, if you are working with denovo data, you could even combine a random sample of anonymous loci together as a sort of pseudo bootstrapping procedure. To select multiple scaffolds you simply provide a list or range of scaffold idxs. 

# In[25]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=[0, 1, 2, 3, 4, 5],
    mincov=0.5,
)

# show stats of the window
ext.stats


# ### Consensus reduction with `imap`
# 
# You can further reduce missing data by condensing data from multiple samples into a single "consensus" representative using the `consensus_reduce=True` option. This uses the `imap` dictionary to group samples into groups and sample the most frequent allele. This can be particularly useful for analyses in which you want dense species-level coverage with little missing data, but it is not particularly important which individual represents the sampled allele for a species at a given locus. For example, if you want to construct many gene trees with one representative per species to use as input to a two-step species tree inference program like ASTRAL. (Note: to automate `window_extracter` calls across many windows of the genome see the `treeslider` tool.)
# 
# In the example below you can see that the consensus option reduces the size of the dataset from 28 samples to 8 samples, and in doing so reduces the missing data from 19% to 1%.

# <img src="https://eaton-lab.org/slides/data-svg/consensus.svg">

# In[22]:


# select a scaffold idx, start, and end positions
ext = ipa.window_extracter(
    data=SEQSFILE,
    scaffold_idxs=0,
    start=200000, 
    end=5000000,
    mincov=0.8,
    imap=imap,
    minmap=minmap,
    consensus_reduce=True,   # <--- uses IMAP info to make consensus calls
)

# show stats of the window
ext.stats


# ### Write selected window to a file
# Once you've chosen the final set of arguments to select the window of interest you can write the alignment to *.phy* format by calling the `.run()` command. If you want to write to nexus format you can simply add the argument `nexus=True`. To change the name and location where file will be written you can set the `name` and `workdir` options on the window_extracter tool, otherwise it will use default options.

# In[24]:


ext.run(force=True)


# ### Accessing the output files
# 
# The output files created by the `.run()` command will be written to the working directory (defaults to "./analysis-window_extracter"). You can either find the full path to that file or access it easily from the extracter object itself as an attribute like below. 

# In[26]:


# path to the phylip file output
ext.outfile