#!/usr/bin/env python3

import argparse
import sys
import os
myBinDir = os.path.normpath(os.path.dirname(sys.argv[0]))
sys.path.insert(0, os.path.join(myBinDir, "../lib"))
from collections import namedtuple, defaultdict
from pycbio.hgdata.bed import Bed
from pycbio.tsv import TsvReader
from cenSat.bedColors import getHg38UcscChromColor

# naive chaining didn't work very well
doChaining = False

def parseArgs():
    usage = """convert output of MUMmer show-coords -T to BED
"""
    parser = argparse.ArgumentParser(description=usage)
    parser.add_argument('--chromAliasTsv', default=None,
                        help="map sequence names using this file")
    parser.add_argument('targetSeq', choices=("seq1", "seq2"),
                        help="target sequence")
    parser.add_argument('inCoords', nargs='?', default=sys.stdin,
                        type=argparse.FileType('r'),
                        help="input coordinates file")
    parser.add_argument('outBed', nargs='?', default=sys.stdout,
                        type=argparse.FileType('w'),
                        help="output BED file")
    return parser.parse_args()

def loadChromAliases(chromAliasTsv):
    return {r.alias: r for r in TsvReader(chromAliasTsv)}

class NumcerCoords(namedtuple("NumcerCoords", ("s1", "e1", "s2", "e2", "len1", "len2",
                                               "perId", "seq1", "seq2", "rev1", "rev2"))):
    def __new__(cls, s1, e1, s2, e2, len1, len2, perId, seq1, seq2, rev1, rev2):
        return super(NumcerCoords, cls).__new__(cls, s1, e1, s2, e2, len1, len2, perId, seq1, seq2, rev1, rev2)

    def swap(self):
        return NumcerCoords(self.s2, self.e2, self.s1, self.e1, self.len2, self.len1, self.perId, self.seq2, self.seq1, self.rev2, self.rev1)

def mkCoordsRight(st, en):
    if st > en:
        return (en, st, True)
    else:
        return (st, en, False)

def parseNumcerCoords(line):
    cols = line.split('\t')
    if len(cols) != 9:
        raise Exception("expected 9 columns: " + line)
    s1, e1, s2, e2, len1, len2 = [int(c) for c in cols[0:6]]
    s1, e1, rev1 = mkCoordsRight(s1, e1)
    s2, e2, rev2 = mkCoordsRight(s2, e2)
    return NumcerCoords(s1 - 1, e1, s2 - 1, e2, len1, len2, float(cols[6]), cols[7], cols[8], rev1, rev2)

def loadNumcerCoords(fh, swapSides):
    # [S1]	[E1]	[S2]	[E2]	[LEN 1]	[LEN 2]	[% IDY]	[TAGS]
    for line in fh:
        if line.startswith("[S1]"):
            break
    coords = []
    for line in fh:
        c = parseNumcerCoords(line[0:-1])
        if swapSides:
            c = c.swap()
        coords.append(c)
    return coords

def partitionCoords(coords):
    "split by (target, query) pairs"
    coordParts = defaultdict(list)
    for c in coords:
        coordParts[c.seq1, c.seq2].append(c)
    return coordParts

def popNextChain(coords):
    prev = coords.pop(0)
    chained = [prev]
    i = 0
    while i < len(coords):
        c = coords[i]
        if (c.s1 > prev.e1) and (c.rev1 == prev.rev1) and (c.s2 > prev.e2) and (c.rev2 == prev.rev2):
            prev = coords.pop(i)
            chained.append(prev)
        else:
            i += 1
    return chained

def chainCoords(coords):
    "combine into list that can be chained"
    coords = sorted(coords, key=lambda c: (c.s1, c.rev1, c.s2, c.rev2, -c.len1))
    chaineds = []
    while len(coords) > 0:
        chaineds.append(popNextChain(coords))
    return chaineds

def chainCoordParts(coordParts):
    chained = []
    for coordPart in coordParts.values():
        chained.extend(chainCoords(coordPart))
    return chained

def transAlias(chromAliases, seq):
    a = chromAliases.get(seq)
    return a.chrom if a is not None else seq

def getColor(seq1, seq2):
    rgb = getHg38UcscChromColor(seq1, useDefault=False)
    if rgb is None:
        rgb = getHg38UcscChromColor(seq2)
    return rgb

def coordToBed(coordsChain, chromAliases):
    coord0 = coordsChain[0]
    seq1 = transAlias(chromAliases, coord0.seq1)
    seq2 = transAlias(chromAliases, coord0.seq2)
    strand = '-' if (coord0.rev1 ^ coord0.rev2) else '+'
    s1 = coord0.s1
    e1 = coordsChain[-1].e1
    s2 = coord0.s2
    e2 = coordsChain[-1].e2
    blocks = [Bed.Block(c.s1, c.e1) for c in coordsChain]
    return Bed(seq1, s1, e1, "{}:{}-{}".format(seq2, s2, e2), 0,
               strand, s1, e1, getColor(seq1, seq2), blocks)

def main(args):
    chromAliases = loadChromAliases(args.chromAliasTsv) if args.chromAliasTsv is not None else {}
    coords = loadNumcerCoords(args.inCoords, (args.targetSeq == "seq1"))
    if doChaining:
        # naive chaining doesn;'t work very well
        coordParts = partitionCoords(coords)
        coordsChains = chainCoordParts(coordParts)
    else:
        coordsChains = [[c] for c in coords]
    for coordsChain in coordsChains:
        coordToBed(coordsChain, chromAliases).write(args.outBed)
    args.outBed.close()

main(parseArgs())
