<commit_before>"""Examine callable regions following genome mapping of short reads.

Identifies callable analysis regions surrounded by larger regions lacking
aligned bases. This allows parallelization of smaller chromosome chunks
through post-processing and variant calling, with each sub-section
mapping handled separately.

Regions are split to try to maintain relative uniformity across the
genome and avoid extremes of large blocks or large numbers of
small blocks.
"""
import contextlib
import copy
from distutils.version import LooseVersion
import operator
import os
import subprocess
import sys

import numpy
import pybedtools
import pysam
import toolz as tz
import yaml

from bcbio import bam, broad, utils
from bcbio.bam import ref
from bcbio.log import logger
from bcbio.distributed import multi, prun
from bcbio.distributed.split import parallel_split_combine
from bcbio.distributed.transaction import file_transaction
from bcbio.pipeline import config_utils, shared
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do, programs
from bcbio.variation import bedutils
from bcbio.variation import multi as vmulti

def parallel_callable_loci(in_bam, ref_file, data):
    config = copy.deepcopy(data["config"])
    num_cores = config["algorithm"].get("num_cores", 1)
    data = {"work_bam": in_bam, "config": config,
            "reference": data["reference"]}
    parallel = {"type": "local", "cores": num_cores, "module": "bcbio.distributed"}
    items = [[data]]
    with prun.start(parallel, items, config, multiplier=int(num_cores)) as runner:
        split_fn = shared.process_bam_by_chromosome("-callable.bed", "work_bam", remove_alts=True)
        out = parallel_split_combine(items, split_fn, runner,
                                     "calc_callable_loci", "combine_bed",
                                     "callable_bed", ["config"])[0]
    return out[0]["callable_bed"]

@multi.zeromq_aware_logging
def calc_callable_loci(data, region=None, out_file=None):
    """Determine callable bases for an input BAM in the given region.
    """
    if out_file is None:
        out_file = "%s-callable.bed" % os.path.splitext(data["work_bam"])[0]
    max_depth = utils.get_in(data, ("config", "algorithm", "coverage_depth_max"), 10000)
    depth = {"max": max_depth * 7 if max_depth > 0 else sys.maxint - 1,
             "min": utils.get_in(data, ("config", "algorithm", "coverage_depth_min"), 4)}
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            ref_file = tz.get_in(["reference", "fasta", "base"], data)
            region_file, calc_callable = _regions_for_coverage(data, region, ref_file, tx_out_file)
            if calc_callable:
                _group_by_ctype(_get_coverage_file(data["work_bam"], ref_file, region, region_file, depth,
                                                   tx_out_file, data),
                                depth, region_file, tx_out_file)
            # special case, do not calculate if we are in a chromosome not covered by BED file
            else:
                os.rename(region_file, tx_out_file)
    return [{"callable_bed": out_file, "config": data["config"], "work_bam": data["work_bam"]}]

def _group_by_ctype(bed_file, depth, region_file, out_file):
    """Group adjacent callable/uncallble regions into defined intervals.

    Uses tips from bedtools discussion:
    https://groups.google.com/d/msg/bedtools-discuss/qYDE6XF-GRA/2icQtUeOX_UJ
    https://gist.github.com/arq5x/b67196a46db5b63bee06
    """
    def assign_coverage(feat):
        feat.name = _get_ctype(float(feat.name), depth)
        return feat
    full_out_file = "%s-full%s" % utils.splitext_plus(out_file)
    with open(full_out_file, "w") as out_handle:
        kwargs = {"g": [1, 4], "c": [1, 2, 3, 4], "ops": ["first", "first", "max", "first"]}
        # back compatible precision https://github.com/chapmanb/bcbio-nextgen/issues/664
        if LooseVersion(programs.get_version_manifest("bedtools", True)) >= LooseVersion("2.22.0"):
            kwargs["prec"] = 21
        for line in open(pybedtools.BedTool(bed_file).each(assign_coverage).saveas()
                                                     .groupby(**kwargs).fn):
            out_handle.write("\t".join(line.split("\t")[2:]))
    pybedtools.BedTool(full_out_file).intersect(region_file, nonamecheck=True).saveas(out_file)

def _get_coverage_file(in_bam, ref_file, region, region_file, depth, base_file, data):
    """Retrieve summary of coverage in a region.
    Requires positive non-zero mapping quality at a position, matching GATK's
    CallableLoci defaults.
    """
    out_file = "%s-genomecov.bed" % utils.splitext_plus(base_file)[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            bam.index(in_bam, data["config"])
            fai_file = ref.fasta_idx(ref_file, data["config"])
            sambamba = config_utils.get_program("sambamba", data["config"])
            bedtools = config_utils.get_program("bedtools", data["config"])
            max_depth = depth["max"] + 1
            cmd = ("{sambamba} view -F 'mapping_quality > 0' -L {region_file} -f bam -l 1 {in_bam} | "
                   "{bedtools} genomecov -split -ibam stdin -bga -g {fai_file} -max {max_depth} "
                   "> {tx_out_file}")
            do.run(cmd.format(**locals()), "bedtools genomecov: %s" % (str(region)), data)
    # Empty output file, no coverage for the whole contig
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for feat in get_ref_bedtool(ref_file, data["config"], region):
                    out_handle.write("%s\t%s\t%s\t%s\n" % (feat.chrom, feat.start, feat.end, 0))
    return out_file

def _get_ctype(count, depth):
    if count == 0:
        return "NO_COVERAGE"
    elif count < depth["min"]:
        return "LOW_COVERAGE"
    elif count > depth["max"]:
        return "EXCESSIVE_COVERAGE"
    else:
        return "CALLABLE"

def _regions_for_coverage(data, region, ref_file, out_file):
    """Retrieve BED file of regions we need to calculate coverage in.
    """
    variant_regions = bedutils.merge_overlaps(dd.get_variant_regions(data), data)
    ready_region = shared.subset_variant_regions(variant_regions, region, out_file)
    custom_file = "%s-coverageregions.bed" % utils.splitext_plus(out_file)[0]
    if not ready_region:
        get_ref_bedtool(ref_file, data["config"]).saveas(custom_file)
        return custom_file, True
    elif os.path.isfile(ready_region):
        return ready_region, True
    elif isinstance(ready_region, (list, tuple)):
        c, s, e = ready_region
        pybedtools.BedTool("%s\t%s\t%s\n" % (c, s, e), from_string=True).saveas(custom_file)
        return custom_file, True
    else:
        with file_transaction(data, custom_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for feat in get_ref_bedtool(ref_file, data["config"], region):
                    out_handle.write("%s\t%s\t%s\t%s\n" % (feat.chrom, feat.start, feat.end, "NO_COVERAGE"))
        return custom_file, variant_regions is None

def sample_callable_bed(bam_file, ref_file, data):
    """Retrieve callable regions for a sample subset by defined analysis regions.
    """
    config = data["config"]
    out_file = "%s-callable_sample.bed" % os.path.splitext(bam_file)[0]
    with shared.bedtools_tmpdir({"config": config}):
        callable_bed = parallel_callable_loci(bam_file, ref_file, data)
        input_regions_bed = config["algorithm"].get("variant_regions", None)
        if not utils.file_uptodate(out_file, callable_bed):
            with file_transaction(config, out_file) as tx_out_file:
                callable_regions = pybedtools.BedTool(callable_bed)
                filter_regions = callable_regions.filter(lambda x: x.name == "CALLABLE")
                if input_regions_bed:
                    if not utils.file_uptodate(out_file, input_regions_bed):
                        input_regions = pybedtools.BedTool(input_regions_bed)
                        filter_regions.intersect(input_regions, nonamecheck=True).saveas(tx_out_file)
                else:
                    filter_regions.saveas(tx_out_file)
    return out_file

def calculate_offtarget(bam_file, ref_file, data):
    """Generate file of offtarget read counts for inputs with variant regions.
    """
    vrs_file = dd.get_variant_regions(data)
    if vrs_file:
        out_file = "%s-offtarget-stats.yaml" % os.path.splitext(bam_file)[0]
        if not utils.file_exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                offtarget_regions = "%s-regions.bed" % utils.splitext_plus(out_file)[0]
                ref_bed = get_ref_bedtool(ref_file, data["config"])
                ref_bed.subtract(pybedtools.BedTool(vrs_file)).saveas(offtarget_regions)
                cmd = ("samtools view -u {bam_file} -L {offtarget_regions} | "
                       "bedtools intersect -sorted -abam - -b {offtarget_regions} -f 1.0 -bed | wc -l")
                offtarget_count = int(subprocess.check_output(cmd.format(**locals()), shell=True))
                cmd = "samtools idxstats {bam_file} | awk '{{s+=$3}} END {{print s}}'"
                mapped_count = int(subprocess.check_output(cmd.format(**locals()), shell=True))
                with open(tx_out_file, "w") as out_handle:
                    yaml.safe_dump({"mapped": mapped_count, "offtarget": offtarget_count}, out_handle,
                                   allow_unicode=False, default_flow_style=False)
        return out_file

def get_ref_bedtool(ref_file, config, chrom=None):
    """Retrieve a pybedtool BedTool object with reference sizes from input reference.
    """
    broad_runner = broad.runner_from_config(config, "picard")
    ref_dict = broad_runner.run_fn("picard_index_ref", ref_file)
    ref_lines = []
    with contextlib.closing(pysam.Samfile(ref_dict, "r")) as ref_sam:
        for sq in ref_sam.header["SQ"]:
            if not chrom or sq["SN"] == chrom:
                ref_lines.append("%s\t%s\t%s" % (sq["SN"], 0, sq["LN"]))
    return pybedtools.BedTool("\n".join(ref_lines), from_string=True)

def _get_nblock_regions(in_file, min_n_size, ref_regions):
    """Retrieve coordinates of regions in reference genome with no mapping.
    These are potential breakpoints for parallelizing analysis.
    """
    out_lines = []
    called_contigs = set([])<commit_msg>Avoid warning messages from bedtools with out of sort BAM and BED files<commit_after>"""Examine callable regions following genome mapping of short reads.

Identifies callable analysis regions surrounded by larger regions lacking
aligned bases. This allows parallelization of smaller chromosome chunks
through post-processing and variant calling, with each sub-section
mapping handled separately.

Regions are split to try to maintain relative uniformity across the
genome and avoid extremes of large blocks or large numbers of
small blocks.
"""
import contextlib
import copy
from distutils.version import LooseVersion
import operator
import os
import subprocess
import sys

import numpy
import pybedtools
import pysam
import toolz as tz
import yaml

from bcbio import bam, broad, utils
from bcbio.bam import ref
from bcbio.log import logger
from bcbio.distributed import multi, prun
from bcbio.distributed.split import parallel_split_combine
from bcbio.distributed.transaction import file_transaction
from bcbio.pipeline import config_utils, shared
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do, programs
from bcbio.variation import bedutils
from bcbio.variation import multi as vmulti

def parallel_callable_loci(in_bam, ref_file, data):
    config = copy.deepcopy(data["config"])
    num_cores = config["algorithm"].get("num_cores", 1)
    data = {"work_bam": in_bam, "config": config,
            "reference": data["reference"]}
    parallel = {"type": "local", "cores": num_cores, "module": "bcbio.distributed"}
    items = [[data]]
    with prun.start(parallel, items, config, multiplier=int(num_cores)) as runner:
        split_fn = shared.process_bam_by_chromosome("-callable.bed", "work_bam", remove_alts=True)
        out = parallel_split_combine(items, split_fn, runner,
                                     "calc_callable_loci", "combine_bed",
                                     "callable_bed", ["config"])[0]
    return out[0]["callable_bed"]

@multi.zeromq_aware_logging
def calc_callable_loci(data, region=None, out_file=None):
    """Determine callable bases for an input BAM in the given region.
    """
    if out_file is None:
        out_file = "%s-callable.bed" % os.path.splitext(data["work_bam"])[0]
    max_depth = utils.get_in(data, ("config", "algorithm", "coverage_depth_max"), 10000)
    depth = {"max": max_depth * 7 if max_depth > 0 else sys.maxint - 1,
             "min": utils.get_in(data, ("config", "algorithm", "coverage_depth_min"), 4)}
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            ref_file = tz.get_in(["reference", "fasta", "base"], data)
            region_file, calc_callable = _regions_for_coverage(data, region, ref_file, tx_out_file)
            if calc_callable:
                _group_by_ctype(_get_coverage_file(data["work_bam"], ref_file, region, region_file, depth,
                                                   tx_out_file, data),
                                depth, region_file, tx_out_file)
            # special case, do not calculate if we are in a chromosome not covered by BED file
            else:
                os.rename(region_file, tx_out_file)
    return [{"callable_bed": out_file, "config": data["config"], "work_bam": data["work_bam"]}]

def _group_by_ctype(bed_file, depth, region_file, out_file):
    """Group adjacent callable/uncallble regions into defined intervals.

    Uses tips from bedtools discussion:
    https://groups.google.com/d/msg/bedtools-discuss/qYDE6XF-GRA/2icQtUeOX_UJ
    https://gist.github.com/arq5x/b67196a46db5b63bee06
    """
    def assign_coverage(feat):
        feat.name = _get_ctype(float(feat.name), depth)
        return feat
    full_out_file = "%s-full%s" % utils.splitext_plus(out_file)
    with open(full_out_file, "w") as out_handle:
        kwargs = {"g": [1, 4], "c": [1, 2, 3, 4], "ops": ["first", "first", "max", "first"]}
        # back compatible precision https://github.com/chapmanb/bcbio-nextgen/issues/664
        if LooseVersion(programs.get_version_manifest("bedtools", True)) >= LooseVersion("2.22.0"):
            kwargs["prec"] = 21
        for line in open(pybedtools.BedTool(bed_file).each(assign_coverage).saveas()
                                                     .groupby(**kwargs).fn):
            out_handle.write("\t".join(line.split("\t")[2:]))
    pybedtools.BedTool(full_out_file).intersect(region_file, nonamecheck=True).saveas(out_file)

def _get_coverage_file(in_bam, ref_file, region, region_file, depth, base_file, data):
    """Retrieve summary of coverage in a region.
    Requires positive non-zero mapping quality at a position, matching GATK's
    CallableLoci defaults.
    """
    out_file = "%s-genomecov.bed" % utils.splitext_plus(base_file)[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            bam.index(in_bam, data["config"])
            fai_file = ref.fasta_idx(ref_file, data["config"])
            sambamba = config_utils.get_program("sambamba", data["config"])
            bedtools = config_utils.get_program("bedtools", data["config"])
            max_depth = depth["max"] + 1
            cmd = ("{sambamba} view -F 'mapping_quality > 0' -L {region_file} -f bam -l 1 {in_bam} | "
                   "{bedtools} genomecov -split -ibam stdin -bga -g {fai_file} -max {max_depth} "
                   "> {tx_out_file}")
            do.run(cmd.format(**locals()), "bedtools genomecov: %s" % (str(region)), data)
    # Empty output file, no coverage for the whole contig
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for feat in get_ref_bedtool(ref_file, data["config"], region):
                    out_handle.write("%s\t%s\t%s\t%s\n" % (feat.chrom, feat.start, feat.end, 0))
    return out_file

def _get_ctype(count, depth):
    if count == 0:
        return "NO_COVERAGE"
    elif count < depth["min"]:
        return "LOW_COVERAGE"
    elif count > depth["max"]:
        return "EXCESSIVE_COVERAGE"
    else:
        return "CALLABLE"

def _regions_for_coverage(data, region, ref_file, out_file):
    """Retrieve BED file of regions we need to calculate coverage in.
    """
    variant_regions = bedutils.merge_overlaps(dd.get_variant_regions(data), data)
    ready_region = shared.subset_variant_regions(variant_regions, region, out_file)
    custom_file = "%s-coverageregions.bed" % utils.splitext_plus(out_file)[0]
    if not ready_region:
        get_ref_bedtool(ref_file, data["config"]).saveas(custom_file)
        return custom_file, True
    elif os.path.isfile(ready_region):
        return ready_region, True
    elif isinstance(ready_region, (list, tuple)):
        c, s, e = ready_region
        pybedtools.BedTool("%s\t%s\t%s\n" % (c, s, e), from_string=True).saveas(custom_file)
        return custom_file, True
    else:
        with file_transaction(data, custom_file) as tx_out_file:
            with open(tx_out_file, "w") as out_handle:
                for feat in get_ref_bedtool(ref_file, data["config"], region):
                    out_handle.write("%s\t%s\t%s\t%s\n" % (feat.chrom, feat.start, feat.end, "NO_COVERAGE"))
        return custom_file, variant_regions is None

def sample_callable_bed(bam_file, ref_file, data):
    """Retrieve callable regions for a sample subset by defined analysis regions.
    """
    config = data["config"]
    out_file = "%s-callable_sample.bed" % os.path.splitext(bam_file)[0]
    with shared.bedtools_tmpdir({"config": config}):
        callable_bed = parallel_callable_loci(bam_file, ref_file, data)
        input_regions_bed = config["algorithm"].get("variant_regions", None)
        if not utils.file_uptodate(out_file, callable_bed):
            with file_transaction(config, out_file) as tx_out_file:
                callable_regions = pybedtools.BedTool(callable_bed)
                filter_regions = callable_regions.filter(lambda x: x.name == "CALLABLE")
                if input_regions_bed:
                    if not utils.file_uptodate(out_file, input_regions_bed):
                        input_regions = pybedtools.BedTool(input_regions_bed)
                        filter_regions.intersect(input_regions, nonamecheck=True).saveas(tx_out_file)
                else:
                    filter_regions.saveas(tx_out_file)
    return out_file

def calculate_offtarget(bam_file, ref_file, data):
    """Generate file of offtarget read counts for inputs with variant regions.
    """
    vrs_file = dd.get_variant_regions(data)
    if vrs_file:
        out_file = "%s-offtarget-stats.yaml" % os.path.splitext(bam_file)[0]
        if not utils.file_exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                offtarget_regions = "%s-regions.bed" % utils.splitext_plus(out_file)[0]
                ref_bed = get_ref_bedtool(ref_file, data["config"])
                ref_bed.subtract(pybedtools.BedTool(vrs_file)).saveas(offtarget_regions)
                cmd = ("samtools view -u {bam_file} -L {offtarget_regions} | "
                       "bedtools intersect -abam - -b {offtarget_regions} -f 1.0 -bed | wc -l")
                offtarget_count = int(subprocess.check_output(cmd.format(**locals()), shell=True))
                cmd = "samtools idxstats {bam_file} | awk '{{s+=$3}} END {{print s}}'"
                mapped_count = int(subprocess.check_output(cmd.format(**locals()), shell=True))
                with open(tx_out_file, "w") as out_handle:
                    yaml.safe_dump({"mapped": mapped_count, "offtarget": offtarget_count}, out_handle,
                                   allow_unicode=False, default_flow_style=False)
        return out_file

def get_ref_bedtool(ref_file, config, chrom=None):
    """Retrieve a pybedtool BedTool object with reference sizes from input reference.
    """
    broad_runner = broad.runner_from_config(config, "picard")
    ref_dict = broad_runner.run_fn("picard_index_ref", ref_file)
    ref_lines = []
    with contextlib.closing(pysam.Samfile(ref_dict, "r")) as ref_sam:
        for sq in ref_sam.header["SQ"]:
            if not chrom or sq["SN"] == chrom:
                ref_lines.append("%s\t%s\t%s" % (sq["SN"], 0, sq["LN"]))
    return pybedtools.BedTool("\n".join(ref_lines), from_string=True)

def _get_nblock_regions(in_file, min_n_size, ref_regions):
    """Retrieve coordinates of regions in reference genome with no mapping.
    These are potential breakpoints for parallelizing analysis.
    """
    out_lines = []
    called_contigs = set([])
{'fraction_non_alphanumeric': 0.11327574875598535, 'fraction_numerical': 0.005445498075298094, 'mean_word_length': 3.1054153015995376, 'pattern_counts': {'":': 34, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3585349', 'n_tokens_mistral': 6960, 'n_tokens_neox': 6697, 'n_words': 1588}