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

:Author: Bob Harris (rsharris@bx.psu.edu)
"""

# DNA reverse complement table

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


class SeqFile:
    """
    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
        ), f"Interval beyond end of sequence ({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):
        return text.translate(DNA_COMP)[::-1]


class SeqReader:
    """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 next sequence
        # (of type SeqFile or a subclass) read from self.file
        return


class SeqReaderIter:
    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