# -*- coding: utf-8 -*-
"""
Created on Tue Jul 28 12:54:17 2020

@author: rjovelin
"""


import pysam
import os
import subprocess
import itertools
from Bio import pairwise2
from Bio.pairwise2 import format_alignment
import regex
import gzip



def read_panel(smmip_panel):
    '''
    (str) -> dict
    
    Returns a dictionary storing information about the smmips using >mip_key as key
    Precondition: smmip_panel is a tab delimited file and expected column names are present
    
    Parameters
    ----------
    - smmip_panel (str): Path to file with smmip information.
                         File generated by MIPGEN. Contains info about the smmips
                         included in the panel 
    '''
    
    D = {}
    infile = open(smmip_panel)
    header = infile.readline().rstrip().split('\t')
    
    for line in infile:
        line = line.rstrip()
        if line != '':
            line = line.split('\t')
            # get smmip key
            mip_key = line[header.index('>mip_key')]
            D[mip_key] = {i: line[header.index(i)] for i in header}
            # convert types
            for i in D[mip_key]:
                if i == 'logistic_score':
                    D[mip_key][i] = float(D[mip_key][i])
                # convert length to int
                elif i in ['target length', 'lig arm length', 'ext arm length',
                           'for primer length', 'rev primer length', 'umi length',
                           'pcr product length']:
                    D[mip_key][i] = int(D[mip_key][i])
                # convert coordinates to 0-based
                elif i in ['ext_probe_start', 'ext_probe_stop', 'lig_probe_start',
                           'lig_probe_stop', 'mip_scan_start_position', 'mip_scan_stop_position',
                           'feature_start_position', 'feature_stop_position']:
                    if 'start' in i:
                        D[mip_key][i] = int(D[mip_key][i]) - 1
                    elif 'stop' in i:
                        D[mip_key][i] = int(D[mip_key][i])
                
                # chromosome field is not a valid chromosome name
                # WARNING. IT WOULD BE BETTER IF THE PANEL FILE HAS CORRECT CHR NAMES INSTEAD OF ADJUSTING
                # FOR ALTERNATIVE CHROMOS, IT MAY CAUSE AN ERROR
                elif i == 'chr':
                    if 'chr' not in D[mip_key][i].lower():
                        D[mip_key][i] = 'chr' + D[mip_key][i]
    infile.close()
    return D


def reverse_complement(S):
    '''
    (str) -> str
    
    Returns the reverse complement of sequence S.
    Nucleotides other than [ACGT] are labeled N
    
    Parameters
    ----------
    - S (str): A nucleotide sequence
    
    Examples
    --------
    >>> reverse_complement('aaAA')
    'TTtt'
    >>> reverse_complement('aaaa')
    'tttt'
    >>> reverse_complement('Atcg')
    'cgaT'
    >>> reverse_complement('AtcX')
    'NgaT'
    >>> reverse_complement('TTCX')
    'NGAA'
    '''
    
    MapNuc = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A',
              'a': 't', 'c': 'g', 'g': 'c', 't': 'a'}
    s = ''
    for i in S:
        if i in MapNuc:
            s = MapNuc[i] + s
        else:
            s = 'N' + s
    return s


def create_tree(outdir):
    '''
    (str) -> list
    
    Creates a directory structure in directory outdir
    and return a list with the paths of the subdirectories created
    
    Parameters
    ----------
    - outdir (str): Path to directory where directory structure is created
    
    Examples
    --------
    >>> create_tree('foo')
    ['foo/out', 'foo/stats', 'foo/alignment']   
        
    '''
    
    # create outdir doesn't exist
    os.makedirs(outdir, exist_ok=True)
    # create subdirectories if they don't exist
    tree = []
    for i in ['out', 'stats', 'alignment']:
        j = os.path.join(outdir, i)
        os.makedirs(j, exist_ok=True)
        tree.append(j)
    return tree



def align_fastqs(fastq1, fastq2, reference, outdir, bwa, prefix, remove):
    '''
    (str, str, str, str, str, str, str, bool) -> None
    
    Align fastq1 and fastq2 using bwa mem into coordinate-sorted and indexed bam
    in the out directory in outdir
    
    Parameters
    ----------
    - fastq1 (str): Path to Fastq1
    - fastq2 (str): Path to Fastq2
    - reference (str): Path to the reference genome
    - outdir (str): Path to output directory where subdirectories are written  
    - bwa (str): Path to the bwa script
    - prefix (str): Name of the aligned bam file
    - remove (bool): Remove intermediate files if True                     
    '''

    # create subdirectories
    finaldir, statsdir, aligndir = create_tree(outdir)
    
    # save samfile in alignment directory
    samfile = os.path.join(aligndir, os.path.basename(prefix) + '.sam') 
    
    # align fastqs
    cmd = "{0} mem {1} {2} {3} > {4}".format(bwa, reference, fastq1, fastq2, samfile)
    subprocess.call(cmd, shell=True)

    # convert sam to bam
    bamfile = os.path.join(aligndir, os.path.basename(prefix) + '.bam') 
    pysam.view('-b', '-h', '-o', bamfile, samfile, catch_stdout=False)
    
    # sort bam on coordinates and index and write to out directory
    sortedbam = os.path.join(finaldir, os.path.basename(prefix) + '.sorted.bam')
    pysam.sort('-o', sortedbam, bamfile)
    pysam.index(sortedbam)

    # remove intermediate files
    if remove == True:
        # remove sam file and non-sorted bam
        os.remove(samfile)
        os.remove(bamfile)        
    


def track_read(read, d):
    '''
    (pysam.AlignedSegment, dict) -> None
    
    Update d with read name, read key, value pair in place
    
    Parameters
    ----------
    - read (pysam.AlignedSegment): An aligned segment
    - d (dict): Dictionary with read name, list of pysam reads key, value pairs
    '''
        
    if read.query_name not in d:
        d[read.query_name] = [read]
    else:
        d[read.query_name].append(read)


def reorder_reads(L):
    '''
    (list) -> None

    Order the list of paired reads L in place with read1, read2

    Parameters
    ----------
    - L (list): A list of paired pysam reads
    '''
    
    # check if paired reads
    assert len(L) == 2
    # reorder such that L = [read1, read2]       
    if L[0].is_read1 == False:
        # instert read1 at beginning
        L.insert(0, L.pop(1))
    assert L[0].is_read1 == True and L[1].is_read2 == True

   
def remove_read(d, readname):
    '''
    (dict, str) -> None
    
    Remove readname from d if d is present. Modifies d in place
    
    Parameters
    ----------
    - d (dict): Dictionary with read names as keys
    - readname (str): The name of a pysam.AlignedSegment read
    '''
    
    if readname in d:
        del d[readname]


def extract_from_regex(read_seq, p):
    '''
    (str, _regex.Pattern) -> str

    Returns the UMI sequence extracted from read_seq if present or the empty string
        
    Parameters
    ----------
    - read_seq (str): Sequence of the read
    - p (_regex.Pattern): Compiled regex pattern used for matching pattern in read sequence
     
    Examples
    --------
    read_seq = 'TCATGTCTGCTAATGGGAAAGAGTGTCCTAACTGTCCCAGATCGTTTTTTCTCACGTCTTTTCTCCTTTCACTTCTCTTTTTCTTTTTCTTTCTTCTTCTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT'
    # UMI starts at the beginning of the read
    >>> extract_from_regex(read_seq, regex.compile('(?<umi_1>.{12})(?<discard_1>ATGGGAAAGAGTGTCC)'))
    'TCATGTCTGCTA'
    # UMI does not start at the beginning of the read
    read_seq = 'ATCATGTCTGCTAATGGGAAAGAGTGTCCTAACTGTCCCAGATCGTTTTTTCTCACGTCTTTTCTCCTTTCACTTCTCTTTTTCTTTTTCTTTCTTCTTCTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT'
    # first nucleotide is part of the new read sequence
    >>> extract_from_regex(read_seq, regex.compile('(?<umi_1>.{12})(?<discard_1>ATGGGAAAGAGTGTCC)'))
    'TCATGTCTGCTA'
    # force UMI to start at the beginning of the read sequence
    >>> extract_from_regex(read_seq, regex.compile('(?<umi_1>^.{12})(?<discard_1>ATGGGAAAGAGTGTCC)'))
    ''
    # discard nuceotides upstream of UMI
    >>> extract_from_regex(read_seq, regex.compile('(?<discard_1>.*)(?<umi_1>.{12})(?<discard_2>ATGGGAAAGAGTGTCC)'))
    'TCATGTCTGCTA'
    '''
    
    # initialize umi_seq
    umi_seq = ''
        
    # scan through the string looking for a match
    m = p.search(read_seq)
    # process if match is found
    if m:
        # collect umi positions
        umi_pos = []
        for i in m.groupdict():
            if 'umi' in i:
                umi_pos.append(m.span(i))
        # sort umi and discard positions
        umi_pos.sort()
        # get umi sequences
        umi_seq = ''.join([read_seq[i[0]:i[1]] for i in umi_pos])
    return umi_seq


def get_template_positions(L):
    '''
    (list) -> tuple
    
    Return a tuple with template start and end coordinates 0-based
    
    Parameters
    ----------
    - L (list): List of pysam reads
    '''
    
    # make a list of read coordinates
    read_positions = []
    for read in L:
        # add start, end positions for each paired read
        read_positions.append(read.reference_start)
        read_positions.append(read.reference_end)
    # sort list --> template coordinates are defined by 1st and last positions
    read_positions.sort()
    
    return read_positions[0], read_positions[-1]


def get_target_coordinates(panel, smmip):
    '''
    (dict, str) -> (int, int)
    
    Return a tuple with the start and end positions (0-based) of the target
    region for the given smmip     
        
    Parameters
    ----------
    - panel (dict): Panel information
    - smmip (str): A specific smmip in the panel
    '''
        
    # get expected target coordinates (including probes)
    L = sorted([panel[smmip]['ext_probe_start'], panel[smmip]['ext_probe_stop'],
         panel[smmip]['lig_probe_start'], panel[smmip]['lig_probe_stop']])
    start, end = L[0], L[-1] 
    return start, end


def compute_overlap(template_start, template_end, target_start, target_end):
    '''
    (int, int, int, int) -> int
    
    Return the length of overlap between a smMIP and a read pair    
    Precondition: The reads are mapped on the chromosome of the target region    
    
    Parameters
    ----------
    - template_start (int): Start position of the template region delimited by read1 and read2 (0-based, included)
    - template_end (int): End position of the template region delimited by read1 and read2 (0-based excluded)
                               Refers to position 1 nucleotide after the last nucleotide of the template
    - target_start (int): Start position of the target region on the reference (0-based, included)
    - target_end (int): End position of the target region of the reference (0-based excluded). 
                             Refers to position 1 nucleotide after the last nucleotide of the target 
    '''
    
    overlap = 0
    if target_start > template_end or template_start > target_end:
        overlap = 0
    elif template_start <= target_start <= template_end:
        if target_end >= template_end:
            overlap = template_end - target_start
        else:
            overlap = target_end - target_start
    elif target_start <= template_start <= target_end:
        if template_end >= target_end:
            overlap = target_end - template_start
        else:
            overlap = template_end - template_start
    return overlap
    

def find_overlapping_candidates(chromo, panel, template_start, template_end):
    '''
    (str, dict, int, int, float) -> list
    
    Return a list candidate smmips overlapping with a read pair reverse-sorted by the length of the overlap
    
    Parameters
    ----------
    - chromo (str): Chromosome name 
    - panel (dict): Panel information
    - template_start (int): Start position of the template region delimited by read1 and read2 (0-based, included)
    - template_end (int): End position of the template region delimited by read1 and read2 (0-based excluded)
                               Refers to position 1 nucleotide after the last nucleotide of the template
    '''
    
    # make a list of mips on chromo
    mips = [i for i in panel if panel[i]['chr'] == chromo]
    
    # find overlapping smmips on chromo
    candidates = []
    if len(mips) != 0:
        for smmips in mips:
            # get expected target coordinates (including probes)
            Start, End = get_target_coordinates(panel, smmips)
            # compute overlap 
            overlap = compute_overlap(template_start, template_end, Start, End)
            if overlap > 0:
                # found a smmip candidate
                candidates.append([overlap, smmips])

        # sort candidates smmips on overlap length
        candidates.sort(key=lambda x:x[0])
        # reverse sort candidate smmips by overlap length
        candidates.reverse()
        
    return candidates


def align_sequences(seq1, seq2, match, mismatch, gap_opening, gap_extension):
    '''
    (str, str, num, num, num, num) -> list
    
    Return the possible local pairwise alignments between
    seq1 and seq2 as a list of tuples each containing 
    the aligned sequences of seq1 and seq2, the alignment score, 
    the start and end of the alignment.
    
    Parameters
    ----------
    - seq1 (str): Sequence 1
    - seq2 (str): Sequence 2
    - match (float or int): Score of identical characters
    - mismatch (float or int): Score of non-identical characters
    - gap_opening (float or int): Score for opening a gap
    - gap_extension (float or int): Score for extending an open gap
    '''
    
    alignments = pairwise2.align.localms(seq1, seq2, match, mismatch, gap_opening, gap_extension)
    return alignments



def get_highest_scoring_aligment(alignments):
    '''
    (list) -> tuple

    Return the alignment with highest score among the possible alignments.
    Alignment is a tuple with the aligned sequences of seq1 and seq2,
    the alignment score, the start and end of the alignment.

    Parameters
    ----------
    - alignments (list): List of possible pairwise alignments
                              generated by AlignSequences
    '''
    
    alignments.sort(key = lambda x: x[1])
    best = alignments[-1]
    return best    


def is_empty_smmip(seq1, seq2, match, mismatch, gap_opening, gap_extension, overlap_threshold, matches_threshold):
    '''
    (str, str, int, float, float, float, float) -> bool    
    
    Return True is the read pair did not capture the target (ie, is empty) and False if not empty
    
    Parameters
    ----------
    - seq1 (str): Sequence of read 1
    - seq2 (str): Sequence of read 2
    - match (float or int): Score of identical characters
    - mismatch (float or int): Score of non-identical characters
    - gap_opening (float or int): Score for opening a gap
    - gap_extension (float or int): Score for extending an open gap
    - overlap_threshold (float or int): Cut-off value for the length of 
                                        the de-gapped overlap between read1 and read2 
    - matches_threshold (flot or int): Cut-off value for the number of matching
                                       positions within the de-gapped overlap
                                       between read1 and read2 
    '''
    # initiate variable
    empty = False
    # align upper cases sequences 
    seq1, seq2 = seq1.strip().upper(), seq2.strip().upper()    
    # do a local alignment between r1 and rev. compl of r2
    alignments = align_sequences(seq1, reverse_complement(seq2), match, mismatch, gap_opening, gap_extension)
    # check if alignemnts exist
    if len(alignments) == 0:
        # consider empty smmip
        empty = True
    else:
        # find the best scoring alignment
        best = get_highest_scoring_aligment(alignments)
        # get overlap length without gapped positions
        ali1, matches, ali2, score, end = format_alignment(*best).split('\n')
        overlap = matches.replace(' ', '')
        # count matching positions within de-gapped overlap
        matching = overlap.count('|')
        # record readnames of empty smmips
        if len(overlap) < overlap_threshold and matching / len(overlap) > matches_threshold:
            # consider empty smmip
            empty = True    
    return empty



def add_custom_tag(read, tag, value):
    '''
    (pysam.AlignedSegment, str, str) -> None
    
    Updates the read in place with tag: value
    
    Parameters
    ----------
    - read (pysam.AlignedSegment): An aligned segment
    - tag (str): 2-letter code of the custom tag. See samtools tags for specs
    - value (str): The value of the tag to be added. Type is string (Z datatype)
    '''
    
    # get the read tags
    read_tags = read.get_tags()
    # replace tags with tag: value (str)   
    read_tags += [(tag, value, 'Z')]
    read.set_tags(read_tags)


def update_read_tags(read, smmip_name, umi_seq):
    '''
    (pysam.AlignedSegment, str, str) -> None
    
    Update the read in place with tag: value
    
    Parameters
    ----------
    - read (pysam.AlignedSegment): An aligned segment
    - smmip_name (str): Name of of the smmip mapping to the read 
    - umi_seq (str): UMI sequence  
    '''
    
    # add smmip tag
    add_custom_tag(read, 'SP', smmip_name)
    # add UMI tag
    add_custom_tag(read, 'MI', umi_seq)
    # add UMI sequence 
    add_custom_tag(read, 'OX', umi_seq)


def remove_bam_extension(bamfile):
    '''
    (str) -> str 
    
    Returns the path to bam file stripped from its file extension if extension
    is ".sorted.bam", ".bam", ".assigned_reads.bam", ".unassigned_reads.bam",
    ".empty_reads.bam" or returns path to bam file
    
    Parameters
    ----------
    - bamfile (str): Path to bam file
    '''
    
    
    if '.temp.' in bamfile:
        return bamfile[:bamfile.rfind('.temp.')]
    else:
        if bamfile[-len('.sorted.bam'):] == '.sorted.bam':
            return bamfile[: -len('.sorted.bam')]
        elif bamfile[-len('.assigned_reads.bam'):] == '.assigned_reads.bam':
            return bamfile[: -len('.assigned_reads.bam')]
        elif bamfile[-len('.unassigned_reads.bam'):] == '.unassigned_reads.bam':
            return bamfile[: -len('.unassigned_reads.bam')]
        elif bamfile[-len('.empty_reads.bam'):] == '.empty_reads.bam':
            return bamfile[: -len('.empty_reads.bam')]
        elif bamfile[-len('.bam'):] == '.bam':
            return bamfile[: -len('.bam')]
        else:
            return bamfile
    


def get_bam_header(bamfile):
    '''
    (str) -> dict
    
    Returns the header of the bam file as a dictionary
    
    Parameters
    ----------
    - bamfile (str): Path to the bam file
    '''
       
    # parse bam header
    infile = pysam.AlignmentFile(bamfile)
    header = dict(infile.header)
    infile.close()
    return header


def find_matching_smmip(candidates, panel, seq1, seq2, upstream_nucleotides, umi_length, max_subs):
    '''
    (list, dict, str, str, int, int, int) -> (list, list)

    Returns a list of mipname matching a given read and a list of extracted umi sequences from the matched read     
    
    Parameters
    ----------
    - candidates (list): List of candidate smMIPs (smMIPs overlapping read1 and read2) reverse-sorted by overlap length
                         Each element of the list is a 2-item list with overlap length and smMIP Id
    - panel (dict): Panel information        
    - seq1 (str): Sequence of read 1
    - seq2 (str): Sequence of read 2
    - upstream_nucleotides (int): Maximum number of nucleotides upstream the UMI sequence
    - umi_length (int): Length of the UMI    
    - max_subs (int): Maximum number of substitutions allowed in the probe sequence
    '''
    
    matching, umi_sequences = [], []
    for i in candidates:
        extension_probe = panel[i[1]]['ext_probe_sequence']
        ligation_probe = panel[i[1]]['lig_probe_sequence']
        p1 = regex.compile('(?P<discard_1>^.{0,%s})' %upstream_nucleotides + '(?P<umi_1>.{%s})' %umi_length + '(' + reverse_complement(ligation_probe) + ')' + '{s<=%s}.*' % max_subs)
        p2 = regex.compile('(?P<discard_1>^.{0,%s})' %upstream_nucleotides + '(?P<umi_1>.{%s})' %umi_length + '(' + extension_probe + ')' + '{s<=%s}.*' % max_subs)
        umi_seq1 = extract_from_regex(seq1, p1)
        umi_seq2 = extract_from_regex(seq2, p2)
        if umi_seq1 and umi_seq2:
            # found matching probe
            matching.append(panel[i[1]]['mip_name'])
            umi_sequences.append(umi_seq1 + umi_seq2)
            # stop searching for matching probes if found
            break
    return matching, umi_sequences                        


    
def record_read(metrics, category, newbam, L):
    '''
    (dict, str, pysam.libcalignmentfile.AlignmentFile, list) -> None
    
    Update read counter category in place for each read in L and write each reads to newbam
    
    Parameters
    ----------
    
    - metrics (dict): Dictionary to track read counts
    - category (str): read category, key in metrics. Values are: 'reads', 'assigned', 'not_assigned', 'assigned_empty', 'assigned_not_empty'
    - newbam (pysam.libcalignmentfile.AlignmentFile): Bam file open for writing in binary mode
    - L (list): list of pysam.AlignedSegment reads
    '''
    
    for read in L:
        # update counter
        metrics[category] += 1
        # write read to newbam
        newbam.write(read)
                
        
def sort_index_bam(filename, suffix):
    '''
    (str, str) -> None    
    
    Sort and index filename. Generate filename + suffix bam and filename + suffix.bai index

    Parameters
    ----------
    
    - filename (str): Path to bam file to be sorted
    - suffix (str): Suffix string to be part of the sorted file name, appended to filename
    '''
    
    sorted_file = remove_bam_extension(filename) + suffix
    pysam.sort('-o', sorted_file, filename)
    pysam.index(sorted_file)


def get_positions(start, end, chromo_length):
    '''
    (int | None, int | None, int) -> (int, int)
    
    Returns a tuple with 0-based start and end positions of a contig region
    
    Parameters
    ----------
    - start (int or None): Start position if defined
    - end (int or None): End position if defined
    - chromo_length (int): Length of the reference (chromosome)
    '''
    
    if start:
        start_pos = start
    else:
        start_pos = 0
    if end:
        # check if end position > chromosome length
        if end > chromo_length:
            end_pos = chromo_length
        else:
            end_pos = end
    else:
        end_pos = chromo_length
    return start_pos, end_pos


def assign_reads_to_smmips(bamfile, assigned_file, empty_file, panel, upstream_nucleotides, umi_length, max_subs, match, mismatch, gap_opening, gap_extension, alignment_overlap_threshold, matches_threshold, chromosome, start, end):
    '''
    (str, pysam.AlignmentFile, pysam.AlignmentFile, dict, int, int, int, float, float, float, float, float, float, bool, str | None) -> (dict, dict)
    
    Return a tuple of dictionaries with read counts. The first dictionary counts
    total reads, assigned and unassigned reads as well as empty smmips.
    The second dictionary tracks the number of reads per smmip
    
    Write assigned reads, assigned but empty smmips to 2 separate output bams
    Assigned reads are tagged with the smMip name and the extracted UMI sequence
        
    Pre-condition: bamfile is sorted on coordinates and indexed (a .bai exists in the same directory)
    
    Parameters
    ----------
    - bamfile (str): Path to the input bam sorted on coordinates
    - assigned_file (pysam.AlignmentFile): Bam file opened to write assigned reads
    - empty_file (pysam.AlignmentFile): Bam file opened to write empty reads
    - panel (dict): Panel information        
    - upstream_nucleotides (int): Maximum number of nucleotides upstream the UMI sequence
    - umi_length (int): Length of the UMI    
    - max_subs (int): Maximum number of substitutions allowed in the probe sequence
    - match (float or int): Score of identical characters
    - mismatch (float or int): Score of non-identical characters
    - gap_opening (float or int): Score for opening a gap
    - gap_extension (float or int): Score for extending an open gap
    - alignment_overlap_threshold (float or int): Cut-off value for the length of the de-gapped overlap between read1 and read2 
    - matches_threshold (float or int): Cut-off value for the number of matching positions within the de-gapped overlap between read1 and read2 
    - chromosome (str | None): Specifies the genomic region in the alignment file where reads are mapped.
                               Examine reads on chromosome if used and on all chromosomes if None
                               Chromosome format must match format in the bam header
    - start (int | None): Start position of region on chromosome if defined
    - end (int | None): End position of region on chromosome if defined
    '''
    
    # count total reads in file, reads in region, assigned and unassigned reads
    metrics = {'assigned': 0, 'assigned_empty': 0, 'assigned_not_empty': 0}
    
    # count the total number of reads in file, excluding unmapped reads, secondary and supplementary alignments
    with pysam.AlignmentFile(bamfile, 'rb') as infile:
        total_count = infile.count(until_eof=False, read_callback='all')
    metrics['total'] = total_count  
    # count the number of reads in region, ignoring unmapped reads, secondary and supplementary alignments 
    with pysam.AlignmentFile(bamfile, 'rb') as infile:
        read_count = infile.count(until_eof=False, read_callback='all')
    metrics['reads'] = read_count
        
    # count smmips
    smmip_counts = {panel[i]['mip_name'] : {'empty':0, 'not_empty':0} for i in panel}
            
    # create AlignmentFile object to read input bam
    infile = pysam.AlignmentFile(bamfile, 'rb')
    
    # make a list of chromosomes in panel
    panel_chromosomes = [panel[i]['chr'] for i in panel]
    # make a list of chromosomes from the bam header
    header = get_bam_header(bamfile)
    bam_chromosomes = [i['SN'] for i in header['SQ']]
    
    # use specific chromosome if defined 
    # check that chromosome is in the bam header
    if chromosome:
        bam_chromosomes = [chromosome] if chromosome in bam_chromosomes else []
    
    # only look at reads expected to map on chromosomes in panel
    # discard reads mapping to other chromosomes
    for contig in bam_chromosomes:
        # create dictionary to keep track of read pairs 
        D = {}
        if contig in panel_chromosomes:
            # get the start and end positions of region of interest. include all contig if start and end are not defined
            chromo_length = infile.get_reference_length(contig)
            start_pos, end_pos = get_positions(start, end, chromo_length)
            # loop over all reads mapping to contig
            for query in infile.fetch(contig=contig,start=start_pos, end=end_pos, until_eof=False):
                qname = query.query_name
                # ignore unmapped reads, secondary and supplementary alignments
                if query.is_secondary == False and query.is_supplementary == False and query.is_unmapped == False:
                    # collect all reads with same query name in D until all are found
                    track_read(query, D)
                if qname in D:
                    # process reads if paired reads have been found
                    if len(D[qname]) == 2:
                        # reorder reads
                        reorder_reads(D[qname])
                        # get template coordinates
                        template_start, template_end = get_template_positions(D[qname])
                        # find candidate overlapping smmips, reverse-sorted by overlap length 
                        candidates = find_overlapping_candidates(contig, panel, template_start, template_end)
                        if len(candidates) == 0:
                            # no overlapping smmip. discard read
                            remove_read(D, qname)
                        else:
                            # get read sequences
                            seq1, seq2 = D[qname][0].get_forward_sequence(), D[qname][1].get_forward_sequence()
                            # find matching probes
                            # make a lists of smmips with matching probes and corresponding umis
                            matching, umi_sequences = find_matching_smmip(candidates, panel, seq1, seq2, upstream_nucleotides, umi_length, max_subs)
                            if len(matching) == 0:
                                # no matching probes. discard read
                                remove_read(D, qname)
                            else:
                                # get the name of matching smmip and corresponding umi
                                matching_smmip = matching[0]
                                umi_seq = umi_sequences[0]
                                # get the name of the smmip with greatest overlap
                                overlapping_smmip = panel[candidates[0][1]]['mip_name']
                                # check if smmip with matching is the same as smmip with greatest overlap
                                if matching_smmip != overlapping_smmip:
                                    # conflict between pre-assigned smmip and overlapping smmip. discard read
                                    remove_read(D, qname)
                                else:
                                    # assign read. add tags
                                    for read in D[qname]:
                                        update_read_tags(read, matching_smmip, umi_seq)
                                        metrics['assigned'] += 1
                                    # check if read is empty
                                    if is_empty_smmip(seq1, seq2, match, mismatch, gap_opening, gap_extension, alignment_overlap_threshold, matches_threshold):
                                        # consider an empty smmip. write reads to empty file
                                        record_read(metrics, 'assigned_empty', empty_file, D[qname])
                                        remove_read(D, qname)
                                        smmip_counts[matching_smmip]['empty'] += 2
                                    else:
                                        # non empty reads. write to outputbam
                                        record_read(metrics, 'assigned_not_empty', assigned_file, D[qname])
                                        remove_read(D, qname)
                                        smmip_counts[matching_smmip]['not_empty'] += 2
    # close bams    
    infile.close()
    
    # update metrics dict
    # unassigned reads are reads not assigned in the region of interest
    metrics['not_assigned'] = metrics['reads'] - metrics['assigned']
    assigned_ratio = round(metrics['assigned'] / metrics['reads'] * 100, 4) if metrics['reads'] != 0 else 0
    unassigned_ratio = round(metrics['not_assigned'] / metrics['reads'] * 100, 4) if metrics['reads'] != 0 else 0
    empty_ratio = round(metrics['assigned_empty'] / metrics['assigned'] * 100, 4) if metrics['assigned'] != 0 else 0
    metrics.update({'percent_assigned': assigned_ratio})
    metrics.update({'percent_not_assigned': unassigned_ratio})
    metrics.update({'percent_empty_smmips': empty_ratio})
      
    return metrics, smmip_counts    


def get_consecutive_items(L):
    '''
    (list) -> generator
    
    Returns a generator with 1st and last item of consecutive items in L
    
    Parameters
    ----------
    - L (list): List
    
    Examples
    --------
    >>> list(get_consecutive_items([1, 2, 3, 6, 7, 8, 10,11, 24, 25, 26, 27, 28, 30, 50]))
    [(1, 3), (6, 8), (10, 11), (24, 28), (30, 30), (50, 50)]
    '''
    
    # remove duplicate and sort L
    L = sorted(set(L))
    for key, group in itertools.groupby(enumerate(L), lambda t: t[1] - t[0]):
        group = list(group)
        # return a generator with 1st and last item of consecutive items in L
        yield group[0][1], group[-1][1]


def make_non_overlapping_regions(panel):
    '''
    (dict) -> dict
    
    Returns a dictionary with lists of coordinates (0-based) for non-overlapping
    regions on each chromosome
       
    Parameters
    ----------
    - panel (dict): Panel information
    '''
    
    # collect all positions per chromosome
    D = {}
    for i in panel:
        chromo = panel[i]['chr']
        # get target coordinates 0-based
        Start, End = get_target_coordinates(panel, i)
        if chromo not in D:
            D[chromo] = []
        # need to adjust End because get_consecituve items returns coordinates
        D[chromo].extend([j for j in range(Start, End+1)])
 
    # remove duplicate positions
    for chromo in D:
        D[chromo] = sorted(list(set(D[chromo])))
    
    # find non-overlapping regions on each chromosome
    regions = {}
    for chromo in D:
        # get consecutive regions
        regions[chromo] = [j for j in get_consecutive_items(D[chromo])]
    
    return regions



def get_base_position(read, pos):
    '''
    (pysam.PileupRead, int) -> (str, str)

    Returns a tuple with reference and alternative bases at position pos
    
    Parameters
    ----------
    - read (pysam.PileupRead): Representation of a read aligned to a particular position in the reference sequence.
    - pos (int): Position in the reference genome (0-based)
    '''
    
    # get aligned read, ref pos and ref base 
    pairs = read.alignment.get_aligned_pairs(with_seq=True)
    # read.indel looks ahead to see if indel at next position(s)  
    if read.indel == 0:
        # no indel, record ref and alt 
        # get index of pileupcolumn pos in aligned pairs
        j = [i[1] for i in pairs].index(pos)
        # record base on ref and read
        ref_base = pairs[j][-1].upper()
        alt_base = read.alignment.query_sequence[read.query_position].upper()
    elif read.indel > 0:
        # next position is an insertion
        # get index of pileupcolumn pos in aligned pairs
        j = [i[1] for i in pairs].index(pos)
        # record base on ref and insertion on read
        ref_base = pairs[j][-1].upper()
        alt_base = read.alignment.query_sequence[read.query_position:read.query_position + abs(read.indel) + 1].upper()
    elif read.indel < 0:
        # next position is deletion
        # get index of pileupcolumn pos in aligned pairs
        j = [i[1] for i in pairs].index(pos)
        # record base on ref at pos + ref bases deleted on read and base on read
        ref_base = ''.join([i[-1] for i in pairs[j: j +  abs(read.indel) + 1]]).upper()
        alt_base = read.alignment.query_sequence[read.query_position]
        
    return ref_base, alt_base



def track_allele_counts(D, pos, mipname, UMI, allele):
    '''
    (dict, int, str, str, tuple) -> None
    
    Update dictionary D tracking allele count per smmip and UMI and position with allele
    information at pos for a given read. Modidy D in place
    
    Parameters
    ----------
    - D (dict): Dictionary with allele count per position and smmip
    - pos (int): Position in the reference genome 
    - mipname (str): Name of a smmip
    - UMI (str): UMI sequence
    - allele (tuple): Tuple with reference nucleotide and base in read at position 
    '''
    
    # get ref and alt bases
    ref_base, alt_base = allele
    
    # keep track of ref_base at pos
    if pos not in D:
        D[pos] = {}
    if 'ref_base' not in D[pos]:
        # check ref_base.
        # ref base is deleted region + ref base if only read with indel overlap positions
        if len(ref_base) == 1:
            D[pos]['ref_base'] = ref_base
        else:
            D[pos]['ref_base'] = ref_base[0]
    # count the number of reads supporting this allele
    if 'smmips' not in D[pos]:
        D[pos]['smmips'] = {}
    if mipname not in D[pos]['smmips']:
        D[pos]['smmips'][mipname] = {}
    if UMI not in D[pos]['smmips'][mipname]:
        D[pos]['smmips'][mipname][UMI] = {}
    if allele in D[pos]['smmips'][mipname][UMI]:
        D[pos]['smmips'][mipname][UMI][allele] += 1
    else:
        D[pos]['smmips'][mipname][UMI][allele] = 1



def get_allele_counts_target(bamfile, contig, region_start, region_end, max_depth, truncate, ignore_orphans, stepper):
    '''
    
    (str, str, int, int, int, bool, bool, str) -> dict
    
    Returns a dictionary keeping track of allele counts at each position per smmip

    Parameters
    ----------    
    - bamfile (str): Path to the coordinate-sorted and indexed bam file with annotated reads with smMIP and UMI tags 
    - contig: Chromosome name, eg. chrN
    - region_start (int): Start index of the region of interest. 0-based half opened
    - region_end (int): End index of the region of interest. 0-based half opened
    - max_depth (int): Maximum read depth
    - truncate: Consider only pileup columns within interval defined by region start and end if True
    - ignore_orphans: Ignore orphan reads (paired reads not in proper pair) if True
    - stepper: Controls how the iterator advances. Accepted values:
               'all': skip reads with following flags: BAM_FUNMAP, BAM_FSECONDARY, BAM_FQCFAIL, BAM_FDUP
               'nofilter': uses every single read turning off any filtering
    '''

    # create a dict to store consensus seq info
    # {pos: {'ref_base': ref_base, 'smmips': {smmip_name: {allele: count}}}}
    D = {}
    
    with pysam.AlignmentFile(bamfile, "rb") as reader:
        # loop over pileup columns
        for pileupcolumn in reader.pileup(contig, region_start, region_end, truncate=truncate, ignore_orphans=ignore_orphans, max_depth=max_depth, stepper=stepper):
            # get column position
            pos = int(pileupcolumn.reference_pos)  
            # restict pileup columns to genomic region
            if region_start <= pos < region_end:
                # loop over pileupreads in pileup column
                for read in pileupcolumn.pileups:
                    # skip unmapped reads and reads mapping to multiple positions 
                    if any([read.alignment.is_unmapped, read.alignment.is_secondary, read.alignment.is_supplementary]) == False:
                        # get smmip name from tags
                        tags = read.alignment.get_tags()
                        tags = {i[0]:i[1] for i in tags}
                        mipname = tags['SP']
                        # get UMI from tags
                        UMI = tags['MI']
                        # skip positions with deletions or ref not defined
                        # events are captured at the position before they occur
                        if not read.is_del and not read.is_refskip:
                            # read.indel looks ahead to see if indel at next position(s) 
                            # 0 --> not indel; > 0 --> insertion; < 0 --> deletion
                            # get ref and alt bases at position pos
                            allele = get_base_position(read, pos)
                            ref_base, alt_base = allele
                            # count allele per smmip and position
                            track_allele_counts(D, pos, mipname, UMI, allele)
    return D    



def count_alleles_across_panel(bamfile, panel, max_depth, truncate, ignore_orphans, stepper):
    '''
    (str, panel, int, bool, bool, str) -> dict

    Parameters
    ----------
    - bamfile (str): Path to the coordinate-sorted and indexed bam with annotated read with smMIP and UMI
    - panel (dict): Panel information
    - max_depth (int): Maximum read depth
    - truncate: Consider only pileup columns within interval defined by region start and end if True
    - ignore_orphans: Ignore orphan reads (paired reads not in proper pair) if True
    - stepper: Controls how the iterator advances. Accepted values:
               'all': skip reads with following flags: BAM_FUNMAP, BAM_FSECONDARY, BAM_FQCFAIL, BAM_FDUP
               'nofilter': uses every single read turning off any filtering
    
    Return a dictionary with allele counts per smmips and UMI and target regions
    '''
    
    # get coordinates of non-overlapping regions on each chromosome
    regions = make_non_overlapping_regions(panel)
    # track allele counts for each region of all chromosomes in panel
    D = {}
    for chromo in regions:
        D[chromo] = []
        for target in regions[chromo]:
            c = get_allele_counts_target(bamfile, chromo, target[0], target[1], max_depth, truncate, ignore_orphans, stepper)
            D[chromo].append(c)
    return D


def sort_chromos(L):
    '''
    (list) -> list
    
    Return a list of chromosome names sorted numerically, with non-numeric
    chromosomes sorted at the end of the list
    
    Parameters
    ----------
    - L (list): List of chromosome names
    '''
        
    # make lists of chromos
    numerical, non_numerical, non_chr = [], [], []
    for chromo in L:
        if chromo.startswith('chr'):
            chromo = chromo.replace('chr', '')
            if chromo.isnumeric() == False:
                non_numerical.append(chromo)
            else:
                numerical.append(int(chromo))
        else:
            non_chr.append(chromo)
    # sort each list separately
    non_chr.sort()
    numerical.sort()
    non_numerical.sort()
    # add back chr
    for i in range(len(non_numerical)):
        non_numerical[i] = 'chr' + non_numerical[i]
    for i in range(len(numerical)):
        numerical[i] = 'chr' + str(numerical[i])
    # concatenate lists
    return numerical + non_numerical + non_chr


def count_variants(d):
    '''
    (dict) -> dict

    Returns a dictionary with single nucleotide and indel counts

    Parameters
    ----------
    - d (dict): Dictionary with allele counts
    '''
    
    # count each allele, initiate with single nucleotides
    counts = {'A': 0, 'C': 0, 'G': 0, 'T': 0, 'N': 0}
    for allele in d:
        # count single nucleotides
        if len(allele[0]) == 1 and len(allele[1]) == 1:
            # snv or no change
            counts[allele[1]] += d[allele]
        else:
            # indel, record allele and its count
            if allele in counts:
                counts[allele] += d[allele]
            else:
                counts[allele] = d[allele]
    return counts


def list_indel_counts(counts):
    '''
    (dict) -> (list, list)
    
    Returns a tuple with lists of 2-item lists (counts and alleles) for inserstions
    and deletions sorted by allele counts
    
    Parameters
    ----------
    - counts (dict): Dictionary of allele counts
    '''
    
    # make lists of indels and indel counts 
    D, I = [], []
    for allele in counts:
        if len(allele) == 2:
            # indel
            if len(allele[1]) > 1:
                # insertions
                I.append([counts[allele], allele])
            elif len(allele[0]) > 1:
                # deletions
                D.append([counts[allele], allele])
    D.sort()
    I.sort()
    return I, D


def get_genomic_positions(D):
    '''
    (dict) -> dict
    
    Returns a dictionary with all positions (1-based) on each chromosome with
    variant information from regions in the panel 
    
    Parameters
    ----------
    - D (dict): Dictionary with allele counts per smmips and UMI and target regions.
                Output from function count_alleles_across_panel
    '''
    
    # create a dict {chromo: [positions]}
    positions = {}
    
    # make a list of positions for each chromosome
    for chromo in D:
        positions[chromo] = []
        # loop dicts with allele counts in target on chromo
        for C in D[chromo]:
            # make a sorted list of positions
            pos = list(map(lambda x: str(int(x) + 1), list(set(C.keys()))))
            positions[chromo].extend(pos)
    return positions


def parse_cosmic(reference, cosmicfile, positions):
    '''
    (str, str, dict) -> dict
    
    Returns a dictionary with mutation information at each position in reference
    
    Parameters
    ----------
    - reference (str): Reference genome. Accepted values: 37 or 38
    - cosmicfile (str): Cosmic file. Tab separated table of all COSMIC coding
                        point mutations from targeted and genome wide screens
    - positions (dict): Dictionary with all positions (1-based) on each chromosome
                        with variant information from regions in the panel 
    '''
    
    infile = gzip.open(cosmicfile, 'rt', errors='ignore')
    header = infile.readline()
    # create a dict {chromo: {pos: [{pos:x, gene_name: x, cosmic_id: x, cosmic_mutation: x, description: x}] 
    D = {}
    for line in infile:
        line = line.rstrip()
        if line != '':
            line = line.split('\t')
            genome = line[24]
            if genome == reference:
                chromo, position = line[25].split(':')
                chromo = 'chr' + chromo
                position = position.split('-')[0]
                if chromo in positions and position in positions[chromo]:
                    gene_name = line[0]
                    cosmic_id = line[17]
                    description = line[21]
                    mutation = line[39]
                    d = {'position': position, 'gene_name': gene_name,
                         'cosmic_id': cosmic_id, 'description': description, 'mutation': mutation}
                    if chromo not in D:
                        D[chromo] = {}
                    # records last mutation if multiple mutations at a given position
                    D[chromo][position] = d
    infile.close()
    return D            



def write_table_variants(D, outputfile, cosmic):
    '''
    (dict, str) -> None
    
    Writes a summary table with all nucleotide and indel counts detected for each smMIP
    
    Parameters
    ----------
    - D (dict): Dictionary with allele counts per smmips and target regions
    - outputfile (str): Path to output file
    - cosmic (dict): Cosmic information about point mutations in reference GRCh37 or CRCh38
    '''
    
    # make a sorted list of chromosomes
    chromosomes = sort_chromos([chromo for chromo in D])
    
    # open file for writing
    newfile = open(outputfile, 'w')
    Header = ['CHROM', 'POS', 'SMMIP', 'UMI', 'REF', 'A', 'C', 'G', 'T', 'N', 'I_(ref,ins)', 'I_counts', 'D_(ref,del)', 'D_counts', 'RAWDP', 'REF_FREQ',
              'GENE_NAME', 'COSMIC_ID', 'DESCRIPTION', 'MUTATION']
    newfile.write('\t'.join(Header) + '\n')

    # loop over chromosomes
    for chromo in chromosomes:
        # loop dicts with allele counts in target on chromo
        for C in D[chromo]:
            # make a sorted list of positions
            positions = sorted(list(map(lambda x: int(x), list(set(C.keys())))))
            # loop over positions
            for pos in positions:
                # loop over smmip
                for smmip_name in sorted(C[pos]['smmips'].keys()):
                    # get the reference
                    ref_base = C[pos]['ref_base']
                    # count snvs and indels for each UMI
                    UMIs = sorted(list(C[pos]['smmips'][smmip_name].keys()))
                    umi_counts = [count_variants(C[pos]['smmips'][smmip_name][umi]) for umi in UMIs]
                    umi_insertions, umi_deletions = [], [] 
                    for i in umi_counts:
                        insertion, deletion = list_indel_counts(i)
                        umi_insertions.append(insertion)
                        umi_deletions.append(deletion)
                    # count snvs and indels at the mip level (adding counts from each UMI)
                    smmip_counts = {}
                    for i in umi_counts:
                        for j in i:
                            if j in smmip_counts:
                                smmip_counts[j] += i[j]
                            else:
                                smmip_counts[j] = i[j]
                    smmip_insertions, smmip_deletions = list_indel_counts(smmip_counts)                
                    # compute raw depth
                    rawdepth = 0
                    for i in smmip_counts:
                        rawdepth += smmip_counts[i]
                    # compute ref frequency
                    ref_freq = smmip_counts[ref_base] / rawdepth
                    # write info for counts at smmip level                
                    line = [chromo, pos + 1, smmip_name, 'all', ref_base, smmip_counts['A'], smmip_counts['C'],
                            smmip_counts['G'], smmip_counts['T'], smmip_counts['N'],
                            ';'.join([str(i[1]) for i in smmip_insertions]), ';'.join([str(i[0]) for i in smmip_insertions]),
                            ';'.join([str(i[1]) for i in smmip_deletions]), ';'.join([str(i[0]) for i in smmip_deletions]),
                            str(rawdepth), str(round(ref_freq, 5))]
                    # add cosmic info
                    if chromo in cosmic and str(pos+1) in cosmic[chromo]:
                        line.extend([cosmic[chromo][str(pos+1)]['gene_name'],
                                 cosmic[chromo][str(pos+1)]['cosmic_id'],
                                 cosmic[chromo][str(pos+1)]['description'],
                                 cosmic[chromo][str(pos+1)]['mutation']])
                    else:
                        line.extend(['NA', 'NA', 'NA', 'NA'])
                    newfile.write('\t'.join(list(map(lambda x: str(x), line))) + '\n')
                    # write info for counts at umi level
                    for i in range(len(umi_counts)):
                        # compte raw depth
                        rawdepth = 0
                        for j in umi_counts[i]:
                            rawdepth += umi_counts[i][j]
                        ref_freq = umi_counts[i][ref_base] / rawdepth
                        line = [chromo, pos + 1, smmip_name, UMIs[i], ref_base, umi_counts[i]['A'], umi_counts[i]['C'],
                            umi_counts[i]['G'], umi_counts[i]['T'], umi_counts[i]['N'],
                            ';'.join([str(k[1]) for k in umi_insertions[i]]), ';'.join([str(k[0]) for k in umi_insertions[i]]),
                            ';'.join([str(k[1]) for k in umi_deletions[i]]), ';'.join([str(k[0]) for k in umi_deletions[i]]),
                            str(rawdepth), str(round(ref_freq, 5))]
                        # add cosmic info
                        if chromo in cosmic and str(pos+1) in cosmic[chromo]:
                            line.extend([cosmic[chromo][str(pos+1)]['gene_name'],
                                     cosmic[chromo][str(pos+1)]['cosmic_id'],
                                     cosmic[chromo][str(pos+1)]['description'],
                                     cosmic[chromo][str(pos+1)]['mutation']])
                        else:
                            line.extend(['NA', 'NA', 'NA', 'NA'])
                        newfile.write('\t'.join(list(map(lambda x: str(x), line))) + '\n')
                    
    # close file after writing
    newfile.close()   


def merge_stats(L):
    '''
    (list) -> D
    
    Returns a dictionary with read counts over all chromosomes
    
    Parameters
    ----------
    - L (list): List of dictionaries with read counts for each chromosome
    '''
    
    # create a dict to count reads over all chromosomes
    D = {"reads": 0, "assigned": 0, "assigned_empty": 0, "assigned_not_empty": 0}
    
    for i in L:
        for j in D.keys():
            D[j] += i[j]
    
    # compute unassigned read count
    D["not_assigned"] = D['reads'] - D['assigned']
    
    # add total number of reads in file
    D['total'] = L[0]['total']
    
    # add ratios
    D['percent_assigned'] = round(D['assigned'] / D['reads'] * 100, 4) if D['reads'] != 0 else 0
    D['percent_not_assigned'] = round(D['not_assigned'] / D['reads'] * 100, 4) if D['reads'] != 0 else 0
    D['percent_empty_smmips'] = round(D['assigned_empty'] / D['assigned'] * 100, 4) if D['assigned'] != 0 else 0
    
    return D


def merge_smmip_counts(L):
    '''
    (list) -> D

    Returns a dictionary with read counts supporting empty and non-empty smmips
    
    Parameters
    ----------
    - L (list): List of dictionaries with empty and not_empty read counts for each smmip
    '''
    
    # create a dict to record reads supporting empty and non-empty smmips for each smmip
    D = {}
    for i in L:
        for smmip in i:
            if smmip not in D:
                D[smmip] = {'empty': 0, 'not_empty': 0}
            D[smmip]['empty'] += i[smmip]['empty']    
            D[smmip]['not_empty'] += i[smmip]['not_empty']
    
    return D

    
def merge_bams(outputfile, L):
    '''
    (str, list) -> None
    
    Merge the input bams in list L to outputfile
    
    Parameters
    ----------
    - outputfile (str): Path to the merged bam
    - L (list): List of input bam bames
    '''
    
    # merge bam files
    args = ['-f', outputfile]
    args.extend(L)
    pysam.merge(*args)
    
    
    
