# -*- coding: utf-8 -*-
"""
Analysis tools used by MicrobeGPS
"""

import glob
import numpy as np
import valcov
import re
import sys
import os
import pysam
import gzip
import taxonomy

class Taxonomer:
    def __init__(self, nodes):
        self.parent = dict()
        self.level = dict()
        with open(nodes, 'r') as f:
            for line in f:
                frags = line.split('\t|\t')
                taxid = int(frags[0])
                parentid = int(frags[1])
                level = frags[2]
                self.parent[taxid] = parentid
                self.level[taxid] = level
    def get_parent(self, taxid):
        try:
            return self.parent[taxid]
        except KeyError:
            print('KeyError occured for key {} in Taxonomer.get_parent'.format(taxid))
            return 1
    def get_level(self, taxid):
        try:
            return self.level[taxid]
        except KeyError:
            print('KeyError occured for key {} in Taxonomer.get_level'.format(taxid))
            return 'no rank'
    def get_species(self, taxid):
        if self.get_level(taxid) == 'species':
            return taxid
        else:
            return self.get_species(self.get_parent(taxid))

class MappedRead:
    def __init__(self):
        self.matches = set()

class Target:
    def __init__(self, name, glen):
        self.name = name
        self.length = glen
        self.cov_homog = 1.
        self.map_qual = 0.
        self.reads = dict()
        self.unique = 0


class Reference:
    def __init__(self, name):
        self.name = name
        self.reads = 0
        self.length = 0
        self.coverage = 0.
        self.validity = 0.
        self.cov_homog = 1.
        self.map_qual = 0.
        self.targets = dict()
        self.unique = 0

class Group:
    def __init__(self):
        self.members = []
        self.usrreads = False
        self.allreads = False
class EGroup:
    def __init__(self):
        self.members = dict()
        self.reads = 0
        self.unique = 0
        self.max_val = 0
        self.cov = 0
        self.map_qual = 0
        self.cov_homog = 0

# Deprecated
def gi_from_name(name):
    """ Extract a gi, if it is hidden in the name.
        Looks for the pattern gi| followed by a number.
        Returns -1 if not found."""
    gi_start = name.find('gi|') + 3 # get the first position after gi|
    if gi_start >= 3:
        gi_end = name.find('|',gi_start)
        if gi_end == -1:
            gi_end = None
        gi = name[gi_start:gi_end]
    else:
        gi = -1
    return gi


def silent(text):
    pass


#def __calculate_mapping_score(read):
#       """ calculate mapping score for a mapped read. """
#       score = 0
#       # get mismatches from SAM tag 'MD'
#       score += sum(map(lambda x: x[1] if x[0]=='NM' else 0, read.tags ))
#       # get insertions/deletions from cigar (no clipping)
#       score += sum(map(lambda x: x[1] if 0<x[0]<3 else 0, read.cigar ))
#       return score/float(read.qlen)


def __calculate_mapping_score_no_pysam(cigar,nm,rlen):
    """ calculate mapping score for a mapped read. """
    score = 0.
    # get mismatches from SAM tag 'MD'
    score += nm
    # get all errors from cigar string
    for ln,er in re.findall('(\d+)([ID])',cigar):
        score += float(ln)
    if rlen == 0:
        print('rlen==0 in __calculate_mapping_score_no_pysam\n')
    return score/float(rlen)

###EX We only need one SAM/BAM/CRAM

def read_alignment_file_pysam(afile, status=None):
    filetype = os.path.splitext(afile)[1].lower().split('.')[1]
    if filetype == 'sam':
        return read_sam_files_no_pysam(afile)
    elif filetype == 'bam':
        samfile = pysam.AlignmentFile(afile, 'rb')
    elif filetype == 'cram':
        samfile = pysam.AlignmentFile(afile, 'rc')
    else:
        raise Exception('Unknown alignment file type: {}'.format(filetype))

    # data structure to store all mapping information
    read_table = dict()
    target_table= dict()
    # stores all reference sequence lengths
    ref_len = dict()

    for i, name in enumerate(samfile.references):
        ref_len[name] = samfile.lengths[i]

    for read in samfile.fetch(until_eof=True):
        if not read.is_unmapped:
            r_name = read.query_name
            t_name = read.reference_name
            pos = read.reference_start + 1
            cigar = read.cigarstring
            try:
                nm = read.get_tag('NM')
            except KeyError:
                continue
            r_len = read.query_length
            if r_len == 0:
                print(r_name)
            try:
                score = __calculate_mapping_score_no_pysam(cigar,nm,r_len)
            except ZeroDivisionError:
                continue
            rd = read_table.setdefault(r_name,MappedRead())
            rd.matches.add(t_name)
            gn = target_table.setdefault(t_name, Target(t_name,ref_len[t_name]))
            gn.reads[r_name]= [pos, r_len, score]
    return target_table,read_table

def read_sam_files_no_pysam(files, status=None):
    """ Extract mapping information from SAM files. 'files' can be either
    a string (interpreted as single file name or glob pattern) or a list of
    strings (interpreted as list of file names).
    """
    if type(files) == str:
        fns = glob.glob(files)
    elif type(files) == list:
        fns = files
    else:
        raise Exception('Could not identify format of file names:',files)
    d = len(fns)

    data_length = sum([os.path.getsize(f) for f in fns])
    data_position = 0

    # data structure to store all mapping information
    read_table = dict()
    target_table = dict()

    # read SAM files one by one
    for i,fn in enumerate(fns):
        if status: status(i+1,d)
        sf = open(fn,'rU')
        ref_len = dict() # stores all reference sequence lengths
        # first read the header of the samfile
        for line in sf:
            data_position += len(line)
            if line[0] == '@':
                frag = line.rstrip().split('\t')
                if not frag[0] == '@SQ':
                    continue
                # assume NM tag in first field
                name = frag[1].split(':')[1]
                # assume sequence length tag LN in second field
                glen = int(frag[2].split(':')[1])
                ref_len[name] = glen
            else:
                break
        if line[0] == '@':
            continue # there are no alignments in this file

        # process the one line that we read too much
        frag = line.split()
        if int(frag[1])&4:
            pass # the unmapped flag is set
        else:
            r_name = frag[0]
            t_name = frag[2]
            pos = int(frag[3])-1
            cigar = frag[5]
            nm = 0
            for i,tag in enumerate(frag[9:]):
                if tag.startswith('NM:i:'):
                    nm = int(tag[5:])
                    break
            r_len = len(frag[9])
            score = __calculate_mapping_score_no_pysam(cigar,nm,r_len)
            rd = read_table.setdefault(r_name,MappedRead()) # search dict for r_name and reutrn value, if not found return MappedRead()
            rd.matches.add(t_name)
            if not t_name in target_table:
                target_table[t_name] = Target(t_name,ref_len[t_name])
            gn = target_table[t_name]
            gn.reads[r_name]= [pos, r_len, score]

        last_status = 0
        # process all other alignments
        try:
            for line in sf:
                data_position += len(line)
                if status:
                    last_status += 1
                    if last_status >= 5000: # update every 5000 reads
                        status((data_position*100)/data_length,100)
                        last_status = 0
                frag = line.rstrip().split('\t')
                if int(frag[1])&4:
                    continue # the unmapped flag is set
                else:
                    r_name = frag[0]
                    t_name = frag[2]
                    pos = int(frag[3])-1
                    cigar = frag[5]
                    nm = 0
                    for i,tag in enumerate(frag[9:]):
                        if tag.startswith('NM:i:'):
                            nm = int(tag[5:])
                            break
                    r_len = len(frag[9])
                    score = __calculate_mapping_score_no_pysam(cigar,nm,r_len)
                    rd = read_table.setdefault(r_name,MappedRead())
                    rd.matches.add(t_name)
                    gn = target_table.setdefault(t_name, Target(t_name,ref_len[t_name]))
                    gn.reads[r_name]= [pos, r_len, score]
        except IndexError:
            continue
    if status: status(100,100)
    return target_table,read_table


def filter_raw(target_table,read_table,max_matches=None,min_support=1, max_error=None, qual_percentile=None, pr=silent):
    """ Filter the raw mapping information according to the following criteria:
            max_matches: Reads matching to more than max_matches genomes are discarded
            min_support: Genomes with less than min_support reads are discarded
            max_error: Read alignments with more than max_error are discarded
    """
    # discard reads with more than max_matches matches and count unique reads
    unique_reads = 0
    del_reads = []
    if max_matches:
        for mp in read_table:
            lrt = len(read_table[mp].matches)
            if lrt > max_matches:
                del_reads.append(mp)
            if lrt == 1:
                target_table[next(iter(read_table[mp].matches))].unique += 1
                unique_reads += 1
    for d in del_reads:
        for gn in read_table[d].matches:
            del target_table[gn].reads[d]
        del read_table[d]
    pr("--- Discarded %i reads with > %i matches."%(len(del_reads),int(max_matches)))
    pr("--- Found %i reads with unique matches."%(unique_reads))

    # discard genomes with insufficient support
    del_genomes = []
    if min_support:
        for gn in target_table:
            if len(target_table[gn].reads) < min_support:
                del_genomes.append(gn)
    for g in del_genomes:
        for rd in target_table[g].reads:
            read_table[rd].matches.remove(g)
            if len(read_table[rd].matches) == 0:
                del read_table[rd]
        del target_table[g]
    pr("--- Discarded %i targets with less than %i reads."%(len(del_genomes),int(min_support)))


def calculate_mapping_statistics(target_table):
    """ calculate coverage homogeneity and average mapping error """
    for gn in target_table:
        start_pos = []
        scores = []
        for r in target_table[gn].reads:
            start_pos.append(target_table[gn].reads[r][0])
            scores.append(target_table[gn].reads[r][2])
        if scores == []:
            print('scores==[] in calculate_mapping_statistics')
        target_table[gn].map_qual = np.mean(scores)

        start_pos.sort()
        g_len = float(target_table[gn].length)
        start_pos.append(g_len)
        n_rds = float(len(start_pos))

        max_dev = 0.
        for i,p in enumerate(start_pos):
            dev1 = abs(((p-1)/g_len) - (i/n_rds))
            dev2 = abs((p/g_len) - ((i+1)/n_rds))
            max_dev = max(max_dev,dev1,dev2)
        target_table[gn].cov_homog = max_dev



def get_reference_table(target_table, read_table, gmap=None):
    """ Put targets in biologically meaningful groups.
    'gmap' can be either a dict, mapping target names to genome names, or a
    file name pointing to a tab delimited file."""
    # first create the mapping function: target name --> genome name
    if not gmap:
        # use the identity mapping
        mp = lambda x: x
    elif hasattr(gmap,'__call__'):
        # gmap is a function. Expect that gmap takes only a reference name as input
        mp = gmap
    else:
        if type(gmap) == str:
            d = dict()
            f = open(gmap)
            for line in f:
                contig_name,ref_name = line.rstrip().split('\t')
                d[contig_name] = ref_name
        elif type(gmap) == dict:
            d = gmap
        else:
            raise Exception('gmap type must be either dict or string (filename).')
        mp = lambda x: d.get(x,x)

    # create the reference table
    ref_table = dict()
    for t in target_table:
        r = mp(t)
        ref = ref_table.setdefault(r,Reference(r))
        ref.targets[t] = target_table[t]

    # update read_table, assign reads to References, not Targets
    for rd in read_table:
        renamed_set = set()
        for match in read_table[rd].matches:
            renamed_set.add(mp(match))
        read_table[rd].matches = renamed_set

    # calculate simple metrics over references
    for r in ref_table:
        lengths = [x.length for x in iter(ref_table[r].targets.values())]
        homog = [x.cov_homog for x in iter(ref_table[r].targets.values())]
        qual = [x.map_qual for x in iter(ref_table[r].targets.values())]
        unique = sum([x.unique for x in iter(ref_table[r].targets.values())])
        t_len = float(sum(lengths))
        if t_len == 0:
            print('t_len==0 in calculate_reference_table')

        readnames = set()
        for t in iter(ref_table[r].targets.values()):
            readnames.update(iter(t.reads.keys()))
        reads = len(readnames)

        t_homog = 0.
        t_qual = 0.
        for l,h,q in zip(lengths,homog,qual):
            t_homog += h*l/t_len
            t_qual += q*l/t_len
        ref_table[r].cov_homog = t_homog
        ref_table[r].map_qual = t_qual
        ref_table[r].unique = unique
        ref_table[r].reads = reads
        ref_table[r].length = sum(lengths)

    return ref_table


def get_reference_table_NCBI(target_table, read_table, catalog, pr=silent):
    """ Put targets in biologically meaningful groups using NCBI taxonomy.
    catalog specifies the file name of the RefSeq-releaseXX.catalog file
    """

    # first collect all accession.versions
    acs = set()
    for t in target_table:
        acs.add(t)

    # read the NCBI files
    ac_to_taxid,ac_to_name = taxonomy.acs_to_taxid_and_name(acs,catalog)

    # create the reference table
    ref_table = dict()
    trg_to_ref = dict()
    not_found = 0
    taxi = Taxonomer(os.path.realpath(os.path.join(os.path.dirname(sys.argv[0]), 'data/nodes.dmp')))
    for t in target_table:
        #gi = gi_from_name(t)
        #taxid = gi_to_taxid.get(gi,'[not found]')
        taxid = ac_to_taxid.get(t,'[not found]')
        if taxid == '[not found]':
            not_found += 1
            ref_name = '{}||{}||{}||{}||{}'.format(ac_to_name.get(t, 'NA'), str(t), '[not found]', '[not found]', '[not found]') 
            ref = ref_table.setdefault(ref_name,Reference(taxid))
            ref.targets[t] = target_table[t]
            trg_to_ref[t] = ref_name
        else:
            ref_name = '{}||{}||{}||{}||{}'.format(ac_to_name.get(t, 'NA'), str(t), str(taxid), str(taxi.get_parent(taxid)), str(taxi.get_species(taxid))) 
            ref = ref_table.setdefault(ref_name,Reference(taxid))
            ref.targets[t] = target_table[t]
            trg_to_ref[t] = ref_name
    pr('--- Found %i IDs for %i targets'%(len(target_table)-not_found,len(target_table)))

    # update read_table, assign reads to References, not Targets
    for rd in read_table:
        renamed_set = set()
        for match in read_table[rd].matches:
            renamed_set.add(trg_to_ref[match])
        read_table[rd].matches = renamed_set

    # calculate simple metrics over references
    for r in ref_table:
        lengths = [x.length for x in iter(ref_table[r].targets.values())]
        homog = [x.cov_homog for x in iter(ref_table[r].targets.values())]
        qual = [x.map_qual for x in iter(ref_table[r].targets.values())]
        unique = sum([x.unique for x in iter(ref_table[r].targets.values())])
        t_len = float(sum(lengths))

        readnames = set()
        for t in iter(ref_table[r].targets.values()):
            readnames.update(iter(t.reads.keys()))
        reads = len(readnames)

        t_homog = 0.
        t_qual = 0.
        for l,h,q in zip(lengths,homog,qual):
            t_homog += h*l/t_len
            t_qual += q*l/t_len
        ref_table[r].cov_homog = t_homog
        ref_table[r].map_qual = t_qual
        ref_table[r].unique = unique
        ref_table[r].reads = reads
        ref_table[r].length = sum(lengths)

    return ref_table



def filter_ref_table(ref_table,read_table,filt=lambda x:True, pr=silent):
    """Delete entries in ref_table for which the filter function filt returns
    False. The unary function filt must be able to handle Reference objects."""
    del_ref = []
    # create the set of genomes to be removed
    for rf in ref_table:
        if not filt(ref_table[rf]):
            del_ref.append(rf)
    pr('--- Removing %i of %i references.'%(len(del_ref),len(ref_table)))
    # remove the selected genomes
    for d in del_ref:
        for t in ref_table[d].targets:
            for rd in ref_table[d].targets[t].reads:
                if not rd in read_table:
                    continue  # this happens if one read matched to 2 targets in one reference --> read_table[rd] was already removed when reaching the second target
                read_table[rd].matches.discard(d)
                if len(read_table[rd].matches) == 0:
                    del read_table[rd]
        del ref_table[d]
    # update unique reads information
    pr('--- Updating Unique Reads information.')
    for rf in ref_table:
        ref_table[rf].unique = 0
    for rd in iter(read_table.values()):
        if len(rd.matches) == 1:
            ref_table[next(iter(rd.matches))].unique += 1



def calculate_valcov(ref_table, status=None):
    """Calculate validity and coverage for all References in ref_table."""
    count = 0
    for rf in ref_table:
        count += 1
        if status: status('Calculate val/cov (%i/%i): %s'%(count,len(ref_table),rf))
        dsc = valcov.DS_cov(ref_table[rf])
        val,cov = valcov.validity_from_coverage(dsc)
        if cov < 0.8:
            # use distance based method if coverage is too low
            dsd = valcov.DS_dst(ref_table[rf])
            if len(dsd.value) > 0:
                val,cov = valcov.validity_from_spaces(dsd)
            else:
                val,cov = 0.,0.
        ref_table[rf].coverage = cov
        ref_table[rf].validity = val

def calculate_valcov_one(ref):
    """Calculate validity and coverage for one reference."""
    if ref.coverage > 0. and ref.validity > 0.:
        # validity and coverage were already calculated
        return ref.validity,ref.coverage
    dsc = valcov.DS_cov(ref)
    val,cov = valcov.validity_from_coverage(dsc)
    if cov < 0.8:
        # use distance based method if coverage is too low
        dsd = valcov.DS_dst(ref)
        if len(dsd.value) > 0:
            val,cov = valcov.validity_from_spaces(dsd)
        else:
            val,cov = 0.,0.
    ref.coverage = cov
    ref.validity = val
    return val,cov


def extract_USR(ref_table, read_table, csim=0.2, pr=silent):
    """extract Unique Source Reads from the read_table, i.e. reads mapping to
    genomes with comparable coverage."""
    usr_table = dict()
    for rd in read_table:
        covs = np.array([ref_table[mt].coverage for mt in read_table[rd].matches])
        #vals = np.array([ref_table[mt].validity for mt in read_table[rd].matches])
        # hier noch die validity mit einbauen!!!
        diff = (np.max(covs)-np.min(covs))/(np.sum(covs)/len(covs))
        if diff < csim:
            usr_table[rd] = read_table[rd]
    pr('--- Found %i Unique Source Reads (total %i reads)'%(len(usr_table),len(read_table)))
    return usr_table


def get_read_matrix(reads, targets):
    """Create a numpy matrix assigning reads to targets (or references).
    Returns the matrix and dictionaries mapping target name -> index and
 index -> target name."""
    mat = np.zeros((len(reads),len(targets)),dtype=np.uint8)
    name_to_index = dict()
    index_to_name = dict()
    for i,t in enumerate(targets):
        name_to_index[t] = i
        index_to_name[i] = t
    for i,r in enumerate(reads):
        for t in reads[r].matches:
            mat[i,name_to_index[t]] = 1
    return mat, name_to_index, index_to_name


def create_groups(usrmat, thr=0.1, status=None):
    """create groups of references from matrix mat.
    Returns a dict with lists of reference IDs as enties (=columns of mat) and
    a dict of weak links between genomes."""

    groups = dict()
    rds,rfs = usrmat.shape

    for i in range(rfs):
        if status: status('--- Grouping reference %i of %i'%(i+1,rfs))
        # calculate shared reads with existing groups
        mapped = np.sum(usrmat[:,i])
        shared = {gr:np.sum(np.logical_and(groups[gr].usrreads,usrmat[:,i])) for gr in groups}
        sig = {gr:shared[gr] for gr in shared if shared[gr]>mapped*thr}
        N = len(sig)
        #print('Current reference: {} Number of groups: {}'.format(i, N))
        if N == 0:
            # no significant similarity to existing group --> create new group
            gr = Group()
            gr.members.append(i)
            gr.usrreads = np.logical_or(gr.usrreads,usrmat[:,i])
            grID = len(groups)
            groups[grID] = gr
        elif N == 1:
            # significant similarity with exactly 1 group --> assign to that group
            grID = sig.keys()[0]
            groups[grID].members.append(i)
            groups[grID].usrreads = np.logical_or(groups[grID].usrreads,usrmat[:,i])
        else:
            # significant similarity with > 1 groups: assign to first group
            grID = sig.keys()[0]
            groups[grID].members.append(i)
            groups[grID].usrreads = np.logical_or(groups[grID].usrreads,usrmat[:,i])
            # ... and add other groups
            for otherID in sig.keys()[1:]:
                for m in groups[otherID].members:
                    groups[grID].members.append(m)
                groups[grID].usrreads = np.logical_or(groups[grID].usrreads,groups[otherID].usrreads)
                groups[otherID].usrreads = False
                groups[otherID].members = []
    g_dict = {g:groups[g].members for g in groups if len(groups[g].members)>0}

    return g_dict


def create_groups_2(usrmat, allmat, thr=0.1, thr_all=0.6, status=None):
    """create groups of references from matrix mat.
    Returns a dict with lists of reference IDs as enties (=columns of mat)"""

    groups = dict()
    rds,rfs = usrmat.shape

    for i in range(rfs):
        if status: status('--- Grouping reference %i of %i'%(i+1,rfs))
        #print('Current reference:{} Number of groups: {}'.format(i, len(groups)))
        sig = dict()

        usr_reads = np.sum(usrmat[:,i])
        all_reads = np.sum(allmat[:,i])
        for gr in groups:
            if type(groups[gr].usrreads) == bool:
                continue # skip groups without usrreads/allreads information
            usr_shared = np.sum(np.logical_and(groups[gr].usrreads,usrmat[:,i]))
            all_shared = np.sum(np.logical_and(groups[gr].allreads,allmat[:,i]))
            gr_usr_reads = np.sum(groups[gr].usrreads)
            gr_all_reads = np.sum(groups[gr].allreads)
            if usr_shared > thr*usr_reads or usr_shared > thr*gr_usr_reads:
                sig[gr] = usr_shared
            if all_shared > thr_all*all_reads or all_shared > thr_all*gr_all_reads:
                sig[gr] = usr_shared

        N = len(sig)
        if N == 0:
            # no significant similarity to existing group --> create new group
            gr = Group()
            gr.members.append(i)
            gr.usrreads = np.logical_or(gr.usrreads,usrmat[:,i])
            gr.allreads = np.logical_or(gr.allreads,allmat[:,i])
            grID = len(groups)
            groups[grID] = gr
        elif N == 1:
            # significant similarity with exactly 1 group --> assign to that group
            grID = sig.keys()[0]
            groups[grID].members.append(i)
            groups[grID].usrreads = np.logical_or(groups[grID].usrreads,usrmat[:,i])
            groups[grID].allreads = np.logical_or(groups[grID].allreads,allmat[:,i])
        else:
            # significant similarity with > 1 groups: assign to best group
            grID = sig.keys()[np.argmax(sig.values())]
            groups[grID].members.append(i)
            groups[grID].usrreads = np.logical_or(groups[grID].usrreads,usrmat[:,i])
            groups[grID].allreads = np.logical_or(groups[grID].allreads,allmat[:,i])

    g_dict = {g:groups[g].members for g in groups if len(groups[g].members)>0}

    return g_dict

def create_groups_dc(usrmat, allmat, n2i, thr=0.1, thr_all=0.6, status=None):
    """create groups of references from matrix mat using a divide&conquer strategy
    Returns a dict with lists of reference IDs as enties (=columns of mat)"""

    groups = dict()
    rds,rfs = usrmat.shape

    # Two stage strategy
    # First stage: search for groups locally among all references sharing the first 3 letters
    if status: status('--- Clustering stage 1')
    level1 = dict()
    for name in n2i:
        item = level1.setdefault(name[0:3], list())
        item.append(n2i[name])
    level2 = dict()
    for b,bucket in enumerate(iter(level1.values())):
        l1_groups = dict()
        for i in bucket:
            sig = dict()

            usr_reads = np.sum(usrmat[:,i])
            all_reads = np.sum(allmat[:,i])
            for gr in l1_groups:
                if type(l1_groups[gr].usrreads) == bool:
                    continue # skip l1_groups without usrreads/allreads information
                usr_shared = np.sum(np.logical_and(l1_groups[gr].usrreads,usrmat[:,i]))
                all_shared = np.sum(np.logical_and(l1_groups[gr].allreads,allmat[:,i]))
                gr_usr_reads = np.sum(l1_groups[gr].usrreads)
                gr_all_reads = np.sum(l1_groups[gr].allreads)
                if usr_shared > thr*usr_reads or usr_shared > thr*gr_usr_reads:
                    sig[gr] = usr_shared
                if all_shared > thr_all*all_reads or all_shared > thr_all*gr_all_reads:
                    sig[gr] = usr_shared

            N = len(sig)
            if N == 0:
                # no significant similarity to existing group --> create new group
                gr = Group()
                gr.members.append(i)
                gr.usrreads = np.logical_or(gr.usrreads,usrmat[:,i])
                gr.allreads = np.logical_or(gr.allreads,allmat[:,i])
                l1ID = 'bucket%i_%i'%(b,len(l1_groups))
                l1_groups[l1ID] = gr
            else:
                # significant similarity with > 1 groups: assign to best group
                l1ID = list(sig)[np.argmax(sig.values())]
                l1_groups[l1ID].members.append(i)
                l1_groups[l1ID].usrreads = np.logical_or(l1_groups[l1ID].usrreads,usrmat[:,i])
                l1_groups[l1ID].allreads = np.logical_or(l1_groups[l1ID].allreads,allmat[:,i])
        level2.update(l1_groups)

    # Second stage: merge all level 1 groups that suffice the conditions
    if status: status('--- Clustering stage 2')
    for l1ID in level2:
        sig = dict()

        usr_reads = np.sum(level2[l1ID].usrreads)
        all_reads = np.sum(level2[l1ID].allreads)
        for gr in groups:
            if type(groups[gr].usrreads) == bool:
                continue # skip groups without usrreads/allreads information
            usr_shared = np.sum(np.logical_and(groups[gr].usrreads,level2[l1ID].usrreads))
            all_shared = np.sum(np.logical_and(groups[gr].allreads,level2[l1ID].allreads))
            gr_usr_reads = np.sum(groups[gr].usrreads)
            gr_all_reads = np.sum(groups[gr].allreads)
            if usr_shared > thr*usr_reads or usr_shared > thr*gr_usr_reads:
                sig[gr] = usr_shared
            if all_shared > thr_all*all_reads or all_shared > thr_all*gr_all_reads:
                sig[gr] = usr_shared

        N = len(sig)
        if N == 0:
            # no significant similarity to existing group --> create new group
            grID = len(groups)
            groups[grID] = level2[l1ID]
        else:
            # significant similarity with > 1 groups: assign to best group
            grID = list(sig)[np.argmax(sig.values())]
            groups[grID].members.extend(level2[l1ID].members)
            groups[grID].usrreads = np.logical_or(groups[grID].usrreads,level2[l1ID].usrreads)
            groups[grID].allreads = np.logical_or(groups[grID].allreads,level2[l1ID].allreads)

    g_dict = {g:groups[g].members for g in groups if len(groups[g].members)>0}

    return g_dict



def enrich_groups(groups,references,reads,i2n):
    """ create list containing groups enriched with names and reference information """
    egroups = dict()
    for g in groups:
        egroups[g] = EGroup()
        readnames = set()
        t_len = 0.
        for m in groups[g]:
            egroups[g].members[i2n[m]] = references[i2n[m]]
            for t in references[i2n[m]].targets:
                readnames.update(iter(references[i2n[m]].targets[t].reads.keys()))
            t_len += float(references[i2n[m]].length)

        egroups[g].reads = len(readnames)
        egroups[g].max_val = max([m.validity for m in iter(egroups[g].members.values())])
        egroups[g].cov = sum([m.coverage*m.length/t_len for m in iter(egroups[g].members.values())])
        egroups[g].cov_homog = sum([m.cov_homog*m.length/t_len for m in iter(egroups[g].members.values())])
        egroups[g].map_qual = sum([m.map_qual*m.length/t_len for m in iter(egroups[g].members.values())])

        # count the group-unique reads, i.e. reads mapping only to that group
        group_unique = set()
        for mbm in iter(egroups[g].members.values()):
            for trg in iter(mbm.targets.values()):
                for rd in trg.reads:
                    if rd in group_unique:
                        continue
                    add = True
                    # check if a read has other matches than in this group
                    for mtc in reads[rd].matches:
                        if not mtc in egroups[g].members:
                            add = False
                            break
                    if add: group_unique.add(rd)

        egroups[g].unique = len(group_unique)

    sgroups = sorted([egroups[g] for g in egroups], key=lambda x: -x.unique)
    return sgroups
