--- a
+++ b/singlecellmultiomics/utils/base_call_covariates.py
@@ -0,0 +1,655 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from multiprocessing import Pool
+from pysam import AlignmentFile, FastaFile
+import pysam
+from singlecellmultiomics.bamProcessing.bamBinCounts import blacklisted_binning_contigs
+from singlecellmultiomics.utils.sequtils import reverse_complement, get_context
+from singlecellmultiomics.utils import prob_to_phred
+from pysamiterators import CachedFasta
+from array import array
+from uuid import uuid4
+from singlecellmultiomics.bamProcessing import merge_bams, get_contigs_with_reads, has_variant_reads
+import argparse
+import pickle
+import gzip
+import pandas as pd
+import numpy as np
+import os
+from dataclasses import dataclass, field
+from collections import defaultdict,Counter
+
+def get_covariate_key(read, qpos, refpos, reference, refbase, cycle_bin_size=3, k_rad=1):
+
+    qual = read.query_qualities[qpos]
+    qbase = read.query_sequence[qpos]
+
+    if qbase == 'N':
+        return None
+
+    context = get_context(read.reference_name, refpos, reference, qbase, k_rad)
+    if 'N' in context or len(context)!=(2*k_rad+1):
+        context=get_context(read.reference_name, refpos, reference, qbase, 0)
+        assert len(context)==1
+
+    if 'N' in context:
+        return None
+
+
+    if read.is_reverse:
+        cycle = len(read.query_sequence) - qpos  # -1
+        context = reverse_complement(context)
+    else:
+        cycle = qpos + 1
+
+    return qual, read.is_read2, int(round(cycle / cycle_bin_size)) * cycle_bin_size, context
+
+
+def add_keys_excluding_context_covar(covariates, covar_phreds, k_rad=1):
+    k_rad = 1
+    for base in 'ACTG':
+        pd.DataFrame({ (key[0],key[1],key[2],key[3][k_rad]):value for key,value in covariates.items() if len(key[3])==k_rad*2+1 and key[3][k_rad]==base }).T
+
+    index = []
+    values = []
+    for key,(t,f) in [((key[0],key[1],key[2],key[3][k_rad]), value) for key,value in covariates.items() if len(key[3])==k_rad*2+1 and key[3][k_rad]==base ]:
+        index.append(key)
+        values.append([t,f])
+
+    for indices, base_data in pd.DataFrame(values,index=pd.MultiIndex.from_tuples(index)).groupby(level=(0,1,2,3)):
+        t,f = base_data.sum()
+        covar_phreds[indices] = prob_to_phred( f/(t+f) )
+
+def covariate_obs_to_phreds(covariates, k_rad):
+    covar_phreds = {}
+
+    for k,(p_base_true,p_base_false) in covariates.items():
+        covar_phreds[k] = prob_to_phred( p_base_false/(p_base_true+p_base_false) )
+
+    # Add values for context not available:
+    add_keys_excluding_context_covar(covariates, covar_phreds,k_rad=k_rad)
+
+    # Add value for when everything fails (just mean):
+    t,f = pd.DataFrame(covariates).mean(1)
+    covar_phreds[None] = prob_to_phred( f/(t+f) )
+
+    return covar_phreds
+
+
+# Molecule base information container:
+# N:reads,
+# N: IVT reactions
+#
+@dataclass
+class BaseInformation:
+    ref_base: str = None
+    majority_base: str = None
+    reads: set = field(default_factory=set)
+    qual: set = field(default_factory=list)
+    cycle_r1: set = field(default_factory=list)
+    cycle_r2: set = field(default_factory=list)
+    ivt_reactions: set = field(default_factory=set)
+
+
+# Nested counter:
+def nested_counter():
+    return defaultdict(Counter)
+
+
+def get_molecule_covariate_dict():
+    return {
+
+    }
+
+
+def get_covariate_dict():
+    return {
+        ('r1', 'qual'): defaultdict(Counter),
+        ('r2', 'qual'): defaultdict(Counter),
+        ('r1', 'cycle'): defaultdict(Counter),
+        ('r2', 'cycle'): defaultdict(Counter),
+        ('m', 'n_ivt'): defaultdict(Counter),
+        ('m', 'n_reads'): defaultdict(Counter),
+        ('m', 'mean_qual'): defaultdict(Counter),
+        ('m', 'mean_cycle_r1'): defaultdict(Counter),
+        ('m', 'mean_cycle_r2'): defaultdict(Counter),
+
+        # Within Ivt duplicate comparison:
+        ('i1', 'qual'): defaultdict(Counter),
+        ('i2', 'qual'): defaultdict(Counter),
+        ('i1', 'cycle'): defaultdict(Counter),
+        ('i2', 'cycle'): defaultdict(Counter),
+        ('i', 'd_start'): defaultdict(Counter),
+
+        # Outer IVT duplicate comparison (comparing the consensus sequences of IVT duplicates)
+        ('O1', 'n_reads'): defaultdict(Counter),
+        ('O', 'conv'): defaultdict(Counter),
+
+        # Inner IVT duplicate comparison (IVT duplicates agree)
+        ('v', 'conv'): defaultdict(Counter),
+
+    }
+
+
+def pool_wrapper(args):
+    func, kwargs = args
+    return func(**kwargs)
+
+
+def get_jobs(alignments, job_size=10_000_000, select_contig=None, **kwargs):
+    for contig, size in zip(alignments.references, alignments.lengths):
+        if select_contig is not None and contig != select_contig:
+            continue
+        for start in range(0, size, job_size):
+            end = min(size, start + job_size)
+            args = {'contig': contig,
+                    'start': start,
+                    'stop': end,
+                    }
+            args.update(kwargs)
+            yield args
+
+"""
+def extract_covariates(bam_path, contig, start, stop):
+    covariates = get_covariate_dict()
+    variant_discoveries = set()
+    no_variants = set()
+    max_show = 1
+    shown = 0
+    with pysam.AlignmentFile(bam_path) as alignments, pysam.AlignmentFile(bulk_bam_path) as bulk, pysam.FastaFile(
+            reference_path) as reference:
+        # reference = CachedFasta(ref)
+
+        read_covariates_enabled = False
+
+        reference_sequence = reference.fetch(contig, start, stop + 1).upper()
+        for i, molecule in enumerate(
+                MoleculeIterator(alignments,
+                                 molecule_class=NlaIIIMolecule,
+                                 fragment_class=NlaIIIFragment,
+                                 fragment_class_args=fragment_class_args,
+                                 contig=contig,
+                                 start=start,
+                                 stop=stop
+
+                                 )
+        ):
+
+            if len(molecule) < 4:
+                continue
+
+            found_allele = False
+            for read in molecule.iter_reads():
+                if read.has_tag('DA'):
+                    found_allele = True
+                    break
+            if not found_allele:
+                continue
+
+            information_container = defaultdict(BaseInformation)
+
+            majority = molecule.get_consensus()
+            for key, base in majority.items():
+                information_container[key].majority_base = base
+
+            ivt_base_obs = []
+            for ivt_id, fragments in molecule.get_rt_reactions().items():
+
+                # Look for inconsistencies between the fragments of the same IVT reaction
+                # Calculate IVT consensus: (Only when more fragments are available)
+                if len(fragments) > 1 and ivt_id[0] is not None:
+                    # IVT consensus:
+                    ivt_molecule = NlaIIIMolecule(fragments)
+                    ivt_base_obs.append(ivt_molecule.get_base_observation_dict())
+                    # ivt_consensus = ivt_molecule.get_consensus()
+
+                # else:
+                #    ivt_consensus = None
+
+                if read_covariates_enabled:
+                    for fragment in fragments:
+
+                        for read in fragment:
+                            if read is None:
+                                continue
+
+                            for qpos, refpos in read.get_aligned_pairs(with_seq=False):
+                                if refpos is None or refpos < start or refpos >= stop:
+                                    continue
+
+                                ref_base = reference_sequence[refpos - start]
+                                # if ref_base is None:
+                                #    continue
+                                # ref_base = ref_base.upper()
+                                if qpos is None or refpos is None or ref_base not in 'ACGT':
+                                    continue
+                                qbase = read.query_sequence[qpos]
+                                if qbase == 'N':
+                                    continue
+
+                                key = (read.reference_name, refpos)
+                                if key in known:
+                                    continue
+
+                                if ref_base != qbase and not key in no_variants:
+                                    if (read.reference_name, refpos, qbase) in variant_discoveries:
+                                        continue
+
+                                    if has_variant_reads(bulk, *key, qbase):
+
+                                        # print(f'Discovered variant at {read.reference_name}:{refpos+1} {ref_base}>{qbase}')
+                                        # known.add((read.reference_name, refpos))
+                                        variant_discoveries.add((read.reference_name, refpos, qbase))
+                                        continue
+                                    else:
+                                        no_variants.add(key)
+
+                                if ref_base is not None:
+                                    information_container[key].ref_base = ref_base
+
+                                information_container[key].reads.add(read)
+                                information_container[key].qual.append(read.query_qualities[qpos])
+
+                                cycle = read.query_length - qpos if read.is_reverse else qpos
+                                if read.is_read1:
+                                    information_container[key].cycle_r1.append(cycle)
+                                else:
+                                    information_container[key].cycle_r2.append(cycle)
+
+                                information_container[key].ivt_reactions.add(ivt_id)
+
+                                # Set read covariates:
+                                call_is_correct = ref_base == qbase
+                                covariates[f'r{"2" if read.is_read2 else "1"}', 'qual'][read.query_qualities[qpos]][
+                                    call_is_correct] += 1
+                                covariates[f'r{"2" if read.is_read2 else "1"}', 'cycle'][cycle][call_is_correct] += 1
+
+                                # Set IVT covariates
+                                if ivt_consensus is None or not key in ivt_consensus:
+                                    continue
+
+                                ivt_majority_base = ivt_consensus[key]
+                                if ivt_majority_base != ref_base or ivt_majority_base == 'N':
+                                    continue
+                                call_matches_ivt_mayority = ivt_majority_base == qbase
+                                covariates[f'i{"2" if read.is_read2 else "1"}', 'qual'][read.query_qualities[qpos]][
+                                    call_matches_ivt_mayority] += 1
+                                covariates[f'i{"2" if read.is_read2 else "1"}', 'cycle'][cycle][
+                                    call_matches_ivt_mayority] += 1
+
+                                # Calculate distance to start of molecule (in bins of 5 bp, to a maximum of 400)
+                                dstart = np.clip(int(abs(refpos - molecule.get_cut_site()[1]) / 5), 0, 400)
+                                covariates[f'i', 'd_start'][dstart][call_matches_ivt_mayority] += 1
+                                # if not call_matches_ivt_mayority and shown<max_show:
+                                #    #print(key, molecule.sample, f'{ivt_majority_base}>{qbase}' )
+                                #    shown+=1
+
+            if len(ivt_base_obs) >= 2:
+                for ivt_matches, gen_pos, base_A, base_B in get_ivt_mismatches(ivt_base_obs):
+                    # Ignore known variation:
+                    if gen_pos in variant_discoveries or gen_pos in known or has_variant_reads(bulk, *gen_pos,
+                                                                                               base_A) or has_variant_reads(
+                            bulk, *gen_pos, base_B):
+                        continue
+
+                    # Obtain reference base:
+
+                    refpos = gen_pos[1]
+                    if refpos is None or refpos < start or refpos >= stop:
+                        continue
+
+                    # ref_base = reference_sequence[refpos-start]
+                    try:
+                        context = reference_sequence[refpos - start - 1: refpos - start + 2]
+                        if 'N' in context:
+                            continue
+
+                        if molecule.strand:  # is reverse..
+                            context = reverse_complement(context)
+                            base_A = complement(base_A)
+                            base_B = complement(base_B)
+
+                        refbase = context[1]
+                    except IndexError:
+                        continue
+                    if ivt_matches:
+                        if refbase != base_A:
+                            covariates['v', 'conv'][f'{context}>{context[0]}{base_A}{context[2]}'][False] += 1
+                        # else:
+                        #    covariates['O','conv'][f'{context}>{context[0]}{base_A}{context[2]}'][True]+=1
+
+                        if refbase != base_B:
+                            covariates['v', 'conv'][f'{context}>{context[0]}{base_B}{context[2]}'][False] += 1
+
+
+                    else:
+                        # IVT not matching:
+                        print('Not matching IVT at', gen_pos)
+
+                        if refbase != base_A:
+                            covariates['O', 'conv'][f'{context}>{context[0]}{base_A}{context[2]}'][False] += 1
+                        # else:
+                        #    covariates['O','conv'][f'{context}>{context[0]}{base_A}{context[2]}'][True]+=1
+
+                        if refbase != base_B:
+                            covariates['O', 'conv'][f'{context}>{context[0]}{base_B}{context[2]}'][False] += 1
+                    # else:
+                    #    covariates['O','conv'][f'{context}>{context[0]}{base_B}{context[2]}'][True]+=1
+
+            if False:
+                # We have accumulated our molecule wisdom.
+                for location, bi in information_container.items():
+                    if location in known:
+                        continue
+
+                    if bi.ref_base is None or bi.majority_base is None or bi.ref_base not in 'ACGT':
+                        continue
+                    call_is_correct = bi.ref_base == bi.majority_base
+
+                    covariates['m', 'n_ivt'][len(bi.ivt_reactions)][call_is_correct] += 1
+                    covariates['m', 'n_reads'][len(bi.reads)][call_is_correct] += 1
+                    covariates['m', 'mean_qual'][int(np.mean(bi.qual))][call_is_correct] += 1
+                    if len(bi.cycle_r2):
+                        covariates['m', 'mean_cycle_r2'][int(np.mean(bi.cycle_r2))][call_is_correct] += 1
+                    if len(bi.cycle_r1):
+                        covariates['m', 'mean_cycle_r1'][int(np.mean(bi.cycle_r1))][call_is_correct] += 1
+
+    return covariates, variant_discoveries
+
+"""
+
+def extract_covariates(bam_path: str,
+                       reference_path: str,
+                       contig: str,
+                       start: int,
+                       end: int,
+                       start_fetch: int,
+                       end_fetch: int,
+                       filter_kwargs: dict,
+                       covariate_kwargs: dict):
+    """
+    Count mismatches and matches for similar base-calls
+
+    Returns:
+        match_mismatch(dict) : dictionary ( covariate_key: [mismatches, matches], .. )
+    """
+    # known is a set() containing locations of known variation (snps)
+    # @todo: extend to indels
+    global known  # <- Locations, set of (contig, position) tuples to ignore
+
+    joined = dict()
+
+    # Filters which select which reads are used to estimate covariates:
+    min_mapping_quality = filter_kwargs.get('min_mapping_quality', 0)
+    deduplicate = filter_kwargs.get('deduplicate', False)
+    filter_qcfailed = filter_kwargs.get('filter_qcfailed', False)
+    variant_blacklist_vcf_files = filter_kwargs.get('variant_blacklist_vcf_files', None)
+
+    # Obtain all variants in the selected range:
+    blacklist = set()
+    if variant_blacklist_vcf_files is not None:
+
+        for path in variant_blacklist_vcf_files:
+            try:
+                with pysam.VariantFile(path) as bf:
+                    for record in bf.fetch(contig, start_fetch, end_fetch):
+                        blacklist.add(record.pos)
+            except ValueError:
+                print(f'Contig {contig} is missing in the vcf file')
+
+
+    with AlignmentFile(bam_path) as alignments,  FastaFile(reference_path) as fa:
+        reference = CachedFasta(fa)  # @todo: prefetch selected region
+        for read in alignments.fetch(contig, start_fetch, end_fetch):
+            if (deduplicate and read.is_duplicate) or \
+                    (read.is_qcfail and filter_qcfailed) or \
+                    (read.mapping_quality < min_mapping_quality):
+                continue
+
+            for qpos, refpos, refbase in read.get_aligned_pairs(matches_only=True, with_seq=True):
+
+                if refpos > end or refpos < start:  # Prevent the same location to be counted multiple times
+                    continue
+
+                if refpos in blacklist:
+                    continue
+
+                refbase = refbase.upper()
+                if refbase == 'N' or (read.reference_name, refpos) in known:
+                    continue
+
+                key = get_covariate_key(read, qpos, refpos, reference, refbase, **covariate_kwargs)
+                if key is None:
+                    continue
+
+                matched = (refbase == read.query_sequence[qpos])
+                try:
+                    joined[key][matched] += 1
+                except KeyError:
+                    if matched:
+                        joined[key] = array('l', [0, 1])
+                    else:
+                        joined[key] = array('l', [1, 0])
+    return joined
+
+
+def extract_covariates_wrapper(kwargs):
+    return extract_covariates(**kwargs)
+
+
+def extract_covariates_from_bam(bam_path, reference_path, known_variants, n_processes=None, bin_size=10_000_000,
+            min_mapping_quality = 40,
+            deduplicate = True,
+            filter_qcfailed = True,
+            variant_blacklist_vcf_files = None
+            ):
+
+    global known
+    known = known_variants
+
+    joined = dict()
+
+    job_generation_args = {
+        'contig_length_resource': bam_path,
+        'bin_size': bin_size,
+        'fragment_size': 0}
+
+    filter_kwargs = {
+        'min_mapping_quality': 40,
+        'deduplicate': True,
+        'filter_qcfailed': True,
+        'variant_blacklist_vcf_files':variant_blacklist_vcf_files
+    }
+
+    covariate_kwargs = {
+        'cycle_bin_size': 3,
+        'k_rad' : 1
+    }
+    jobs_total = sum(1 for _ in (blacklisted_binning_contigs(**job_generation_args)))
+
+    with Pool(n_processes) as workers:
+
+        for i, r in enumerate(
+                        workers.imap_unordered(extract_covariates_wrapper, (
+                                   {
+                                         'bam_path': bam_path,
+                                         'reference_path': reference_path,
+                                         'contig': contig,
+                                         'start': start,
+                                         'end': end,
+                                         'start_fetch': start_fetch,
+                                         'end_fetch': end_fetch,
+                                         'filter_kwargs': filter_kwargs,
+                                         'covariate_kwargs': covariate_kwargs
+                                    }
+                                   for contig, start, end, start_fetch, end_fetch in
+                                             blacklisted_binning_contigs(**job_generation_args)))):
+            print(round(100 * (i / jobs_total), 1), end='\r')
+
+            for key, tf in r.items():
+                try:
+                    joined[key][0] += tf[0]
+                    joined[key][1] += tf[1]
+                except KeyError:
+                    joined[key] = array('l', [0, 0])
+                    joined[key][0] += tf[0]
+                    joined[key][1] += tf[1]
+    return joined
+
+
+def recalibrate_base_calls(read, reference, joined_prob, covariate_kwargs):
+    # @todo: make copy to save phred scores of soft-clipped bases
+    # This array will contain all recalibrated phred scores:
+    new_qualities = array('B', [0] * len(read.query_qualities))
+
+    # Iterate all aligned pairs and replace phred score:
+
+    for qpos, refpos in read.get_aligned_pairs(matches_only=True, with_seq=False):
+
+        key = get_covariate_key(read, qpos, refpos, reference, None, **covariate_kwargs)
+        try:
+            phred = joined_prob[key]
+        except KeyError:
+            phred = joined_prob[None]
+        new_qualities[qpos] = phred
+
+    read.query_qualities = new_qualities
+
+
+def _recalibrate_reads(bam_path, reference_path, contig, start, end, covariate_kwargs, **kwargs):
+    # Recalibrate the reads in bam_path
+
+    global joined_prob  # Global to share over multiprocessing
+    # joined_prob contains  P(error| d), where d is a descriptor generated by get_covariate_key
+
+    o_path = f'out_{uuid4()}.bam'
+
+    # Open source bam file:
+    with AlignmentFile(bam_path) as alignments, FastaFile(reference_path) as fa:
+        # @todo: extract only selected region from fasta file:
+        reference = CachedFasta(fa)
+        # Open target bam file:
+        with AlignmentFile(o_path, header=alignments.header, mode='wb') as out:
+            # Iterate all reads in the source bam file:
+            for read in alignments.fetch(contig, start, end):
+                recalibrate_base_calls(read, reference, joined_prob, covariate_kwargs)
+                out.write(read)
+
+    pysam.index(o_path)
+    return o_path
+
+
+def __recalibrate_reads(kwargs):
+    return _recalibrate_reads(**kwargs)
+
+
+def recalibrate_reads(bam_path, target_bam_path, reference_path, n_processes, covariates,  covariate_kwargs, intermediate_bam_size=20_000_000):
+    job_generation_args = {
+        'contig_length_resource': bam_path,
+        'bin_size': intermediate_bam_size,
+        'fragment_size': 0
+    }
+    global joined_prob
+    joined_prob = covariates
+
+    print(len(covariates), 'discrete elements')
+    with Pool(n_processes) as workers:
+        intermediate_bams = list( workers.imap_unordered(__recalibrate_reads, (
+                        {
+                            'bam_path': bam_path,
+                            'reference_path': reference_path,
+                            'contig': contig,
+                            'start': None,
+                            'end': None,
+                           # 'start_fetch': start_fetch,
+                            #'end_fetch': end_fetch,
+                            'covariate_kwargs': covariate_kwargs
+                        }
+                        for contig in list(get_contigs_with_reads(bam_path)))))
+                        #for contig, start, end, start_fetch, end_fetch in
+                        #blacklisted_binning_contigs(**job_generation_args))))
+
+        merge_bams(intermediate_bams, target_bam_path)
+
+
+
+
+if __name__=='__main__':
+
+    argparser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="""Obtain base calling biases using complete posterior distribution and perform corrections. Both identifying the covariates and the recallibration are multithreaded""")
+    argparser.add_argument('bamfile', metavar='bamfile', type=str)
+    argparser.add_argument('-reference', help="Path to reference fasta file used to generate the bamfile", required=True)
+    argparser.add_argument('-known', help="vcf file with known variation", required=False)
+
+
+    argparser.add_argument('-threads', help="Amount of threads to use. Uses all when not set")
+
+    argparser.add_argument(
+        '-covariates_out',
+        type=str,
+        help='Write covariates to this file, ends in .pickle.gz. ')
+
+    argparser.add_argument(
+        '-covariates_in',
+        type=str,
+        help='Read in existing covariates, ends in .pickle.gz')
+
+    argparser.add_argument(
+        '-bam_out',
+        type=str,
+        required=False,
+        help='Write corrected bam file here')
+
+    argparser.add_argument(
+        '--f',
+        action= 'store_true')
+
+    args = argparser.parse_args()
+
+    assert args.bam_out!=args.bamfile, 'The input bam file name cannot match the output bam file'
+
+    if args.covariates_in is None and args.covariates_out is not None and args.bam_out is None and os.path.exists(args.covariates_out) and not args.f:
+        print('Output covariate file already exists. Use -f to overwrite.')
+        exit()
+
+    # Set defaults  when nothing is supplied
+    if args.covariates_in is None and args.bam_out is None and args.covariates_out is None:
+        args.bam_out = args.bamfile.replace('.bam','.recall.bam')
+        if args.covariates_out is None:
+            args.covariates_out = args.bamfile.replace('.bam','.covariates.pickle.gz')
+
+
+    if args.covariates_in is not None:
+        print(f'Loading covariates from {args.covariates_in} ')
+        with gzip.open(args.covariates_in,'rb') as i:
+            covariates = pickle.load(i)
+    else:
+        covariates = extract_covariates_from_bam(args.bamfile,
+                            reference_path=args.reference,
+                               variant_blacklist_vcf_files= [] if args.known is None else [args.known],
+                                         known_variants=set()
+                           )
+        if args.covariates_out is not None:
+
+            print(f'Writing covariates to {args.covariates_out}')
+            with gzip.open(args.covariates_out,'wb') as o:
+                pickle.dump(covariates, o)
+
+    if args.bam_out is not None:
+        # Create cov probs
+
+        covar_phreds =  covariate_obs_to_phreds(covariates ,k_rad=1)
+
+        print(f'Writing corrected bam file to {args.bam_out}')
+        recalibrate_reads(bam_path=args.bamfile,
+                            target_bam_path=args.bam_out,
+                            reference_path=args.reference,
+                            n_processes=args.threads,
+                            covariates=covar_phreds,
+                            covariate_kwargs = {
+                                'cycle_bin_size': 3,
+                                'k_rad' : 1
+                            })