"""Functionality to query and extract information from aligned BAM files.
"""
import collections
import contextlib
import os
import itertools
import signal
import subprocess
import numpy

import pysam
import toolz as tz

from bcbio import utils
from bcbio.bam import ref
from bcbio.distributed import objectstore
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger
from bcbio.pipeline import config_utils
import bcbio.pipeline.datadict as dd
from bcbio.provenance import do

def is_paired(bam_file):
    """Determine if a BAM file has paired reads.

    Works around issues with head closing the samtools pipe using signal trick from:
    http://stackoverflow.com/a/12451083/252589
    """
    bam_file = objectstore.cl_input(bam_file)
    cmd = ("set -o pipefail; "
           "sambamba view -h {bam_file} | head -50000 | "
           "sambamba view -S -F paired /dev/stdin  | head -1 | wc -l")
    p = subprocess.Popen(cmd.format(**locals()), shell=True,
                         executable=do.find_bash(),
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                         preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL))
    stdout, stderr = p.communicate()
    if p.returncode == 0 or p.returncode == 141 and stderr.strip() == "":
        return int(stdout) > 0
    else:
        raise ValueError("Failed to check paired status of BAM file: %s" % str(stderr))

def index(in_bam, config, check_timestamp=True):
    """Index a BAM file, skipping if index present.

    Centralizes BAM indexing providing ability to switch indexing approaches.
    """
    assert is_bam(in_bam), "%s in not a BAM file" % in_bam
    index_file = "%s.bai" % in_bam
    alt_index_file = "%s.bai" % os.path.splitext(in_bam)[0]
    if check_timestamp:
        bai_exists = utils.file_uptodate(index_file, in_bam) or utils.file_uptodate(alt_index_file, in_bam)
    else:
        bai_exists = utils.file_exists(index_file) or utils.file_exists(alt_index_file)
    if not bai_exists:
        # Remove old index files and re-run to prevent linking into tx directory
        for fname in [index_file, alt_index_file]:
            utils.remove_safe(fname)
        sambamba = _get_sambamba(config)
        samtools = config_utils.get_program("samtools", config)
        num_cores = config["algorithm"].get("num_cores", 1)
        with file_transaction(config, index_file) as tx_index_file:
            assert tx_index_file.find(".bam.bai") > 0
            tx_bam_file = tx_index_file.replace(".bam.bai", ".bam")
            utils.symlink_plus(in_bam, tx_bam_file)
            if sambamba:
                cmd = "{sambamba} index -t {num_cores} {tx_bam_file}"
            else:
                cmd = "{samtools} index {tx_bam_file}"
            do.run(cmd.format(**locals()), "Index BAM file: %s" % os.path.basename(in_bam))
    return index_file if utils.file_exists(index_file) else alt_index_file

def remove(in_bam):
    """
    remove bam file and the index if exists
    """
    if utils.file_exists(in_bam):
        utils.remove_safe(in_bam)
    if utils.file_exists(in_bam + ".bai"):
        utils.remove_safe(in_bam + ".bai")

def idxstats(in_bam, data):
    """Return BAM index stats for the given file, using samtools idxstats.
    """
    index(in_bam, data["config"])
    AlignInfo = collections.namedtuple("AlignInfo", ["contig", "length", "aligned", "unaligned"])
    samtools = config_utils.get_program("samtools", data["config"])
    idxstats_out = subprocess.check_output([samtools, "idxstats", in_bam])
    out = []
    for line in 
idxstats_out.split("\n"):
        if line.strip():
            contig, length, aligned, unaligned = line.split("\t")
            out.append(AlignInfo(contig, int(lengt
h), int(aligned), int(unaligned)))
    return out

def get_downsample_pct(in_bam, target_counts, data):
    """Retrieve percentage of file to downsample to get to target counts.
    """
    total = sum(x.aligned for x in idxstats(in_bam, data))
    with contextlib.closing(pysam.Samfile(in_bam, "rb")) as work_bam:
        n_rgs = max(1, len(work_bam.header.get("RG", [])))
    rg_target = n_rgs * target_counts
    if total > rg_target:
        return float(rg_target) / float(total)

def get_aligned_reads(in_bam, data):
    index(in_bam, data["config"])
    bam_stats = idxstats(in_bam, data)
    align = sum(x.aligned for x in bam_stats)
    unaligned = sum(x.unaligned for x in bam_stats)
    total = float(align + unaligned)
    return 1.0 * align / total

def downsample(in_bam, data, target_counts, read_filter="", always_run=False,
               work_dir=None):
    """Downsample a BAM file to the specified number of target counts.
    """
    index(in_bam, data["config"])
    ds_pct = get_downsample_pct(in_bam, target_counts, data)
    if always_run and not ds_pct:
        ds_pct = 1.0
    if ds_pct:
        out_file = "%s-downsample%s" % os.path.splitext(in_bam)
        if work_dir:
            out_file = os.path.join(work_dir, os.path.basename(out_file))
        if not utils.file_exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                sambamba = config_utils.get_program("sambamba", data["config"])
                num_cores = dd.get_num_cores(data)
                cmd = ("{sambamba} view -t {num_cores} {read_filter} -f bam -o {tx_out_file} "
                       "--subsample={ds_pct:.3} --subsampling-seed=42 {in_bam}")
                do.run(cmd.format(**locals()), "Downsample BAM file: %s" % os.path.basename(in_bam))
        return out_file

def check_header(in_bam, rgnames, ref_file, config):
    """Ensure passed in BAM header matches reference file and read groups names.
    """
    _check_bam_contigs(in_bam, ref_file, config)
    _check_sample(in_bam, rgnames)

def _check_sample(in_bam, rgnames):
    """Ensure input sample name matches expected run group names.
    """
    with contextlib.closing(pysam.Samfile(in_bam, "rb")) as bamfile:
        rg = bamfile.header.get("RG", [{}])
    msgs = []
    warnings = []
    if len(rg) > 1:
        warnings.append("Multiple read groups found in input BAM. Expect single RG per BAM.")
    elif len(rg) == 0:
        msgs.append("No read groups found in input BAM. Expect single RG per BAM.")
    elif rg[0].get("SM") != rgnames["sample"]:
        msgs.append("Read group sample name (SM) does not match configuration `description`: %s vs %s"
                    % (rg[0].get("SM"), rgnames["sample"]))
    if len(msgs) > 0:
        raise ValueError("Problems with pre-aligned input BAM file: %s\n" % (in_bam)
                         + "\n".join(msgs) +
                         "\nSetting `bam_clean: picard` in the configuration can often fix this issue.")
    if warnings:
        print("*** Potential problems in input BAM compared to reference:\n%s\n" %
              "\n".join(warnings))

def _check_bam_contigs(in_bam, ref_file, config):
    """Ensure a pre-aligned BAM file matches the expected reference genome.
    """
    ref_contigs = [c.name for c in ref.file_contigs(ref_file, config)]
    with contextlib.closing(pysam.Samfile(in_bam, "rb")) as bamfile:
        bam_contigs = [c["SN"] for c in bamfile.header["SQ"]]
    problems = []
    warnings = []
    for bc, rc in itertools.izip_longest(bam_contigs, ref_contigs):
        if bc != rc:
            if bc and rc:
                problems.append("Reference mismatch. BAM: %s Reference: %s" % (bc, rc))
            elif bc:
                warnings.append("Extra BAM chromosomes: %s" % bc)
            elif rc:
                warnings.append("Extra reference chromosomes: %s" % rc)
    if problems:
        raise ValueError("Unexpected order, name or contig mismatches between input BAM and reference file:\n%s\n"
                         "Setting `bam_clean: picard` in the configuration can often fix this issue."
                         % "\n".join(problems))
    if warnings:
        print("*** Potential problems in input BAM compared to reference:\n%s\n" %
              "\n".join(warnings))


def open_samfile(in_file):
    if is_bam(in_file):
        return pysam.Samfile(in_file, "rb")
    elif is_s