<filename>SLpackage/private/pacbio/pythonpkgs/pbreports/lib/python2.7/site-packages/pbreports/report/summarize_consensus.py
#! python

"""
Augment the alignment_summary.gff file with consensus and variants information.
"""
from __future__ import absolute_import, division, print_function

from collections import namedtuple, defaultdict
import argparse
import logging
import bisect
import json
import gzip
import sys

import numpy as np
import ast
import math
import os.path
import itertools

from pbcommand.utils import setup_log
from pbcommand.cli import pbparser_runner, get_default_argparser
from pbcommand.models import FileTypes, get_pbparser
from pbcommand.common_options import add_resolved_tool_contract_option
from pbcore.io import GffReader, GffWriter, Gff3Record
from pbreports.__init__ import VERSION


class DieException(Exception):
    """By deriving from BaseException instead of Exception,
    this will not be trapped by normal "except Exception:" blocks.
    But it will still not trigger a system-exit, which could
    subvert the logging clean-up in pbcommand.
    SE-1012
    """

def die(msg):
    print(msg, file=sys.stderr)
    raise DieException(msg)
    #sys.exit(-1)

class CommonEqualityMixin(object):
    def __eq__(self, other):
        return (isinstance(other, self.__class__)
            and self.__dict__ == other.__dict__)

    def __ne__(self, other):
        return not self.__eq__(other)


# An exception for incompatible cmp.h5 files
class IncompatibleDataException(Exception):
    pass

# We truncate QVs at 93 because the FASTQ format downstream can only
# support QVs in the range [0, 93] without lossage.

def error_probability_to_qv(error_probability, cap=93):
    """
    Convert an error probability to a phred-scaled QV.
    """
    if error_probability==0:
        return cap
    else:
        return min(cap, int(round(-10*math.log10(error_probability))))


_complement = { "A" : "T",
                "C" : "G",
                "G" : "C",
                "T" : "A",
                "-" : "-" }

def complement(s):
    cStr = "".join(_complement[c] for c in s)
    if type(s) == str:
        return cStr
    else:
        return np.fromstring(cStr, "S1")

def reverseComplement(s):
    return complement(s)[::-1]

def fileFormat(filename):
    if filename.endswith(".gz"):
        ext = os.path.splitext(filename[:-3])[1]
    else:
        ext = os.path.splitext(filename)[1]
    ext = ext.lower()
    if   ext in [".fa", ".fasta"]: return "FASTA"
    elif ext in [".fq", ".fastq"]: return "FASTQ"
    elif ext in [".gff" ]:         return "GFF"
    elif ext in [".vcf" ]:         return "VCF"
    elif ext in [".csv" ]:         return "CSV"
    else:
        raise ValueError("Unrecognized file format for {f}".format(f=filename))

def rowNumberIsInReadStratum(readStratum, rowNumber):
    n, N = readStratum
    return (rowNumber % N) == n

def readsInWindow(alnFile, window, depthLimit=None,
                  minMapQV=0, strategy="fileorder",
                  stratum=None, barcode=None):
    """
    Return up to `depthLimit` reads (as row numbers integers) where
    the mapped reference intersects the window.  If depthLimit is None,
    return all the reads meeting the criteria.

    `strategy` can be:
      - "longest" --- get the reads with the longest length in the window
      - "spanning" --- get only the reads spanning the window
      - "fileorder" --- get the reads in file order
    """
    assert strategy in {"longest", "spanning", "fileorder",
                        "long-and-strand-balanced"}

    if stratum is not None:
        raise ValueError("stratum needs to be reimplemented")

    def depthCap(iter):
        if depthLimit is not None:
            return alnFile[list(itertools.islice(iter, 0, depthLimit))]
        else:
            return alnFile[list(iter)]

    def lengthInWindow(hit):
        return (min(alnFile.index.tEnd[hit], winEnd) -
                max(alnFile.index.tStart[hit], winStart))

    winId, winStart, winEnd = window
    alnHits = np.array(list(alnFile.readsInRange(winId, winStart, winEnd,
                                                 justIndices=True)))
    if len(alnHits) == 0:
        return []

    if barcode == None:
        alnHits = alnHits[alnFile.mapQV[alnHits] >= minMapQV]
    else:
        # this wont work with CmpH5 (no bc in index):
        barcode = ast.literal_eval(barcode)
        alnHits = alnHits[(alnFile.mapQV[alnHits] >= minMapQV) &
                          (alnFile.index.bcLeft[alnHits] == barcode[0]) &
                          (alnFile.index.bcRight[alnHits] == barcode[1])]

    if strategy == "fileorder":
        return depthCap(alnHits)
    elif strategy == "spanning":
        winLen = winEnd - winStart
        return depthCap( hit for hit in alnHits
                         if lengthInWindow(hit) == winLen )
    elif strategy == "longest":
        return depthCap(sorted(alnHits, key=lengthInWindow, reverse=True))
    elif strategy == "long-and-strand-balanced":
        # Longest (in window) is great, but bam sorts by tStart then strand.
        # With high coverage, this bias resulted in variants. Here we lexsort
        # by tStart and tEnd. Longest in window is the final criteria in
        # either case.

        # lexical sort:
        ends = alnFile.index.tEnd[alnHits]
        starts = alnFile.index.tStart[alnHits]
        lex_sort = np.lexsort((ends, starts))

        # reorder based on sort:
        sorted_ends = ends[lex_sort]
        sorted_starts = starts[lex_sort]
        sorted_alnHits = alnHits[lex_sort]

        # get lengths in window:
        post = sorted_ends > winEnd
        sorted_ends[post] = winEnd
        pre = sorted_starts < winStart
        sorted_starts[pre] = winStart
        lens = sorted_ends - sorted_starts

        # coerce a descending sort:
        win_sort = ((winEnd - winStart) - lens).argsort(kind="mergesort")
        return depthCap(sorted_alnHits[win_sort])


def datasetCountExceedsThreshold(alnFile, threshold):
    """
    Does the file contain more than `threshold` datasets?  This
    impacts whether or not we should disable the chunk cache.
    """
    total = 0
    for i in np.unique(alnFile.AlnGroupID):
        total += len(alnFile._alignmentGroup(i))
        if total > threshold:
            return True
    return False

#
# Some lisp functions we want
#
fst   = lambda t: t[0]
snd   = lambda t: t[1]
third = lambda t: t[2]

def nub(it):
    """
    Unique entries in an iterable, preserving order
    """
    seen = set()
    for x in it:
        if x not in seen: yield(x)
        seen.add(x)

#
# Note: GFF-style coordinates
#
Region = namedtuple("Region", ("seqid", "start", "end"))

log = logging.getLogger(__name__)
__VERSION__ = '.'.join([str(v) for v in VERSION])

class Constants(object):
    TOOL_ID = "pbreports.tasks.summarize_consensus"
    DRIVER_EXE = "python -m pbreports.report.summarize_consensus --resolved-tool-contract "


def get_contract_parser():
    p = get_pbparser(
        Constants.TOOL_ID,
        __VERSION__,
        "Summarize Consensus",
        __doc__,
        Constants.DRIVER_EXE,
        default_level="ERROR")
    p.add_input_file_type(FileTypes.GFF, "alignment_summary",
        "Alignment summary GFF", "Alignment summary GFF file")
    p.tool_contract_parser.add_input_file_type(FileTypes.GFF, "variants",
        "Variants GFF", "Variants GFF file")
    p.arg_parser.parser.add_argument("--variantsGff",
        type=str,
        help="Input variants.gff or variants.gff.gz filename",
        required=True)
    p.tool_contract_parser.add_output_file_type(FileTypes.GFF, "output",
        name="Coverage and Variant Call Summary",
        description="Coverage and variant call summary for regions (bins) spanning the reference",
        default_name="alignment_summary_variants")
    p.arg_parser.parser.add_argument("-o", "--output",
        type=str,
        help="Output alignment_summary.gff filename")
    return p

def get_args_from_resolved_tool_contract(resolved_tool_contract):
    rtc = resolved_tool_contract
    p = get_contract_parser().arg_parser.parser
    args = [
        rtc.task.input_files[0],
        "--variantsGff", rtc.task.input_files[1],
        "--output", rtc.task.output_files[0],
    ]
    return p.parse_args(args)

def run(options):
    headers = [
        ("source", "pbreports summarize_consensus {}".format(__VERSION__)),
        ("pacbio-alignment-summary-version", "0.6"),
        ("source-commandline", " ".join(sys.argv)),
        ]

    inputVariantsGff = GffReader(options.variantsGff)
    inputAlignmentSummaryGff = GffReader(options.alignment_summary)

    summaries = {}
    for gffRecord in inputAlignmentSummaryGff:
        region = Region(gffRecord.seqid, gffRecord.start, gffRecord.end)
        summaries[region] = { "ins" : 0,
                              "del" : 0,
                              "sub" : 0,
                              # TODO: base consensusQV on effective coverage
                              "cQv" : (20, 20, 20)
                             }
    inputAlignmentSummaryGff.close()

    counterNames = { "insertion"    : "ins",
                     "deletion"     : "del",
                     "substitution" : "sub" }
    regions_by_contig = defaultdict(list)
    for region in summaries:
        regions_by_contig[region.seqid].append(region)
    for seqid in regions_by_contig.keys():
        r = regions_by_contig[seqid]
        regions_by_contig[seqid] = sorted(r, lambda a,b: cmp(a.start, b.start))
    logging.info("Processing variant records")
    i = 0
    have_contigs = set(regions_by_contig.keys())
    for variantGffRecord in inputVariantsGff:
        if not variantGffRecord.seqid in have_contigs:
            raise KeyError(
                "Can't find alignment summary for contig '{s}".format(
                s=variantGffRecord.seqid))
        positions = [r.start for r in regions_by_contig[variantGffRecord.seqid]]
        idx = bisect.bisect_right(positions, variantGffRecord.start) - 1
        # XXX we have to be a little careful here - an insertion at the start
        # of a contig will have start=0 versus start=1 for the first region
        if idx < 0:
            idx = 0
        region = regions_by_contig[variantGffRecord.seqid][idx]
        assert ((region.start <= variantGffRecord.start <= region.end) or
                (region.start == 1 and variantGffRecord.start == 0 and
                 variantGffRecord.type == "insertion")), \
            (variantGffRecord.seqid, region.start, variantGffRecord.start,
             region.end, variantGffRecord.type, idx)
        summary = summaries[region]
        counterName = counterNames[variantGffRecord.type]
        variantLength = max(len(variantGffRecord.reference),
                            len(variantGffRecord.variantSeq))
        summary[counterName] += variantLength
        i += 1
        if i % 1000 == 0:
            logging.info("{i} records...".format(i=i))


    inputAlignmentSummaryGff = open(options.alignment_summary)
    outputAlignmentSummaryGff = open(options.output, "w")

    inHeader = True

    for line in inputAlignmentSummaryGff:
        line = line.rstrip()

        # Pass any metadata line straight through
        if line[0] == "#":
            print(line.strip(), file=outputAlignmentSummaryGff)
            continue

        if inHeader:
            # We are at the end of the header -- write the tool-specific headers
            for k, v in headers:
                print("##%s %s" % (k, v), file=outputAlignmentSummaryGff)
            inHeader = False

        # Parse the line
        rec = Gff3Record.fromString(line)

        if rec.type == "region":
            summary = summaries[(rec.seqid, rec.start, rec.end)]
            if "cQv" in summary:
                cQvTuple = summary["cQv"]
                line += ";%s=%s" % ("cQv", ",".join(str(int(f)) for f in cQvTuple))
            for counterName in counterNames.values():
                if counterName in summary:
                    line += ";%s=%d" % (counterName, summary[counterName])
            print(line, file=outputAlignmentSummaryGff)
    return 0


def args_runner(args):
    return run(options=args)


def resolved_tool_contract_runner(resolved_tool_contract):
    args = get_args_from_resolved_tool_contract(resolved_tool_contract)
    return run(options=args)


def main(argv=sys.argv):
    return pbparser_runner(argv[1:],
                           get_contract_parser(),
                           args_runner,
                           resolved_tool_contract_runner,
                           log,
                           setup_log)

if __name__ == "__main__":
    sys.exit(main())
