"""
Classes to support "biological sequence" files.

:Author: Bob Harris (rsharris@bx.psu.edu)
"""
from six import Iterator

# DNA reverse complement table

DNA_COMP = "                                             -                  " \
           " TVGH  CD  M KN   YSA BWXR       tvgh  cd  m kn   ysa bwxr      " \
           "                                                                " \
           "                                                                "


class SeqFile(object):
    """
    A biological sequence is a sequence of bytes or characters.  Usually these
    represent DNA (A,C,G,T), proteins, or some variation of those.

    class attributes:

        file:    file object containing the sequence
        revcomp: whether gets from this sequence should be reverse-complemented
                 False => no reverse complement
                 True  => (same as "-5'")
                 "maf" => (same as "-5'")
                 "+5'" => minus strand is from plus strand's 5' end (same as "-3'")
                 "+3'" => minus strand is from plus strand's 3' end (same as "-5'")
                 "-5'" => minus strand is from its 5' end (as per MAF file format)
                 "-3'" => minus strand is from its 3' end (as per genome browser,
                          but with origin-zero)
        name:    usually a species and/or chromosome name (e.g. "mule.chr5");  if
                 the file contains a name, that overrides this one
        gap:     gap character that aligners should use for gaps in this sequence
    """

    def __init__(self, file=None, revcomp=False, name="", gap=None):

        self.file = file
        if revcomp:
            self.revcomp = "-5'"
        elif revcomp == "+3'":
            self.revcomp = "-5'"
        elif revcomp == "+5'":
            self.revcomp = "-3'"
        elif revcomp == "maf":
            self.revcomp = "-5'"
        else:
            self.revcomp = revcomp
        self.name = name
        if gap is None:
            self.gap = "-"
        else:
            self.gap = gap

        self.text = None  # (subclasses fill in text and
        self.length = 0  # length or they most override get())

    def close(self):
        assert (self.file is not None)
        self.file.close()
        self.file = None

    def extract_name(self, line):
        try:
            return line.split()[0]
        except Exception:
            return ""

    def set_text(self, text):
        self.text = text
        self.length = len(text)

    def __str__(self):
        text = ""
        if self.name is not None:
            text += self.name + " "
        text += self.get(0, self.length)
        return text

    def get(self, start, length):
        """
        Fetch subsequence starting at position `start` with length `length`.
        This method is picky about parameters, the requested interval must
        have non-negative length and fit entirely inside the NIB sequence,
        the returned string will contain exactly 'length' characters, or an
        AssertionError will be generated.
        """
        # Check parameters
        assert length >= 0, "Length must be non-negative (got %d)" % length
        assert start >= 0, "Start must be greater than 0 (got %d)" % start
        assert start + length <= self.length, \
            "Interval beyond end of sequence (%s..%s > %s)" % (start, start + length, self.length)
        # Fetch sequence and reverse complement if necesary
        if not self.revcomp:
            return self.raw_fetch(start, length)
        if self.revcomp == "-3'":
            return self.reverse_complement(self.raw_fetch(start, length))
        assert self.revcomp == "-5'", "unrecognized reverse complement scheme"
        start = self.length - (start+length)
        return self.reverse_complement(self.raw_fetch(start, length))

    def raw_fetch(self, start, length):
        return self.text[start:start+length]

    def reverse_complement(self, text):
        comp = [ch for ch in text.translate(DNA_COMP)]
        comp.reverse()
        return "".join(comp)


class SeqReader(Iterator):
    """Iterate over all sequences in a file in order"""

    def __init__(self, file, revcomp=False, name="", gap=None):
        self.file = file
        self.revcomp = revcomp
        self.name = name
        self.gap = gap
        self.seqs_read = 0

    def close(self):
        self.file.close()

    def __iter__(self):
        return SeqReaderIter(self)

    def __next__(self):  # subclasses should override this method and return the
        return   # .. next sequence (of type SeqFile or a subclass) read from self.file


class SeqReaderIter(Iterator):
    def __init__(self, reader):
        self.reader = reader

    def __iter__(self):
        return self

    def __next__(self):
        v = next(self.reader)
        if not v:
            raise StopIteration
        return v
