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

# FASTA
# =====
# 
# This notebook briefly explores the [FASTA] format, a very common format for storing DNA sequences.  FASTA is the preferred format for storing *reference genomes*.
# 
# FASTA and FASTQ are rather similar, but FASTQ is almost always used for storing *sequencing reads* (with associated quality values), whereas FASTA is used for storing all kinds of DNA, RNA or protein sequencines (without associated quality values).
# 
# Before delving into the format, I should mention that the [BioPython] project makes parsing and using many file formats, including FASTA, quite easy.  See the BioPython [SeqIO] module in particular.  As far as I know, though, [SeqIO] does not use FASTA indexes, discussed toward the bottom, which is a disadvantage.
# 
# [FASTA]: http://en.wikipedia.org/wiki/FASTA_format
# [BioPython]: http://biopython.org/wiki/Main_Page
# [SeqIO]: http://biopython.org/wiki/SeqIO

# ### Basic format
# 
# Here is the basic format:
# 
#     >sequence1_short_name with optional additional info after whitespace
#     ACATCACCCCATAAACAAATAGGTTTGGTCCTAGCCTTTCTATTAGCTCTTAGTAAGATTACACATGCAA
#     GCATCCCCGTTCCAGTGAGTTCACCCTCTAAATCACCACGATCAAAAGGAACAAGCATCAAGCACGCAGC
#     AATGCAGCTCAAAACGCTTAGCCTAGCCACACCCCCACGGGAAACAGCAGTGAT
#     >sequence2_short_name with optional additional info after whitespace
#     GCCCCAAACCCACTCCACCTTACTACCAGACAACCTTAGCCAAACCATTTACCCAAATAAAGTATAGGCG
#     ATAGAAATTGAAACCTGGCGCAATAGATATAGTACCGCAAGGGAAAGATGAAAAATTATAACCAAGCATA
#     ATATAG

# A line starting with a `>` (greater-than) sign indicates the beginning of a new sequence and specifies its name.  Take the first line above.  Everything after the `>` up to and excluding the first whitespace character (`sequence1_short_name`), is the "short name."  Everything after the `>` up to the end of the *line* (`sequence1_short_name with optional additional info after whitespace`) is the "long name."  We usually use the short name when referring to FASTA sequences.
# 
# The next three lines consists of several nucleotides.  There is a maximum number of nucleotides permitted per line; in this case, it is 70.  If the sequence is longer then 70 nucleotides, it "wraps" down to the next line.  Not every FASTA file uses the same maximum, but a given FASTA file must use the same maximum throughout the file.
# 
# The sequences above are made up.  Here's a real-world reference sequence (the [human mitochondrial genome]) in FASTA format:
# 
# [human mitochondrial genome]: http://en.wikipedia.org/wiki/Human_mitochondrial_genetics

# In[1]:


import gzip
import urllib.request
url = 'ftp://ftp.ncbi.nlm.nih.gov/genomes/archive/old_genbank/Eukaryotes/vertebrates_mammals/Homo_sapiens/GRCh38/non-nuclear/assembled_chromosomes/FASTA/chrMT.fa.gz'
response = urllib.request.urlopen(url)
print(gzip.decompress(response.read()).decode('UTF8'))


# This FASTA file shown above has just one sequence in it.  As we saw in the first example above, it's also possible for one FASTA file to contain multiple sequences.  These are sometimes called multi-FASTA files.  When you write code to interpret FASTA files, it's a good idea to *always* allow for the possibility that the FASTA file might contain multiple sequences.
# 
# FASTA files are often stored with the `.fa` file name extension, but this is not a rule.  `.fasta` is another popular extenson.  You may also see `.fas`, `.fna`, `.mfa` (for multi-FASTA), and others.

# ### Parsing FASTA
# 
# Here is a simple function for parsing a FASTA file into a Python dictionary.  The dictionary maps short names to corresponding nucleotide strings (with whitespace removed).

# In[2]:


def parse_fasta(fh):
    fa = {}
    current_short_name = None
    # Part 1: compile list of lines per sequence
    for ln in fh:
        if ln[0] == '>':
            # new name line; remember current sequence's short name
            long_name = ln[1:].rstrip()
            current_short_name = long_name.split()[0]
            fa[current_short_name] = []
        else:
            # append nucleotides to current sequence
            fa[current_short_name].append(ln.rstrip())
    # Part 2: join lists into strings
    for short_name, nuc_list in fa.items():
        # join this sequence's lines into one long string
        fa[short_name] = ''.join(nuc_list)
    return fa


# The first part accumulates a list of strings (one per line) for each sequence.  The second part joins those lines together so that we end up with one long string per sequence.  Why divide it up this way?  Mainly to avoid the [poor performance](http://www.skymind.com/~ocrow/python_string/) of repeatedly concatenating (immutable) Python strings.
# 
# I'll test it by running it on the simple multi-FASTA file we saw before:

# In[3]:


from io import StringIO
fasta_example = StringIO(
'''>sequence1_short_name with optional additional info after whitespace
ACATCACCCCATAAACAAATAGGTTTGGTCCTAGCCTTTCTATTAGCTCTTAGTAAGATTACACATGCAA
GCATCCCCGTTCCAGTGAGTTCACCCTCTAAATCACCACGATCAAAAGGAACAAGCATCAAGCACGCAGC
AATGCAGCTCAAAACGCTTAGCCTAGCCACACCCCCACGGGAAACAGCAGTGAT
>sequence2_short_name with optional additional info after whitespace
GCCCCAAACCCACTCCACCTTACTACCAGACAACCTTAGCCAAACCATTTACCCAAATAAAGTATAGGCG
ATAGAAATTGAAACCTGGCGCAATAGATATAGTACCGCAAGGGAAAGATGAAAAATTATAACCAAGCATA
ATATAG''')
parsed_fa = parse_fasta(fasta_example)
parsed_fa


# Note that only the short names survive.  This is usually fine, but it's not hard to modify the function so that information relating short names to long names is also retained.

# ### Indexed FASTA
# 
# Say you have one or more *big* FASTA files (e.g. the entire human reference genome) and you'd like to access those files "randomly," peeking at substrings here and there without any regular access pattern.  Maybe you're mimicking a sequencing machine, reading snippets of DNA here and there.
# 
# You could start by using the `parse_fasta` function defined above to parse the FASTA files.  Then, to access a substring, do as follows:

# In[4]:


parsed_fa['sequence2_short_name'][100:130]


# Accessing a substring in this way is very fast and simple.  The downside is that you've stored all of the sequences in memory.  If the FASTA files are really big, this takes lots of valuable memory.  This may or may not be a good trade.
# 
# An alternative is to load only the *portions* of the FASTA files that you need, when you need them.  For this to be practical, we have to have a way of "jumping" to the specific part of the specific FASTA file that you're intersted in.
# 
# Fortunately, there is a standard way of indexing a FASTA file, popularized by the `faidx` tool in [SAMtools].  When you have such an index, it's easy to calculate exactly where to jump to when you want to extract a specific substring.  Here is some Python to create such an index:
# 
# [SAMtools]: http://samtools.sourceforge.net

# In[5]:


def index_fasta(fh):
    index = []
    current_short_name = None
    current_byte_offset, running_seq_length, running_byte_offset = 0, 0, 0
    line_length_including_ws, line_length_excluding_ws = 0, 0
    for ln in fh:
        ln_stripped = ln.rstrip()
        running_byte_offset += len(ln)
        if ln[0] == '>':
            if current_short_name is not None:
                index.append((current_short_name, running_seq_length,
                              current_byte_offset, line_length_excluding_ws,
                              line_length_including_ws))
            long_name = ln_stripped[1:]
            current_short_name = long_name.split()[0]
            current_byte_offset = running_byte_offset
            running_seq_length = 0
        else:
            line_length_including_ws = max(line_length_including_ws, len(ln))
            line_length_excluding_ws = max(line_length_excluding_ws, len(ln_stripped))
            running_seq_length += len(ln_stripped)
    if current_short_name is not None:
        index.append((current_short_name, running_seq_length,
                      current_byte_offset, line_length_excluding_ws,
                      line_length_including_ws))
    return index


# Here we use it to index a small multi-FASTA file.  We print out the index at the end.

# In[6]:


fasta_example = StringIO(
'''>sequence1_short_name with optional additional info after whitespace
ACATCACCCCATAAACAAATAGGTTTGGTCCTAGCCTTTCTATTAGCTCTTAGTAAGATTACACATGCAA
GCATCCCCGTTCCAGTGAGTTCACCCTCTAAATCACCACGATCAAAAGGAACAAGCATCAAGCACGCAGC
AATGCAGCTCAAAACGCTTAGCCTAGCCACACCCCCACGGGAAACAGCAGTGAT
>sequence2_short_name with optional additional info after whitespace
GCCCCAAACCCACTCCACCTTACTACCAGACAACCTTAGCCAAACCATTTACCCAAATAAAGTATAGGCG
ATAGAAATTGAAACCTGGCGCAATAGATATAGTACCGCAAGGGAAAGATGAAAAATTATAACCAAGCATA
ATATAG''')
idx = index_fasta(fasta_example)
idx


# What do the fields in those two records mean?  Take the first record: `('sequence1_short_name', 194, 69, 70, 71)`.  The fields from left to right are (1) the short name, (2) the length (in nucleotides), (3) the byte offset in the FASTA file of the first nucleotide of the sequence, (4) the maximum number of nucleotides per line, and (5) the maximum number of *bytes* per line, including whitespace.  It's not hard to convince yourself that, if you know all these things, it's not hard to figure out the byte offset of any position in any of the sequences.  (This is what the `get` member of the `FastaIndexed` class defined below does.)
# 
# A typical way to build a FASTA index like this is to use [SAMtools], specifically the `samtools faidx` command.  This and all the other `samtools` commands are documented in [its manual](http://samtools.sourceforge.net/samtools.shtml).
# 
# [SAMtools]: http://samtools.sourceforge.net
# 
# When you use a tool like this to index a FASTA file, a new file containing the index is written with an additional `.fai` extension.  E.g. if the FASTA file is named `hg19.fa`, then running `samtools faidx hg19.fa` will create a new file `hg19.fa.fai` containing the index.
# 
# The following Python class shows how you might use the FASTA file together with its index to extract arbitrary substrings without loading all of the sequences into memory:

# In[7]:


import re

class FastaOOB(Exception):
    """ Out-of-bounds exception for FASTA sequences """
    
    def __init__(self, value):
        self.value = value
    
    def __str__(self):
        return repr(self.value)

class FastaIndexed(object):
    """ Encapsulates a set of indexed FASTA files.  Does not load the FASTA
        files into memory but still allows the user to extract arbitrary
        substrings, with the help of the index. """
    
    __removeWs = re.compile(r'\s+')
    
    def __init__(self, fafns):
        self.fafhs = {}
        self.faidxs = {}
        self.chr2fh = {}
        self.offset = {}
        self.lens = {}
        self.charsPerLine = {}
        self.bytesPerLine = {}
        
        for fafn in fafns:
            # Open FASTA file
            self.fafhs[fafn] = fh = open(fafn, 'r')
            # Parse corresponding .fai file
            with open(fafn + '.fai') as idxfh:
                for ln in idxfh:
                    toks = ln.rstrip().split()
                    if len(toks) == 0:
                        continue
                    assert len(toks) == 5
                    # Parse and save the index line
                    chr, ln, offset, charsPerLine, bytesPerLine = toks
                    self.chr2fh[chr] = fh
                    self.offset[chr] = int(offset) # 0-based
                    self.lens[chr] = int(ln)
                    self.charsPerLine[chr] = int(charsPerLine)
                    self.bytesPerLine[chr] = int(bytesPerLine)
    
    def __enter__(self):
        return self
    
    def __exit__(self, type, value, traceback):
        # Close all the open FASTA files
        for fafh in self.fafhs.values():
            fafh.close()
    
    def has_name(self, refid):
        return refid in self.offset
    
    def name_iter(self):
        return self.offset.iterkeys()
    
    def length_of_ref(self, refid):
        return self.lens[refid]
    
    def get(self, refid, start, ln):
        ''' Return the specified substring of the reference. '''
        assert refid in self.offset
        if start + ln > self.lens[refid]:
            raise ReferenceOOB('"%s" has length %d; tried to get [%d, %d)' % (refid, self.lens[refid], start, start + ln))
        fh, offset, charsPerLine, bytesPerLine = \
            self.chr2fh[refid], self.offset[refid], \
            self.charsPerLine[refid], self.bytesPerLine[refid]
        byteOff = offset
        byteOff += (start // charsPerLine) * bytesPerLine
        into = start % charsPerLine
        byteOff += into
        fh.seek(byteOff)
        left = charsPerLine - into
        # Count the number of line breaks interrupting the rest of the
        # string we're trying to read
        if ln < left:
            return fh.read(ln)
        else:
            nbreaks = 1 + (ln - left) // charsPerLine
            res = fh.read(ln + nbreaks * (bytesPerLine - charsPerLine))
            res = re.sub(self.__removeWs, '', res)
        return res


# Here's an example of how to use the class defined above.

# In[8]:


# first we'll write a new FASTA file
with open('tmp.fa', 'w') as fh:
    fh.write('''>sequence1_short_name with optional additional info after whitespace
ACATCACCCCATAAACAAATAGGTTTGGTCCTAGCCTTTCTATTAGCTCTTAGTAAGATTACACATGCAA
GCATCCCCGTTCCAGTGAGTTCACCCTCTAAATCACCACGATCAAAAGGAACAAGCATCAAGCACGCAGC
AATGCAGCTCAAAACGCTTAGCCTAGCCACACCCCCACGGGAAACAGCAGTGAT
>sequence2_short_name with optional additional info after whitespace
GCCCCAAACCCACTCCACCTTACTACCAGACAACCTTAGCCAAACCATTTACCCAAATAAAGTATAGGCG
ATAGAAATTGAAACCTGGCGCAATAGATATAGTACCGCAAGGGAAAGATGAAAAATTATAACCAAGCATA
ATATAG''')
with open('tmp.fa') as fh:
    idx = index_fasta(fh)
with open('tmp.fa.fai', 'w') as fh:
    fh.write('\n'.join(['\t'.join(map(str, x)) for x in idx]))
with FastaIndexed(['tmp.fa']) as fa_idx:
    print(fa_idx.get('sequence2_short_name', 100, 30))


# ### Other resources
# 
# * [Wikipedia page for FASTA format](http://en.wikipedia.org/wiki/Fasta_format)
# * The [original FASTA paper](http://www.sciencedirect.com/science/article/pii/007668799083007V) by Bill Pearson.  This is the software tool that made the format popular.
# * [BioPython], which has [its own ways of parsing FASTA](http://biopython.org/wiki/SeqIO)
# * [Many](https://github.com/brentp/pyfasta) [other](https://github.com/lh3/seqtk) [libraries](https://github.com/mdshw5/pyfaidx) and [tools]((http://hannonlab.cshl.edu/fastx_toolkit/)
# 
# [BioPython]: http://biopython.org/wiki/Main_Page
# [SeqIO]: http://biopython.org/wiki/SeqIO