"""
Support for reading and writing the LAV format produced by the `blastz`_
pairwise aligner.

.. _blastz: http://www.bx.psu.edu/miller_lab/
"""
from __future__ import print_function

import sys

from six import Iterator, StringIO

import bx.seq
from bx.align import (
    Alignment,
    Component,
    src_merge,
    src_split
)


class Reader(Iterator):
    """Iterate over all lav blocks in a file in order"""

    def __init__(self, file, path_subs=None, fail_to_ns=False):
        self.file = file
        self.lineNumber = 0
        self.path_subs = path_subs   # list of (prefix,replacement) to allow
        if self.path_subs is None:  # .. redirection of sequence file paths
            self.path_subs = []       # .. on different machines
        self.fail_to_ns = fail_to_ns  # True => if sequences fail to open, create a fake file of all Ns

        self.d_stanza_text = None

        self.seq1_filename = None
        self.seq1_file = None
        self.seq1_header = None
        self.seq1_start = None
        self.seq1_end = None
        self.seq1_strand = None
        self.seq1_contig = None
        self.seq1_src = None
        self.seq1_gap = None

        self.seq2_filename = None
        self.seq2_file = None
        self.seq2_header = None
        self.seq2_start = None
        self.seq2_end = None
        self.seq2_strand = None
        self.seq2_contig = None
        self.seq2_src = None
        self.seq2_gap = None

    def __next__(self):
        while (True):
            line = self.fetch_line(strip=None, requireLine=False)
            assert (line), "unexpected end of file (missing #:eof)"
            line = line.rstrip()
            if line == "":  # (allow blank lines between stanzas)
                continue
            if line == "#:eof":
                line = self.file.readline().rstrip()
                assert (not line), "extra line after #:eof (line %d, \"%s\")" % (self.lineNumber, line)
                return None
            if line == "#:lav":
                continue
            if line.startswith("d {"):
                self.d_stanza_text = self.parse_unknown_stanza()
                continue
            if line.startswith("s {"):
                self.parse_s_stanza()
                continue
            if line.startswith("h {"):
                self.parse_h_stanza()
                continue
            if line.startswith("a {"):
                (score, pieces) = self.parse_a_stanza()
                break
            if line.endswith("{"):
                self.parse_unknown_stanza()
                continue
            assert (False), "incomprehensible line (line %d, \"%s\")" % (self.lineNumber, line)
        return self.build_alignment(score, pieces)

    def __iter__(self):
        return ReaderIter(self)

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

    def open_seqs(self):
        if self.seq1_file is not None and self.seq2_file is not None:
            return

        if self.seq1_file is None:
            if self.seq1_strand == "+":
                revcomp = False
            else:
                revcomp = "-5'"
            if self.seq1_contig == 1:
                contig = None
            else:
                contig = self.seq1_contig
            try:
                f = open(self.seq1_filename, "rb")
            except Exception:
                if self.fail_to_ns:
                    f = StringIO(">seq1\n" + ("n" * (self.seq1_end - self.seq1_start)))
                    revcomp = False
                    contig = 1
                else:
                    assert (False), "failed to open %s" % self.seq1_filename
            self.seq1_file = bx.seq.seq_file(f, revcomp=revcomp, contig=contig)
            self.seq1_gap = self.seq1_file.gap
            try:
                name1 = self.header_to_src_name(self.seq1_header)
            except ValueError:
                try:
                    name1 = self.path_to_src_name(self.seq1_filename)
                except ValueError:
                    name1 = "seq1"
            (species1, chrom1) = src_split(name1)
            self.seq1_src = src_merge(species1, chrom1, contig)
            if contig is not None:
                chrom1 += "[%s]" % contig

        if self.seq2_file is None:
            if self.seq2_strand == "+":
                revcomp = False
            else:
                revcomp = "-5'"
            if self.seq2_contig == 1:
                contig = None
            else:
                contig = self.seq2_contig
            try:
                f = open(self.seq2_filename, "rb")
            except Exception:
                if self.fail_to_ns:
                    f = StringIO(">seq2\n" + ("n" * (self.seq2_end - self.seq2_start)))
                    revcomp = False
                    contig = 1
                else:
                    assert (False), "failed to open %s" % self.seq1_filename
            self.seq2_file = bx.seq.seq_file(f, revcomp=revcomp, contig=contig)
            self.seq2_gap = self.seq2_file.gap
            try:
                name2 = self.header_to_src_name(self.seq2_header)
            except ValueError:
                try:
                    name2 = self.path_to_src_name(self.seq2_filename)
                except ValueError:
                    name2 = "seq2"
            (species2, chrom2) = src_split(name2)
            self.seq2_src = src_merge(species2, chrom2, contig)
            if contig is not None:
                chrom2 += "[%s]" % contig

        length1 = self.seq1_file.length
        length2 = self.seq2_file.length
        assert (species1 != species2) or (chrom1 != chrom2) or (length1 == length2), \
            "conflicting lengths for %s (%d and %d)" % (self.seq1_src, length1, length2)

        self.species_to_lengths = {}
        self.species_to_lengths[species1] = {}
        self.species_to_lengths[species2] = {}  # (OK if it clobbers line above)
        self.species_to_lengths[species1][chrom1] = self.seq1_file.length
        self.species_to_lengths[species2][chrom2] = self.seq2_file.length

    def close_seqs(self):
        if self.seq1_file is not None:
            self.seq1_file.close()
            self.seq1_file = None
        if self.seq2_file is not None:
            self.seq2_file.close()
            self.seq2_file = None

    def parse_s_stanza(self):
        self.close_seqs()
        line = self.fetch_line(report=" in s-stanza")
        (self.seq1_filename,
         self.seq1_start,
         self.seq1_end,
         self.seq1_strand,
         self.seq1_contig) = self.parse_s_seq(line)

        line = self.fetch_line(report=" in s-stanza")
        (self.seq2_filename,
         self.seq2_start,
         self.seq2_end,
         self.seq2_strand,
         self.seq2_contig) = self.parse_s_seq(line)

        line = self.fetch_line(report=" in s-stanza")
        assert (line == "}"), "improper s-stanza terminator (line %d, \"%s\")" \
            % (self.lineNumber, line)

    def parse_s_seq(self, line):
        fields = line.split()
        filename = fields[0].strip('"')
        start = int(fields[1]) - 1
        end = int(fields[2])
        contig = int(fields[4])
        if fields[3] == "1":
            strand = "-"
        else:
            strand = "+"
        if filename.endswith("-"):
            assert (strand == "-"), "strand mismatch in \"%s\"" % line
            filename = filename[:-1]
        filename = do_path_subs(filename, self.path_subs)
        return (filename, start, end, strand, contig)

    def parse_h_stanza(self):
        line = self.fetch_line(strip='"', report=" in h-stanza")
        self.seq1_header = line
        self.seq1_header_prefix = ""
        if line.startswith(">"):
            self.seq1_header = line[1:].strip()
            self.seq1_header_prefix = ">"
        self.seq1_header = self.seq1_header.split(None, 1)
        if len(self.seq1_header) > 0:
            self.seq1_header = self.seq1_header[0]
        else:
            self.seq1_header = "seq1"

        line = self.fetch_line(strip='"', report=" in h-stanza")
        self.seq2_header = line
        self.seq2_header_prefix = ""
        if line.startswith(">"):
            self.seq2_header = line[1:].strip()
            self.seq2_header_prefix = ">"
        self.seq2_header = self.seq2_header.split(None, 1)
        if len(self.seq2_header) > 0:
            self.seq2_header = self.seq2_header[0]
        else:
            self.seq2_header = "seq2"

        line = self.fetch_line(report=" in h-stanza")
        assert (line == "}"), "improper h-stanza terminator (line %d, \"%s\")" \
            % (self.lineNumber, line)

    def parse_a_stanza(self):
        """returns the pair (score,pieces)
           where pieces is a list of ungapped segments (start1,start2,length,pctId)
           with start1,start2 origin-0"""
        # 's' line -- score, 1 field
        line = self.fetch_line(report=" in a-stanza")
        fields = line.split()
        assert (fields[0] == "s"), "s line expected in a-stanza (line %d, \"%s\")" \
            % (self.lineNumber, line)
        try:
            score = int(fields[1])
        except ValueError:
            score = float(fields[1])

        # 'b' line -- begin positions in seqs, 2 fields
        line = self.fetch_line(report=" in a-stanza")
        fields = line.split()
        assert (fields[0] == "b"), "b line expected in a-stanza (line %d, \"%s\")" \
            % (self.lineNumber, line)

        # 'e' line -- end positions in seqs, 2 fields
        line = self.fetch_line(report=" in a-stanza")
        fields = line.split()
        assert (fields[0] == "e"), "e line expected in a-stanza (line %d, \"%s\")" \
            % (self.lineNumber, line)

        # 'l' lines
        pieces = []
        while (True):
            line = self.fetch_line(report=" in a-stanza")
            fields = line.split()
            if fields[0] != "l":
                break
            start1 = int(fields[1]) - 1
            start2 = int(fields[2]) - 1
            length = int(fields[3]) - start1
            length2 = int(fields[4]) - start2
            try:
                pctId = int(fields[5])
            except ValueError:
                pctId = float(fields[5])
            assert (length2 == length), "length mismatch in a-stanza"
            pieces.append((start1+self.seq1_start, start2+self.seq2_start, length, pctId))
        assert (line == "}"), "improper a-stanza terminator (line %d, \"%s\")" \
            % (self.lineNumber, line)
        return (score, pieces)

    def parse_unknown_stanza(self):
        lines = []
        while (True):
            line = self.fetch_line()
            assert (line), "unexpected end of file (missing #:eof)"
            if line == "}":
                break
            lines.append(line)
        return "  " + "\n  ".join(lines) + "\n"

    def fetch_line(self, strip=True, requireLine=True, report=""):
        if strip is None:
            line = self.file.readline()
        elif strip is True:
            line = self.file.readline().strip()
        else:
            line = self.file.readline().strip().strip(strip)
        self.lineNumber += 1
        if requireLine:
            assert (line), "unexpected blank line or end of file%s (line %d)" \
                % (report, self.lineNumber)
        return line

    def d_stanza(self):
        if self.d_stanza_text is None:
            return ""
        return "d {\n%s}" % self.d_stanza_text

    def s_stanza(self):
        if self.seq1_filename is None:
            return ""

        if self.seq1_strand == "-":
            seq1_strand = "1"
        else:
            seq1_strand = "0"
        if self.seq2_strand == "-":
            seq2_strand = "1"
        else:
            seq2_strand = "0"

        s = "  \"%s\" %d %d %s %d\n"\
            % (self.seq1_filename, self.seq2_start+1, self.seq1_end,
               seq1_strand, self.seq1_contig)
        s += "  \"%s\" %d %d %s %d\n"\
            % (self.seq2_filename, self.seq2_start+1, self.seq2_end,
               seq2_strand, self.seq2_contig)

        return "s {\n%s}" % s

    def h_stanza(self):
        if self.seq1_header is None:
            return ""
        s = "  \"%s%s\"\n" % (self.seq1_header_prefix, self.seq1_header)
        s += "  \"%s%s\"\n" % (self.seq2_header_prefix, self.seq2_header)
        return "h {\n%s}" % s

    def build_alignment(self, score, pieces):
        """converts a score and pieces to an alignment"""
        # build text
        self.open_seqs()
        text1 = text2 = ""
        end1 = end2 = None
        for (start1, start2, length, pctId) in pieces:
            if end1 is not None:
                if start1 == end1:  # insertion in sequence 2
                    text1 += self.seq1_gap * (start2-end2)
                    text2 += self.seq2_file.get(end2, start2-end2)
                else:  # insertion in sequence 1
                    text1 += self.seq1_file.get(end1, start1-end1)
                    text2 += self.seq2_gap * (start1-end1)

            text1 += self.seq1_file.get(start1, length)
            text2 += self.seq2_file.get(start2, length)
            end1 = start1 + length
            end2 = start2 + length
        # create alignment
        start1 = pieces[0][0]
        start2 = pieces[0][1]
        end1 = pieces[-1][0] + pieces[-1][2]
        end2 = pieces[-1][1] + pieces[-1][2]
        size1 = end1 - start1
        size2 = end2 - start2
        a = Alignment(score=score, species_to_lengths=self.species_to_lengths)
        # if (self.seq1_strand == "-"): start1 = self.seq1_file.length - end1
        a.add_component(Component(self.seq1_src, start1, size1, self.seq1_strand, text=text1))
        # if (self.seq2_strand == "-"): start2 = self.seq2_file.length - end2
        a.add_component(Component(self.seq2_src, start2, size2, self.seq2_strand, text=text2))
        return a

    def path_to_src_name(self, path_name):
        # converts, e.g. ".../hg18/seq/chr13.nib" to "hg18.chr13"
        if path_name is None or path_name == "":
            raise ValueError
        if path_name.endswith(".nib"):
            path_name = path_name[:-4]
        if path_name.endswith(".fa"):
            path_name = path_name[:-3]
        if path_name.endswith(".fasta"):
            path_name = path_name[:-6]
        slash = path_name.rfind("/")
        if slash == -1:
            return path_name
        name = path_name[slash+1:]
        path_name = path_name[:slash]
        if path_name.endswith("/seq"):
            path_name = path_name[:-4]
        slash = path_name.rfind("/")
        if slash != -1:
            path_name = path_name[slash+1:]
        return path_name + "." + name

    def header_to_src_name(self, header):
        # converts, e.g. "hg18.chr13:115404472-117281897" to "hg18.chr13"
        if header is None or header == "":
            raise ValueError
        colon = header.rfind(":")
        if colon != -1:
            header = header[:colon]
        if "/" in header:
            raise ValueError
        if header.count(".") == 0:
            return header
        header = header.split(".")
        if header[0] == "" or header[1] == "":
            raise ValueError
        return ".".join(header)


class ReaderIter(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


class LavAsPiecesReader(Reader):
    """Iterate over all lav blocks in a file in order, returning alignments
       as score and pieces, as returned by Reader.parse_a_stanza"""

    def build_alignment(self, score, pieces):
        return (score, pieces)


class Writer(object):

    # blockHash is a hash from (src1,strand1,src2,strand2) to a list of blocks;
    # the blocks are collected on each call to write(), but the actual writing
    # does not occur until close().

    def __init__(self, file, attributes={}):
        self.file = file
        self.fname1 = None
        self.fname2 = None
        self.block = 0
        self.blockHash = {}  # (see note above)

        if "name_format_1" in attributes:
            self.fname1 = attributes["name_format_1"]
        if "name_format_2" in attributes:
            self.fname2 = attributes["name_format_2"]

        if "d_stanza" in attributes:
            self.write_lav_marker()
            print("d {", file=self.file)
            print(attributes["d_stanza"], file=self.file)
            print("}", file=self.file)

    def write(self, alignment):
        if len(alignment.components) != 2:
            raise ValueError(
                "%d-component alignment is not compatible with lav" %
                len(alignment.components))

        c1 = alignment.components[0]
        c2 = alignment.components[1]

        key = (c1.src, c1.strand, c2.src, c2.strand)
        if key not in self.blockHash:
            self.blockHash[key] = []
        self.blockHash[key].append(alignment)
        self.block += 1

    def close(self):
        keys = [key for key in self.blockHash]
        keys = sort_keys_by_chrom(keys)
        for key in keys:
            (src1, strand1, src2, strand2) = key
            alignment = self.blockHash[key][0]
            self.src1 = src1
            self.strand1 = strand1
            self.length1 = alignment.src_size(src1)
            self.src2 = src2
            self.strand2 = strand2
            self.length2 = alignment.src_size(src2)
            self.write_s_stanza()
            self.write_h_stanza()
            for alignment in self.blockHash[key]:
                self.write_a_stanza(alignment)
        self.write_trailer()
        if self.file != sys.stdout:
            self.file.close()

    def write_s_stanza(self):
        self.write_lav_marker()
        (strand1, flag1) = minus_or_nothing(self.strand1)
        (strand2, flag2) = minus_or_nothing(self.strand2)
        fname1 = build_filename(self.fname1, self.src1)
        fname2 = build_filename(self.fname2, self.src2)
        print("s {", file=self.file)
        print("  \"%s%s\" 1 %d %d 1" % (fname1, strand1, self.length1, flag1), file=self.file)
        print("  \"%s%s\" 1 %d %d 1" % (fname2, strand2, self.length2, flag2), file=self.file)
        print("}", file=self.file)

    def write_h_stanza(self):
        strand1 = rc_or_nothing(self.strand1)
        strand2 = rc_or_nothing(self.strand2)
        print("h {", file=self.file)
        print("  \"> %s%s\"" % (self.src1, strand1), file=self.file)
        print("  \"> %s%s\"" % (self.src2, strand2), file=self.file)
        print("}", file=self.file)

    def write_a_stanza(self, alignment):
        c1 = alignment.components[0]
        pos1 = c1.start
        text1 = c1.text.upper()
        c2 = alignment.components[1]
        pos2 = c2.start
        text2 = c2.text.upper()

        # collect ungapped pieces

        pieces = []
        piece1 = None

        for ix in range(len(text1)):
            ch1 = text1[ix]
            ch2 = text2[ix]

            nonGap = (ch1 != "-") and (ch2 != "-")
            if nonGap:
                if piece1 is None:  # new piece starts
                    (piece1, piece2, idCount) = (pos1, pos2, 0)
                if ch1 == ch2:
                    idCount += 1
            elif piece1 is not None:  # new gap starts
                size = pos1 - piece1
                pctId = (200*idCount + size) / (2*size)
                pieces.append((piece1, piece2, size, pctId))
                piece1 = None

            if ch1 != "-":
                pos1 += 1
            if ch2 != "-":
                pos2 += 1

        if piece1 is not None:
            size = pos1 - piece1
            pctId = (200*idCount + size) / (2*size)
            pieces.append((piece1, piece2, size, pctId))

        # write the block

        (start1, start2, size, pctId) = pieces[-1]  # get end of final piece
        end1 = start1 + size
        end2 = start2 + size

        (start1, start2, size, pctId) = pieces[0]   # get start of first piece

        score = int(round(alignment.score))

        print("a {", file=self.file)
        print("  s %s" % score, file=self.file)
        print("  b %d %d" % (start1 + 1, start2 + 1), file=self.file)
        print("  e %d %d" % (end1, end2), file=self.file)
        for (start1, start2, size, pctId) in pieces:
            print("  l %d %d %d %d %d" % (start1 + 1, start2 + 1, start1 + size, start2 + size, pctId), file=self.file)
        print("}", file=self.file)

    def write_lav_marker(self):
        print("#:lav", file=self.file)

    def write_trailer(self):
        print("#:eof", file=self.file)


def sort_keys_by_chrom(keys):
    decorated = sorted((chrom_key(src1), strand1, chrom_key(src2), strand2, (src1, strand1, src2, strand2))
                       for (src1, strand1, src2, strand2) in keys)
    return [key for (src1, strand1, src2, strand2, key) in decorated]


def chrom_key(src):
    (species, chrom) = src_split(src)
    if chrom.startswith("chr"):
        chrom = chrom[3:]
    try:
        chrom = int(chrom)
    except ValueError:
        pass
    return chrom


def build_filename(fmt, src):
    if fmt is None:
        return src
    num = fmt.count("%s")
    if num == 0:
        return fmt
    (species, chrom) = src_split(src)
    if num == 1:
        return fmt % chrom
    return fmt % (species, chrom)


def minus_or_nothing(strand):
    if strand == "-":
        return ("-", 1)
    else:
        return ("", 0)


def rc_or_nothing(strand):
    if strand == "-":
        return " (reverse complement)"
    else:
        return ""


def do_path_subs(path, path_subs):
    for (prefix, replacement) in path_subs:
        if path.startswith(prefix):
            return replacement + path[len(prefix):]
    return path
