#!/usr/bin/env python

# -*- coding: utf-8 -*-

from pysam.libcalignmentfile import IteratorColumnRegion

from collections import Counter

def pileup_truncated(bam,contig, start, stop,**kwargs):

    """

    Obtain Pysam columns only at selected region

    contig(str) : contig to perform pileup

    start(int) : coordinate of first colummn (zero based)

    stop(int) : coordinate of last column (exclusive)

    **kwargs : arguments to pass to pysam.AlignmentFile.IteratorColumn

    """

    has_coord, rtid, rstart, rstop = bam.parse_region(contig, start, stop )

    yield from IteratorColumnRegion(bam,

                                    tid=rtid,

                                    start=rstart,

                                    stop=rstop,truncate=True,**kwargs)

def has_variant_reads(pysam_alignment_file, chrom, pos, alt, min_reads=2, stepper='nofilter'):

    """

    Check if the alignment file contains evidence for the supplied base

    Args:

        pysam_alignment_file(pysam.AlignmentFile) : file to check location

        chrom(str): name of contig

        pos(int) : position to check (zero based)

        alt(str): base to check

    """

    obs = Counter()

    for pile in pileup_truncated(pysam_alignment_file,chrom,pos,pos+1,stepper=stepper):

        if pos!=pile.reference_pos:

            continue

        for read in pile.pileups:

            if not read.is_del and not read.is_refskip:

                obs[read.alignment.query_sequence[read.query_position]]+=1

    return obs[alt]>=min_reads

def mate_pileup(alignments, contig, position,**kwargs):

    """

    Extract all fragments (R1, and R2) which overlap with the supplied position

    Example:

        >>> alignments = pysam.AlignmentFile('example.bam')

        >>> piled_reads = dict()

        >>> obs = collections.Counter()

        >>> pile_pos = 8774

        >>> pile_contig = '1'

        >>> add_pile_mates(alignments, piled_reads, pile_contig, pile_pos, obs )

        [[<pysam.libcalignedsegment.AlignedSegment at 0x7f8a7eca4ac8>,

          <pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecc83a8>],

         [<pysam.libcalignedsegment.AlignedSegment at 0x7f8a7d10ce28>,

          <pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecc8468>],

         [<pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecc8348>,

          <pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecae4c8>],

         [<pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecc8948>,

          <pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecae588>],

         [<pysam.libcalignedsegment.AlignedSegment at 0x7f8a7ecc4ee8>, ...

    Args:

        alignments(pysam.AlignmentFile) : Alignment file to extract reads from

        contig(str) : contig to perform pileup

        position(int) : coordinate to perform pileup

    Returns:

        list

    """

    piled_reads = dict()

    _mate_pileup(alignments=alignments,

                 piled_reads=piled_reads,

                 contig=contig,

                 position=position, obs=None, add_missing_only=False,**kwargs)

    return list(piled_reads.values())

def _mate_pileup(alignments, piled_reads, contig, position, obs=None, add_missing_only=False, max_depth=50000, stepper='nofilter',ignore_overlaps=False,ignore_orphans=False):

    """

    Extract all fragments (R1, and R2) which overlap with the supplied position

    Args:

        alignments(pysam.AlignmentFile) : Alignment file to extract reads from

        piled_reads(dict) : (Empty) Dictionary to which the reads are written

        contig(str) : contig to perform pileup

        position(int) : coordinate to perform pileup

        obs(collections.Counter) : Store base observation frequencies at the location of the pileup in this dictionary

        add_missing_only(bool): Only add missing mates to the existing dictionary (piled_reads)

    """

    for pile in alignments.pileup(contig,position,position+1,stepper=stepper,ignore_overlaps=ignore_overlaps,ignore_orphans=ignore_orphans,max_depth=max_depth):

        if position!=pile.reference_pos:

            continue

        for read in pile.pileups:

            if not add_missing_only and (read.is_del or read.is_refskip):

                continue

            if read.alignment.is_supplementary:

                continue

            if add_missing_only and not read.alignment.qname in piled_reads:

                continue

            if not read.alignment.qname in piled_reads:

                piled_reads[read.alignment.qname] = [None, None]

            piled_reads[read.alignment.qname][read.alignment.is_read2] = read.alignment

            if obs is not None:

                obs[read.alignment.query_sequence[read.query_position]]+=1

    if not add_missing_only:

        find_missing_mates(alignments, piled_reads)

def find_missing_mates(alignments, piled_reads):

    """

    Obtain missing mates

    Args:

        piled_reads(dict) : { query_name:[R1, R2], ... }

    """

    added=True

    piled_locations = set()

    while added:

        added=False

        for qname, reads in piled_reads.items():

            R1, R2 = reads

            if R1 is None and not R2.mate_is_unmapped:

                # Find R1 ..

                if look_for_mate(alignments, R2, piled_reads, piled_locations):

                    added=True

                    break

            if R2 is None and not R1.mate_is_unmapped:

                # Find R2 ..

                if look_for_mate(alignments, R1, piled_reads, piled_locations):

                    added=True

                    break

def look_for_mate(alignments, read, piled_reads, piled_locations):

    """

    Given a read, find the mate and possibly other missing reads in the piled_reads dictionary.

         Skips coordinates in piled_locations set.

     Args:

        alignments(pysam.AlignmentFile) : Alignment file to extract reads from

        read(pysam.AlignedSegment) : read to find mate for

        piled_reads(dict) : Query dictionary  { query_name:[R1, R2], ... }

        piled_locations(set) : Locations which have already been searched

    Returns:

        has_searched(bool) : Boolean which indicates if the location of the mate was already present in piled_locations and the location was not checked again.

    """

    location_descriptor = (read.next_reference_name, read.next_reference_start)

    if location_descriptor in piled_locations:

        return False

    piled_locations.add( location_descriptor )

    _mate_pileup(

                    alignments,

                    piled_reads,

                    contig = location_descriptor[0],

                    position = location_descriptor[1],

                    add_missing_only=True

                )

    if piled_reads[read.qname ][ read.is_read1] is None:

        raise(ValueError( f'Failed to find {read.qname} at {location_descriptor}' ))

    return True