<commit_before>#!/usr/bin/env python
from __future__ import absolute_import, print_function

import pysam
import cyvcf2
import os
import sys
import collections
import re
from unidecode import unidecode
from .config import read_gemini_config
from . import gemini_utils as util
try:
    basestring
except NameError:
    basestring = str

# dictionary of anno_type -> open Tabix file handles
annos = {}

def get_anno_files(args):
    config = read_gemini_config(args=args)
    anno_dirname = config["annotation_dir"]
    # Default annotations -- always found
    annos = {
     'pfam_domain': os.path.join(anno_dirname, 'hg19.pfam.ucscgenes.bed.gz'),
     'cytoband': os.path.join(anno_dirname, 'hg19.cytoband.bed.gz'),
     'dbsnp': os.path.join(anno_dirname, 'dbsnp.151.20180423.tidy.vcf.gz'),
     'clinvar': os.path.join(anno_dirname, 'clinvar_20190102.tidy.vcf.gz'),
     'gwas': os.path.join(anno_dirname, 'hg19.gwas.bed.gz'),
     'rmsk': os.path.join(anno_dirname, 'hg19.rmsk.bed.gz'),
     'segdup': os.path.join(anno_dirname, 'hg19.segdup.bed.gz'),
     'conserved': os.path.join(anno_dirname, '29way_pi_lods_elements_12mers.chr_specific.fdr_0.1_with_scores.txt.hg19.merged.bed.gz'),
     'cpg_island': os.path.join(anno_dirname, 'hg19.CpG.bed.gz'),
     'dgv': os.path.join(anno_dirname, 'hg19.dgv.bed.gz'),
     'esp': os.path.join(anno_dirname,
                         'ESP6500SI.all.snps_indels.tidy.v2.vcf.gz'),
     '1000g': os.path.join(anno_dirname,
                           'ALL.wgs.phase3_shapeit2_mvncall_integrated_v5a.20130502.sites.tidy.vcf.gz'),
     'recomb': os.path.join(anno_dirname,
                            'genetic_map_HapMapII_GRCh37.gz'),
     'gms': os.path.join(anno_dirname,
                         'GRCh37-gms-mappability.vcf.gz'),
     'grc': os.path.join(anno_dirname, 'GRC_patch_regions.bed.gz'),
     'cse': os.path.join(anno_dirname, "cse-hiseq-8_4-2013-02-20.bed.gz"),
     'encode_tfbs': os.path.join(anno_dirname,
                                 'wgEncodeRegTfbsClusteredV2.cell_count.20130213.bed.gz'),
     'encode_dnase1': os.path.join(anno_dirname,
                                   'stam.125cells.dnaseI.hg19.bed.gz'),
     'encode_consensus_segs': os.path.join(anno_dirname,
                                           'encode.6celltypes.consensus.bedg.gz'),
     'gerp_elements': os.path.join(anno_dirname, 'hg19.gerp.elements.bed.gz'),
     'vista_enhancers': os.path.join(anno_dirname, 'hg19.vista.enhancers.20131108.bed.gz'),
     'fitcons': os.path.join(anno_dirname, "hg19_fitcons_fc-i6-0_V1-01.bed.gz"),
     'cosmic': os.path.join(anno_dirname, 'cosmic-v68-GRCh37.tidy.vcf.gz'),
     'exac': os.path.join(anno_dirname, 'ExAC.r0.3.sites.vep.tidy.vcf.gz'),
     'gnomad': os.path.join(anno_dirname, 'gnomad.exomes.r2.1.tidy.bcf'),
     'geno2mp': os.path.join(anno_dirname, 'geno2mp.variants.tidy.vcf.gz'),
    }
    # optional annotations
    if os.path.exists(os.path.join(anno_dirname, 'hg19.gerp.bw')):
        annos['gerp_bp'] = os.path.join(anno_dirname, 'hg19.gerp.bw')
    if os.path.exists(os.path.join(anno_dirname, 'caddv1.4.compressed.gz')):
        annos['cadd_score'] = os.path.join(anno_dirname, 'caddv1.4.compressed.gz')
    return annos

class ClinVarInfo(object):
    def __init__(self):
        self.clinvar_origin = None
        self.clinvar_sig = None
        self.clinvar_dsdb = None
        self.clinvar_dsdbid = None
        self.clinvar_disease_name = None
        self.clinvar_in_omim = None

        self.origin_code_map = {'0': 'unknown',
                                '1': 'germline',
                                '2': 'somatic',
                                '4': 'inherited',
                                '8': 'paternal',
                                '16': 'maternal',
                                '32': 'de-novo',
                                '64': 'biparental',
                                '128': 'uniparental',
                                '256': 'not-tested',
                                '512': 'tested-inconclusive',
                                '1073741824': 'other'}
        # 0 - Uncertain significance, 1 - not provided, 2 - Benign, 3 - Likely benign, 4 - Likely pathogenic, 5 - Pathogenic, 6 - drug response, 7 - histocompatibility, 255 - other

        self.sig_code_map = {'0': 'uncertain',
                             '1': 'not-provided',
                             '2': 'benign',
                             '3': 'likely-benign',
                             '4': 'likely-pathogenic',
                             '5': 'pathogenic',
                             '6': 'drug-response',
                             '7': 'histocompatibility',
                             '255': 'other'}

    def __repr__(self):
        return '\t'.join(map(str, [
                          self.clinvar_origin,
                          self.clinvar_sig,
                          self.clinvar_dsdb,
                          self.clinvar_dsdbid,
                          self.clinvar_disease_name,
                          str(self.clinvar_in_omim),
                          str(self.clinvar_causal_allele)]))

    def lookup_clinvar_origin(self, origin_code):
        try:
            return self.origin_code_map[origin_code]
        except KeyError:
            return None

    def lookup_clinvar_significance(self, sig_code):
        sigs = []
        for s in sig_code.split('|'):
            sigs.extend(s.split(","))

        return ",".join(self.sig_code_map[s] for s in set(sigs) if s != ".")


ESPInfo = collections.namedtuple("ESPInfo",
                                 "found \
                                  aaf_EA \
                                  aaf_AA \
                                  aaf_ALL \
                                  exome_chip")
ENCODEDnaseIClusters = collections.namedtuple("ENCODEDnaseIClusters",
                                              "cell_count \
                                         cell_list")
ENCODESegInfo = collections.namedtuple("ENCODESegInfo",
                                       "gm12878 \
                                         h1hesc \
                                         helas3 \
                                         hepg2 \
                                         huvec \
                                         k562")
ThousandGInfo = collections.namedtuple("ThousandGInfo",
                                       "found \
                                        aaf_ALL \
                                        aaf_AMR \
                                        aaf_EAS \
                                        aaf_SAS \
                                        aaf_AFR \
                                        aaf_EUR")

GmsTechs = collections.namedtuple("GmsTechs", "illumina solid iontorrent")

ExacInfo = collections.namedtuple("ExacInfo",
                                  "found \
                                   aaf_ALL \
                                   adj_aaf_ALL \
                                   aaf_AFR \
                                   aaf_AMR \
                                   aaf_EAS \
                                   aaf_FIN \
                                   aaf_NFE \
                                   aaf_OTH \
                                   aaf_SAS \
                                   num_het \
                                   num_hom_alt \
                                   num_chroms")

EXAC_EMPTY = ExacInfo(False, -1, -1, -1, -1, -1,
                     -1, -1, -1, -1, -1, -1, -1)

GnomadInfo = collections.namedtuple('GnomadInfo',
                                   "aaf_ALL \
                                   aaf_AFR \
                                   aaf_AMR \
                                   aaf_ASJ \
                                   aaf_EAS \
                                   aaf_FIN \
                                   aaf_NFE \
                                   aaf_OTH \
                                   aaf_SAS \
                                   aaf_non_neuro \
                                   aaf_controls \
                                   aaf_non_cancer \
                                   popmax_AF \
                                   num_het \
                                   num_hom_alt \
                                   num_chroms")

GNOMAD_EMPTY = GnomadInfo(-1, -1, -1, -1, -1, -1, -1, -1,
                          -1, -1, -1, -1, -1, -1, -1, -1)

def load_annos(args):
    """
    Populate a dictionary of Tabixfile handles for
    each annotation file.  Other modules can then
    access a given handle and fetch data from it
    as follows:

    dbsnp_handle = annotations.annos['dbsnp']
    hits = dbsnp_handle.fetch(chrom, start, end)
    """
    anno_files = get_anno_files(args)
    for anno in anno_files:
        try:
            # .gz denotes Tabix files.
            if anno_files[anno].endswith(".gz"):
                if anno == "clinvar":
                    annos[anno] = pysam.Tabixfile(anno_files[anno],
                                                  encoding='utf8')
                else:
                    annos[anno] = pysam.Tabixfile(anno_files[anno])

            elif anno_files[anno].endswith(".bcf"):
                annos[anno] = cyvcf2.VCF(anno_files[anno])
            # .bw denotes BigWig files.
            elif anno_files[anno].endswith(".bw"):
                from bx.bbi.bigwig_file import BigWigFile
                annos[anno] = BigWigFile(open(anno_files[anno]))

        except IOError:
            raise IOError("Gemini cannot open this annotation file: %s. \n"
                          "Have you installed the annotation files?  If so, "
                          "have they been moved or deleted? Exiting...\n\n"
                          "For more details:\n\t"
                          "http://gemini.readthedocs.org/en/latest/content/"
                          "#installation.html\#installing-annotation-files\n"
                          % anno_files[anno])

# ## Standard access to Tabix indexed files


PARSERS = {"bed": pysam.asBed(),
           "vcf": pysam.asVCF(),
           "tuple": pysam.asTuple(),
           None: None}

def _get_hits(coords, annotation, parser_type, _parsers=PARSERS):
    """Retrieve BED information, recovering if BED annotation file does have a chromosome.
    """
    try:
        parser = _parsers[parser_type]
    except KeyError:
        raise ValueError("Unexpected parser type: %s" % parser)
    chrom, start, end = coords
    if isinstance(annotation, pysam.VariantFile):
        return annotation.fetch(chrom, start, end)
    elif isinstance(annotation, cyvcf2.VCF):
        return annotation("%s:%d-%d" % (chrom, start - 1, end))
    try:
        hit_iter = annotation.fetch(str(chrom), start, end, parser=parser)
    # catch invalid region errors raised by ctabix
    except ValueError:
        hit_iter = []
    # recent versions of pysam return KeyError
    except KeyError:
        hit_iter = []
    except:
        print(annotation.__class__, file=sys.stderr)
        raise
    return hit_iter

def _get_bw_summary(coords, annotation):
    """Return summary of BigWig scores in an interval
    """
    chrom, start, end = coords
    try:
        return annotation.summarize(str(chrom), start, end, end-start).min_val[0]
    except AttributeError:
        return None


def _get_chr_as_grch37(chrom):
    if chrom in ["chrM"]:
        return "MT"
    return chrom if not chrom.startswith("chr") else chrom[3:]


def _get_chr_as_ucsc(chrom):
    return chrom if chrom.startswith("chr") else "chr" + chrom


def guess_contig_naming(anno):
    """Guess which contig naming scheme a given annotation file uses.
    """
    chr_names = [x for x in anno.contigs if x.startswith("chr")]
    if len(chr_names) > 0:
        return "ucsc"
    else:
        return "grch37"


def _get_var_coords(var, naming):
    """Retrieve variant coordinates from multiple input objects.
    """
    if isinstance(var, cyvcf2.Variant):
      chrom, start, end = var.CHROM, var.start, var.end
    else:
        try:
            # todo: check isinstance resultproxy?
            chrom = var["chrom"]
            start = int(var["start"])
            end = int(var["end"])
        except TypeError:
            chrom = var.CHROM
            start = var.start
            end = var.end
    if naming == "ucsc":
        chrom = _get_chr_as_ucsc(chrom)
    elif naming == "grch37":
        chrom = _get_chr_as_grch37(chrom)
    return chrom, start, end

def _get_var_ref_and_alt(var):
    """Retrieve variant reference and alternate alleles from multiple input objects.
    """
    if isinstance(var, cyvcf2.Variant):
      return var.REF, var.ALT
    if isinstance(var, pysam.VariantRecord):
      return var.ref, var.alts
    if isinstance(var, basestring):
        # Assume var is a line from a VCF.
        ref, alt = var.split('\t', 6)[3:5]
    else:
        try:
            ref = var["ref"]
            alt = var["alt"]
        except (TypeError, AttributeError):
            try:
                ref = var.REF
                alt = var.ALT
            except KeyError:
                # For Pysam reader:
                ref = var.ref
                alt = var.alt

    if isinstance(alt, basestring):
        alt = alt.split(",")
    elif isinstance(alt, (tuple, list)):
        alt = [x for x in alt if x]
    return ref, alt

def _get_cadd_scores(var, labels, hit):
    """
    get cadd scores
    """
    raw = hit[3].split(",")
    scaled = hit[4].split(",")

    p = re.compile(str(var.ALT[0]))
    for m in p.finditer(str(labels[hit[2]])):
        pos = m.start()
        return raw[pos], scaled[pos]


def annotations_in_region(var, anno, parser_type=None, naming="ucsc"):
    """Iterator of annotations found in a genomic region.

    - var: PyVCF object or database query with chromosome, start and end.
    - anno: pysam Tabix annotation file or string to reference
            a standard annotation
    - parser_type: string specifying the filetype of the tabix file
    - naming: chromosome naming scheme used, ucsc or grch37
    """
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    return _get_hits(coords, anno, parser_type)

def annotations_in_vcf(var, anno, parser_type=None, naming="ucsc", region_only=False, warnings=False):
    """Iterator of annotations found in a VCF. For variants with multiple alleles,
       match using intersection and warn that decomposition, etc. is recommended.

    - var: PyVCF object or database query with chromosome, start and end.
    - anno: pysam Tabix annotation file or string to reference
            a standard annotation
    - parser_type: string specifying the filetype of the tabix file
    - naming: chromosome naming scheme used, ucsc or grch37
    - region_only: match using only region coordinates, not variant reference
                   and alternate; only used for VCF annotations
    """

    # Get hits by region only.
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    hits = _get_hits(coords, anno, parser_type)

    # Now filter by allele.
    if not region_only:
        def multiallele_warning(chrom, start, alt, is_annotation):
            """
            Print warnings for multi-allelic sites and recommend decomposition, etc.
            """

            # Check for multiple alleles and warnings flag.
            if not warnings:
                return
            if len(alt) == 1 or isinstance(alt, basestring):
                return

            variant_text = 'variant'
            if is_annotation:
                variant_text = 'annotation variant'

            sys.stderr.write("warning: %s with multiple alternate alleles found at %s:%i (alt: %s)\n"
                             "in order to reduce the number of false negatives we recommend splitting multiple alts. see:\n"
                             "http://gemini.readthedocs.org/en/latest/content/preprocessing.html#preprocess\n"
                             % (variant_text, chrom, start, ','.join(alt)))

        # Get variant ref, alt.
        var_ref, var_alt = _get_var_ref_and_alt(var)
        # no alternative alleles, return the original region_only hits
        if not var_alt or len(var_alt) == 0:
            return hits
        var_alt = set(var_alt)

        # Warn for multiple alleles.
        chrom, start, end = coords
        multiallele_warning(chrom, start, ','.join(var_alt), False)

        # Filter hits to those that match ref and alt.
        matched_hits = []
        for h in hits:
            if isinstance(h, (cyvcf2.Variant, pysam.VariantRecord)):
                start = h.start
            elif isinstance(h, basestring):
                start = int(h.split('\t', 2)[1]) - 1
            else:
                start = h.pos
            if start != coords[1]: continue

            anno_ref, anno_alt = _get_var_ref_and_alt(h)
            anno_alt = set(anno_alt)

            multiallele_warning(chrom, start, anno_alt, True)

            # Match via ref and set intersection of alternates.
            # the mappability uses "." as the alt for all rows. so
            if var_ref == anno_ref and (len(var_alt & anno_alt) >= 1 \
                    or anno_alt == set(".")):
                matched_hits.append(h)
        hits = matched_hits
    return hits


def bigwig_summary(var, anno, naming="ucsc"):
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    return _get_bw_summary(coords, anno)


# ## Track-specific annotations
def get_cpg_island_info(var):
    """
    Returns a boolean indicating whether or not the
    variant overlaps a CpG island
    """
    for hit in annotations_in_region(var, "cpg_island", "bed"):
        return True
    return False

# def get_dbNSFP_info(var, impacts):
#     """
#     Returns Polyphen, SIFT, etc. from dbNSFP annotation file.
#     One prediction per transcript.

#     LIMITATION: only handles bi-allelic loci
#     """

#     # is this variant predicted to be nonsynonymous for any of the transcripts?
#     # if not, we can skip dnNSFP.
#     non_syn_impacts = [imp for imp in impacts \
#                                if imp.consequence == 'non_syn_coding']

#     if len(non_syn_impacts) > 0:
#         for hit in annotations_in_region(var, "dbnsfp", parser_type="tuple", naming="grch37"):

#             if var.POS == int(hit[1]) and \
#                var.REF == hit[2] and \
#                var.ALT[0] == hit[3]:

#                 transcripts = hit[7].split(';')
#                 aapos = hit[8].split(';')
#                 pp_scores = hit[11].split(';')

#                 if len(transcripts) != len(pp_scores):
#                     print var.POS, var.REF, var.ALT[0], [i.transcript for i in non_syn_impacts], \
#                           [i.polyphen_pred for i in non_syn_impacts], [i.polyphen_score for i in non_syn_impacts], \
#                           hit[7], hit[8], hit[11], hit[12]
#     else:
#         pass

def get_cyto_info(var):
    """
    Returns a comma-separated list of the chromosomal
    cytobands that a variant overlaps.
    """
    cyto_band = ''
    for hit in annotations_in_region(var, "cytoband", "bed"):
        if len(cyto_band) > 0:
            cyto_band += "," + hit.contig + hit.name
        else:
            cyto_band += hit.contig + hit.name
    return cyto_band if len(cyto_band) > 0 else None

def get_gerp_bp(var):
    """
    Returns a summary of the GERP scores for the variant.
    """
    if "gerp_bp" not in annos:
        raise IOError("Need to download BigWig file with GERP scores per base pair. "
                      "Run `gemini update --dataonly --extra gerp_bp")
    gerp = bigwig_summary(var, "gerp_bp")
    return gerp

def get_gerp_elements(var):
    """
    Returns the GERP element information.
    """
    p_vals = []
    for hit in annotations_in_region(var, "gerp_elements", "tuple"):
        p_vals.append(hit[3])
    if len(p_vals) == 1:
        return p_vals[0]
    elif len(p_vals) > 1:
        return min(float(p) for p in p_vals)
    else:
        return None

def get_vista_enhancers(var):
    """
    Returns the VISTA enhancer information.
    """
    vista_enhancers = []
    for hit in annotations_in_region(var, "vista_enhancers", "tuple"):
        vista_enhancers.append(hit[4])
    return ",".join(vista_enhancers) if len(vista_enhancers) > 0 else None

def get_fitcons(var):
    hmax = float('nan')
    for hit in annotations_in_region(var, "fitcons", None, "ucsc"):
        _, val = hit.rsplit("\t", 1)
        v = float(val)
        if not hmax > v:
            hmax = v
    return hmax

def get_cadd_scores(var):
    """
    Returns the C-raw scores & scaled scores (CADD) to predict deleterious
    variants. Implemented only for SNV's
    """
    if "cadd_score" not in annos:
        raise IOError("Need to download the CADD data file for deleteriousness."
                      "Run `gemini update --dataonly --extra cadd_score")

    cadd_raw = cadd_scaled = None
    labels = {"A":"CGT", "C":"AGT", "G":"ACT", "T":"ACG", "R":"ACGT", "M":"ACGT"}

    for hit in annotations_in_region(var, "cadd_score", "tuple", "grch37"):
        # we want exact position mapping here and not a range (end-start) as
        # returned in hit (e.g. indels) & we do not want to consider del & ins
        if str(hit[1]) == str(var.POS) and var.REF and var.ALT[0] and \
           len(var.REF) == 1 and len(var.ALT[0]) == 1:

            if str(hit[2]) == var.REF and str(var.ALT[0]) in labels[hit[2]]:
               (cadd_raw, cadd_scaled) = _get_cadd_scores(var, labels, hit)

            # consider ref cases with ambiguity codes R (G,A) and M (A,C)
            elif ((str(hit[2]) == 'R'  and var.REF in('G','A')) or \
                (str(hit[2]) == 'M'  and var.REF in('A','C'))) and \
                str(var.ALT[0]) in labels[hit[2]]:
                (cadd_raw, cadd_scaled) = _get_cadd_scores(var, labels, hit)

    return (cadd_raw, cadd_scaled)


def get_pfamA_domains(var):
    """
    Returns pfamA domains that a variant overlaps
    """
    pfam_domain = []
    for hit in annotations_in_region(var, "pfam_domain", "bed"):
        pfam_domain.append(hit.name)
    return ",".join(pfam_domain) if len(pfam_domain) > 0 else None


def get_cosmic_info(var):
    """
    Returns a list of COSMIC ids associated with given variant

    E.g. from COSMIC VCF
    #CHROM  POS     ID      REF     ALT     QUAL    FILTER  INFO
    chrM    1747    COSN408408      G       A       .       .       .
    chrM    2700    COSN408409      G       A       .       .       .
    chr1    42880262    COSM464635  G   C   .   .   AA=p.D224H;CDS=c.670G>C;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880269    COSM909628  G   A   .   .   AA=p.G226D;CDS=c.677G>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880284    COSM1502979 G   T   .   .   AA=p.C231F;CDS=c.692G>T;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880318    COSM681351  T   A   .   .   AA=p.F242L;CDS=c.726T>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880337    COSM464636  G   A   .   .   AA=p.D249N;CDS=c.745G>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880384    COSM909629  T   C   .   .   AA=p.N264N;CDS=c.792T>C;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880415    COSM909630  G   C   .   .   AA=p.G275R;CDS=c.823G>C;CNT=1;GENE=RIMKLA;STRAND=+
    """
    # report the first overlapping ClinVar variant Most often, just one).
    cosmic_ids = []
    for hit in annotations_in_vcf(var, "cosmic", "vcf", "grch37"):
        cosmic_ids.append(hit.id)
    return ",".join(cosmic_ids) if len(cosmic_ids) > 0 else None


def get_clinvar_info(var):
    """
    Returns a suite of annotations from ClinVar

    ClinVarInfo named_tuple:
    --------------------------------------------------------------------------
    # clinvar_dbsource         = CLNSRC=OMIM Allelic Variant;
    # clinvar_dbsource_id      = CLNSRCID=103320.0001;
    # clinvar_origin           = CLNORIGIN=1
    # clinvar_sig              = CLNSIG=5
    # clinvar_dsdb             = CLNDSDB=GeneReviews:NCBI:OMIM:Orphanet;
    # clinvar_dsdbid           = CLNDSDBID=NBK1168:C1850792:254300:590;
    # clinvar_disease_name     = CLNDBN=Myasthenia\x2c limb-girdle\x2c familial;
    # clinvar_disease_acc      = CLNACC=RCV000019902.1
    # clinvar_in_omim          = OM
    # clinvar_causal_allele    = CLNALLE=1
    """

    clinvar = ClinVarInfo()

    # report the first overlapping ClinVar variant Most often, just one).
    for hit in annotations_in_vcf(var, "clinvar", "vcf", "grch37"):
        # load each VCF INFO key/value pair into a DICT
        info_map = {}
        vals = hit.info.split(';')

        for info in vals:
            if info.find("=") > 0:
                (key, value) = info.split("=")
                info_map[key] = value
            else:
                info_map[info] = True
        clinvar.clinvar_origin = clinvar.lookup_clinvar_origin(info_map.get('ORIGIN', '0'))
        clinvar.clinvar_sig = info_map.get('CLNSIG', info_map.get('CLNDISDBINCL', '')).lower()
        clinvar.clinvar_dsdb = info_map.get('CLNDISDB') or None
        # Remap all unicode characters into plain text string replacements
        raw_disease_name = info_map.get('CLNDN', '')
        try:
            clinvar.clinvar_disease_name = unidecode(raw_disease_name.decode('utf-8')).decode('string_escape')
        except:
            clinvar.clinvar_disease_name = unidecode(raw_disease_name.encode('utf-8').decode())
        # Clinvar represents commas as \x2c.  Make them commas.

        clinvar.clinvar_in_omim = 1 if 'OMIM' in hit.info else 0

        clinvar.clinvar_causal_allele = hit.alt
    return clinvar


def get_dbsnp_info(var):
    """
    Returns a suite of annotations from dbSNP
    """
    rs_ids = []
    for hit in annotations_in_vcf(var, "dbsnp", "vcf", "grch37"):
        rs_ids.append(hit.id)
    return ",".join(rs_ids) if len(rs_ids) > 0 else None


def get_esp_info(var):
    """
    Returns a suite of annotations from the ESP project

    ESP reports the minor allele frequency (MAF), not the
    alternate allele frequency (AAF). We must therefore figure
    out whther the reference or alternate allele is the minor allele.

    1       69496   rs150690004     G       A       .       PASS    DBSNP=dbSNP_134;EA_AC=2,6764;AA_AC=23,3785;TAC=25,10549;MAF=0.0296,0.604,0.2364;GTS=AA,AG,GG;EA_GTC=0,2,3381;AA_GTC=5,13,1886;GTC=5,15,5267;DP=91;GL=OR4F5;CP=0.5;CG=2.3;AA=G;CA=.;EXOME_CHIP=no;GWAS_PUBMED=.;GM=NM_001005484.1;FG=missense;AAC=SER/GLY;PP=136/306;CDP=406;GS=56;PH=benign
    1       69511   rs75062661      A       G       .       PASS    DBSNP=dbSNP_131;EA_AC=5337,677;AA_AC=1937,1623;TAC=7274,2300;MAF=11.2571,45.5899,24.0234;GTS=GG,GA,AA;EA_GTC=2430,477,100;AA_GTC=784,369,627;GTC=3214,846,727;DP=69;GL=OR4F5;CP=1.0;CG=1.1;AA=G;CA=.;EXOME_CHIP=no;GWAS_PUBMED=.;GM=NM_001005484.1;FG=missense;AAC=ALA/THR;PP=141/306;CDP=421;GS=58;PH=benign
    """
    fetched = []
    exome_chip = False
    found = False
    info_map = {}
    acs = {}
    for hit in annotations_in_vcf(var, "esp", "vcf", "grch37"):
        if hit.contig not in ['Y']:
            fetched.append(hit)
            # We need a single ESP entry for a variant
            if fetched != None and len(fetched) == 1 and \
                    hit.alt == var.ALT[0] and hit.ref == var.REF:
                found = True
                # loads each VCF INFO key/value pair into a DICT
                for info in hit.info.split(";"):
                    if info.find("=") > 0:
                    # splits on first occurence of '='
                    # useful to handle valuerror: too many values to unpack (e.g (a,b) = split(",", (a,b,c,d)) for cases like
                    # SA=http://www.ncbi.nlm.nih.gov/sites/varvu?gene=4524&amp%3Brs=1801131|http://omim.org/entry/607093#0004
                        (key, value) = info.split("=", 1)
                        info_map[key] = value

                # NOTE:, if we start ot use GTS, need to update preprocessing
                # script to handle weirdness on X, Y
                # get the allele counts so that we can compute alternate allele frequencies
                # example: EA_AC=2,6764;AA_AC=23,3785;TAC=25,10549
                for key in ('EA_AC', 'AA_AC', 'TAC'):
                    if info_map.get(key) is not None:
                        lines = info_map[key].split(",")
                        denom = float(lines[0]) + float(lines[1])
                        if denom == 0:
                            acs[key] = 0
                        else:
                            # alt allele is stored as 2nd.
                            acs[key] = float(lines[1]) / denom
                    else:
                        acs[key] = -1

                # Is the SNP on an human exome chip?
                if info_map.get('EXOME_CHIP') is not None and \
                        info_map['EXOME_CHIP'] == "no":
                    exome_chip = 0
                elif info_map.get('EXOME_CHIP') is not None and \
                        info_map['EXOME_CHIP'] == "yes":
                    exome_chip = 1
                break
    return ESPInfo(found, acs.get('EA_AC', -1), acs.get("AA_AC", -1),
            acs.get("TAC", -1), exome_chip)


EMPTY_1000G = ThousandGInfo(False, -1, -1, -1, -1, -1, -1)
def get_1000G_info(var, empty=EMPTY_1000G):
    """
    Returns a suite of annotations from the 1000 Genomes project
    """
    #fetched = []
    info_map = {}

    for hit in annotations_in_vcf(var, "1000g", "vcf", "grch37"):
        # We need to ensure we are dealing with the exact sample variant
        # based on position and the alleles present.
        # var.start is used since the chromosomal pos in pysam.asVCF is zero based (hit.pos)
        # and would be equivalent to (POS-1) i.e. var.start
        if var.start == hit.pos and \
           var.ALT and len(var.ALT) > 0 and \
           var.ALT[0] == hit.alt and \
           hit.ref == var.REF:
            for info in hit.info.split(";"):
                if "=" in info:
                    (key, value) = info.split("=", 1)
                    info_map[key] = value

            return ThousandGInfo(True, float(info_map.get('AF', -1)), float(info_map.get('AMR_AF', -1)),
                         float(info_map.get('EAS_AF', -1)), float(info_map.get('SAS_AF', -1)),
                         float(info_map.get('AFR_AF', -1)), float(info_map.get('EUR_AF', -1)))
    return empty

def get_geno2mp_ct(var):

    for hit in annotations_in_vcf(var, "geno2mp", "vcf", "grch37"):
        if not (var.start == hit.pos and var.REF == hit.ref):
            continue
        if not var.ALT[0] in hit.alt.split(","): continue

        ct = next(x for x in hit.info.split(";") if x.startswith("HPO_CT="))
        val = int(ct.split("=")[1])
        return val
    # missing is -1
    return -1

def get_gnomad_info(var, empty=GNOMAD_EMPTY):
    afs = {}
    for hit in annotations_in_vcf(var, "gnomad", "vcf", "grch37"):
        # Does not handle anything beyond var.ALT[0] in the VCF (in case of multi-allelic variants)
        # var.start is used since the chromosomal pos in pysam.asVCF is zero based (hit.pos)
        # and would be equivalent to (POS-1) i.e var.start
        # TODO: change to REF
        if var.start != hit.start or var.REF != hit.REF or var.ALT[0] != hit.ALT[0]:
            continue

        # This would look for var.ALT[0] matches to
        # any of the multiple alt alleles represented in the EXAC file
        # Population independent raw (non-adjusted) allele frequencies given by AF
        info_map = hit.INFO
        aaf_ALL = info_map.get("AF", -1.0)

        for grp in ('_afr', '_amr', '_asj', '_eas', '_fin', '_nfe', '_oth', '_sas'):
            ac = info_map.get('AC' + grp)
            if ac is None: continue

            an = info_map.get('AN' + grp)
            if an is None: continue

            if an == 0:
                afs[grp] = 0
                continue

            afs[grp] = float(ac) / float(an)
        if ac is None or an is None: continue

        num_homs = int(info_map.get("nhomalt", -1))
        num_hets = int(info_map.get("AC", -1))
        if num_hets != -1:
          num_hets -= 2 * num_homs

        called_chroms = int(info_map.get('AN', -1))

        return GnomadInfo(aaf_ALL, float(afs['_afr']),
                              float(afs['_amr']), float(afs['_asj']),
                              float(afs['_eas']), float(afs['_fin']),
                              float(afs['_nfe']), float(afs['_oth']),
                              float(afs['_sas']),
                              info_map.get('non_neuro_AF', -1.0),
                              info_map.get('controls_AF', -1.0),
                              info_map.get('non_cancer_AF', -1.0),
                              info_map.get("popmax_AF", -1.0),
                              num_hets, num_homs,
                              called_chroms)

    return empty

<commit_msg>correct AF_popmax<commit_after>#!/usr/bin/env python
from __future__ import absolute_import, print_function

import pysam
import cyvcf2
import os
import sys
import collections
import re
from unidecode import unidecode
from .config import read_gemini_config
from . import gemini_utils as util
try:
    basestring
except NameError:
    basestring = str

# dictionary of anno_type -> open Tabix file handles
annos = {}

def get_anno_files(args):
    config = read_gemini_config(args=args)
    anno_dirname = config["annotation_dir"]
    # Default annotations -- always found
    annos = {
     'pfam_domain': os.path.join(anno_dirname, 'hg19.pfam.ucscgenes.bed.gz'),
     'cytoband': os.path.join(anno_dirname, 'hg19.cytoband.bed.gz'),
     'dbsnp': os.path.join(anno_dirname, 'dbsnp.151.20180423.tidy.vcf.gz'),
     'clinvar': os.path.join(anno_dirname, 'clinvar_20190102.tidy.vcf.gz'),
     'gwas': os.path.join(anno_dirname, 'hg19.gwas.bed.gz'),
     'rmsk': os.path.join(anno_dirname, 'hg19.rmsk.bed.gz'),
     'segdup': os.path.join(anno_dirname, 'hg19.segdup.bed.gz'),
     'conserved': os.path.join(anno_dirname, '29way_pi_lods_elements_12mers.chr_specific.fdr_0.1_with_scores.txt.hg19.merged.bed.gz'),
     'cpg_island': os.path.join(anno_dirname, 'hg19.CpG.bed.gz'),
     'dgv': os.path.join(anno_dirname, 'hg19.dgv.bed.gz'),
     'esp': os.path.join(anno_dirname,
                         'ESP6500SI.all.snps_indels.tidy.v2.vcf.gz'),
     '1000g': os.path.join(anno_dirname,
                           'ALL.wgs.phase3_shapeit2_mvncall_integrated_v5a.20130502.sites.tidy.vcf.gz'),
     'recomb': os.path.join(anno_dirname,
                            'genetic_map_HapMapII_GRCh37.gz'),
     'gms': os.path.join(anno_dirname,
                         'GRCh37-gms-mappability.vcf.gz'),
     'grc': os.path.join(anno_dirname, 'GRC_patch_regions.bed.gz'),
     'cse': os.path.join(anno_dirname, "cse-hiseq-8_4-2013-02-20.bed.gz"),
     'encode_tfbs': os.path.join(anno_dirname,
                                 'wgEncodeRegTfbsClusteredV2.cell_count.20130213.bed.gz'),
     'encode_dnase1': os.path.join(anno_dirname,
                                   'stam.125cells.dnaseI.hg19.bed.gz'),
     'encode_consensus_segs': os.path.join(anno_dirname,
                                           'encode.6celltypes.consensus.bedg.gz'),
     'gerp_elements': os.path.join(anno_dirname, 'hg19.gerp.elements.bed.gz'),
     'vista_enhancers': os.path.join(anno_dirname, 'hg19.vista.enhancers.20131108.bed.gz'),
     'fitcons': os.path.join(anno_dirname, "hg19_fitcons_fc-i6-0_V1-01.bed.gz"),
     'cosmic': os.path.join(anno_dirname, 'cosmic-v68-GRCh37.tidy.vcf.gz'),
     'exac': os.path.join(anno_dirname, 'ExAC.r0.3.sites.vep.tidy.vcf.gz'),
     'gnomad': os.path.join(anno_dirname, 'gnomad.exomes.r2.1.tidy.bcf'),
     'geno2mp': os.path.join(anno_dirname, 'geno2mp.variants.tidy.vcf.gz'),
    }
    # optional annotations
    if os.path.exists(os.path.join(anno_dirname, 'hg19.gerp.bw')):
        annos['gerp_bp'] = os.path.join(anno_dirname, 'hg19.gerp.bw')
    if os.path.exists(os.path.join(anno_dirname, 'caddv1.4.compressed.gz')):
        annos['cadd_score'] = os.path.join(anno_dirname, 'caddv1.4.compressed.gz')
    return annos

class ClinVarInfo(object):
    def __init__(self):
        self.clinvar_origin = None
        self.clinvar_sig = None
        self.clinvar_dsdb = None
        self.clinvar_dsdbid = None
        self.clinvar_disease_name = None
        self.clinvar_in_omim = None

        self.origin_code_map = {'0': 'unknown',
                                '1': 'germline',
                                '2': 'somatic',
                                '4': 'inherited',
                                '8': 'paternal',
                                '16': 'maternal',
                                '32': 'de-novo',
                                '64': 'biparental',
                                '128': 'uniparental',
                                '256': 'not-tested',
                                '512': 'tested-inconclusive',
                                '1073741824': 'other'}
        # 0 - Uncertain significance, 1 - not provided, 2 - Benign, 3 - Likely benign, 4 - Likely pathogenic, 5 - Pathogenic, 6 - drug response, 7 - histocompatibility, 255 - other

        self.sig_code_map = {'0': 'uncertain',
                             '1': 'not-provided',
                             '2': 'benign',
                             '3': 'likely-benign',
                             '4': 'likely-pathogenic',
                             '5': 'pathogenic',
                             '6': 'drug-response',
                             '7': 'histocompatibility',
                             '255': 'other'}

    def __repr__(self):
        return '\t'.join(map(str, [
                          self.clinvar_origin,
                          self.clinvar_sig,
                          self.clinvar_dsdb,
                          self.clinvar_dsdbid,
                          self.clinvar_disease_name,
                          str(self.clinvar_in_omim),
                          str(self.clinvar_causal_allele)]))

    def lookup_clinvar_origin(self, origin_code):
        try:
            return self.origin_code_map[origin_code]
        except KeyError:
            return None

    def lookup_clinvar_significance(self, sig_code):
        sigs = []
        for s in sig_code.split('|'):
            sigs.extend(s.split(","))

        return ",".join(self.sig_code_map[s] for s in set(sigs) if s != ".")


ESPInfo = collections.namedtuple("ESPInfo",
                                 "found \
                                  aaf_EA \
                                  aaf_AA \
                                  aaf_ALL \
                                  exome_chip")
ENCODEDnaseIClusters = collections.namedtuple("ENCODEDnaseIClusters",
                                              "cell_count \
                                         cell_list")
ENCODESegInfo = collections.namedtuple("ENCODESegInfo",
                                       "gm12878 \
                                         h1hesc \
                                         helas3 \
                                         hepg2 \
                                         huvec \
                                         k562")
ThousandGInfo = collections.namedtuple("ThousandGInfo",
                                       "found \
                                        aaf_ALL \
                                        aaf_AMR \
                                        aaf_EAS \
                                        aaf_SAS \
                                        aaf_AFR \
                                        aaf_EUR")

GmsTechs = collections.namedtuple("GmsTechs", "illumina solid iontorrent")

ExacInfo = collections.namedtuple("ExacInfo",
                                  "found \
                                   aaf_ALL \
                                   adj_aaf_ALL \
                                   aaf_AFR \
                                   aaf_AMR \
                                   aaf_EAS \
                                   aaf_FIN \
                                   aaf_NFE \
                                   aaf_OTH \
                                   aaf_SAS \
                                   num_het \
                                   num_hom_alt \
                                   num_chroms")

EXAC_EMPTY = ExacInfo(False, -1, -1, -1, -1, -1,
                     -1, -1, -1, -1, -1, -1, -1)

GnomadInfo = collections.namedtuple('GnomadInfo',
                                   "aaf_ALL \
                                   aaf_AFR \
                                   aaf_AMR \
                                   aaf_ASJ \
                                   aaf_EAS \
                                   aaf_FIN \
                                   aaf_NFE \
                                   aaf_OTH \
                                   aaf_SAS \
                                   aaf_non_neuro \
                                   aaf_controls \
                                   aaf_non_cancer \
                                   popmax_AF \
                                   num_het \
                                   num_hom_alt \
                                   num_chroms")

GNOMAD_EMPTY = GnomadInfo(-1, -1, -1, -1, -1, -1, -1, -1,
                          -1, -1, -1, -1, -1, -1, -1, -1)

def load_annos(args):
    """
    Populate a dictionary of Tabixfile handles for
    each annotation file.  Other modules can then
    access a given handle and fetch data from it
    as follows:

    dbsnp_handle = annotations.annos['dbsnp']
    hits = dbsnp_handle.fetch(chrom, start, end)
    """
    anno_files = get_anno_files(args)
    for anno in anno_files:
        try:
            # .gz denotes Tabix files.
            if anno_files[anno].endswith(".gz"):
                if anno == "clinvar":
                    annos[anno] = pysam.Tabixfile(anno_files[anno],
                                                  encoding='utf8')
                else:
                    annos[anno] = pysam.Tabixfile(anno_files[anno])

            elif anno_files[anno].endswith(".bcf"):
                annos[anno] = cyvcf2.VCF(anno_files[anno])
            # .bw denotes BigWig files.
            elif anno_files[anno].endswith(".bw"):
                from bx.bbi.bigwig_file import BigWigFile
                annos[anno] = BigWigFile(open(anno_files[anno]))

        except IOError:
            raise IOError("Gemini cannot open this annotation file: %s. \n"
                          "Have you installed the annotation files?  If so, "
                          "have they been moved or deleted? Exiting...\n\n"
                          "For more details:\n\t"
                          "http://gemini.readthedocs.org/en/latest/content/"
                          "#installation.html\#installing-annotation-files\n"
                          % anno_files[anno])

# ## Standard access to Tabix indexed files


PARSERS = {"bed": pysam.asBed(),
           "vcf": pysam.asVCF(),
           "tuple": pysam.asTuple(),
           None: None}

def _get_hits(coords, annotation, parser_type, _parsers=PARSERS):
    """Retrieve BED information, recovering if BED annotation file does have a chromosome.
    """
    try:
        parser = _parsers[parser_type]
    except KeyError:
        raise ValueError("Unexpected parser type: %s" % parser)
    chrom, start, end = coords
    if isinstance(annotation, pysam.VariantFile):
        return annotation.fetch(chrom, start, end)
    elif isinstance(annotation, cyvcf2.VCF):
        return annotation("%s:%d-%d" % (chrom, start - 1, end))
    try:
        hit_iter = annotation.fetch(str(chrom), start, end, parser=parser)
    # catch invalid region errors raised by ctabix
    except ValueError:
        hit_iter = []
    # recent versions of pysam return KeyError
    except KeyError:
        hit_iter = []
    except:
        print(annotation.__class__, file=sys.stderr)
        raise
    return hit_iter

def _get_bw_summary(coords, annotation):
    """Return summary of BigWig scores in an interval
    """
    chrom, start, end = coords
    try:
        return annotation.summarize(str(chrom), start, end, end-start).min_val[0]
    except AttributeError:
        return None


def _get_chr_as_grch37(chrom):
    if chrom in ["chrM"]:
        return "MT"
    return chrom if not chrom.startswith("chr") else chrom[3:]


def _get_chr_as_ucsc(chrom):
    return chrom if chrom.startswith("chr") else "chr" + chrom


def guess_contig_naming(anno):
    """Guess which contig naming scheme a given annotation file uses.
    """
    chr_names = [x for x in anno.contigs if x.startswith("chr")]
    if len(chr_names) > 0:
        return "ucsc"
    else:
        return "grch37"


def _get_var_coords(var, naming):
    """Retrieve variant coordinates from multiple input objects.
    """
    if isinstance(var, cyvcf2.Variant):
      chrom, start, end = var.CHROM, var.start, var.end
    else:
        try:
            # todo: check isinstance resultproxy?
            chrom = var["chrom"]
            start = int(var["start"])
            end = int(var["end"])
        except TypeError:
            chrom = var.CHROM
            start = var.start
            end = var.end
    if naming == "ucsc":
        chrom = _get_chr_as_ucsc(chrom)
    elif naming == "grch37":
        chrom = _get_chr_as_grch37(chrom)
    return chrom, start, end

def _get_var_ref_and_alt(var):
    """Retrieve variant reference and alternate alleles from multiple input objects.
    """
    if isinstance(var, cyvcf2.Variant):
      return var.REF, var.ALT
    if isinstance(var, pysam.VariantRecord):
      return var.ref, var.alts
    if isinstance(var, basestring):
        # Assume var is a line from a VCF.
        ref, alt = var.split('\t', 6)[3:5]
    else:
        try:
            ref = var["ref"]
            alt = var["alt"]
        except (TypeError, AttributeError):
            try:
                ref = var.REF
                alt = var.ALT
            except KeyError:
                # For Pysam reader:
                ref = var.ref
                alt = var.alt

    if isinstance(alt, basestring):
        alt = alt.split(",")
    elif isinstance(alt, (tuple, list)):
        alt = [x for x in alt if x]
    return ref, alt

def _get_cadd_scores(var, labels, hit):
    """
    get cadd scores
    """
    raw = hit[3].split(",")
    scaled = hit[4].split(",")

    p = re.compile(str(var.ALT[0]))
    for m in p.finditer(str(labels[hit[2]])):
        pos = m.start()
        return raw[pos], scaled[pos]


def annotations_in_region(var, anno, parser_type=None, naming="ucsc"):
    """Iterator of annotations found in a genomic region.

    - var: PyVCF object or database query with chromosome, start and end.
    - anno: pysam Tabix annotation file or string to reference
            a standard annotation
    - parser_type: string specifying the filetype of the tabix file
    - naming: chromosome naming scheme used, ucsc or grch37
    """
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    return _get_hits(coords, anno, parser_type)

def annotations_in_vcf(var, anno, parser_type=None, naming="ucsc", region_only=False, warnings=False):
    """Iterator of annotations found in a VCF. For variants with multiple alleles,
       match using intersection and warn that decomposition, etc. is recommended.

    - var: PyVCF object or database query with chromosome, start and end.
    - anno: pysam Tabix annotation file or string to reference
            a standard annotation
    - parser_type: string specifying the filetype of the tabix file
    - naming: chromosome naming scheme used, ucsc or grch37
    - region_only: match using only region coordinates, not variant reference
                   and alternate; only used for VCF annotations
    """

    # Get hits by region only.
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    hits = _get_hits(coords, anno, parser_type)

    # Now filter by allele.
    if not region_only:
        def multiallele_warning(chrom, start, alt, is_annotation):
            """
            Print warnings for multi-allelic sites and recommend decomposition, etc.
            """

            # Check for multiple alleles and warnings flag.
            if not warnings:
                return
            if len(alt) == 1 or isinstance(alt, basestring):
                return

            variant_text = 'variant'
            if is_annotation:
                variant_text = 'annotation variant'

            sys.stderr.write("warning: %s with multiple alternate alleles found at %s:%i (alt: %s)\n"
                             "in order to reduce the number of false negatives we recommend splitting multiple alts. see:\n"
                             "http://gemini.readthedocs.org/en/latest/content/preprocessing.html#preprocess\n"
                             % (variant_text, chrom, start, ','.join(alt)))

        # Get variant ref, alt.
        var_ref, var_alt = _get_var_ref_and_alt(var)
        # no alternative alleles, return the original region_only hits
        if not var_alt or len(var_alt) == 0:
            return hits
        var_alt = set(var_alt)

        # Warn for multiple alleles.
        chrom, start, end = coords
        multiallele_warning(chrom, start, ','.join(var_alt), False)

        # Filter hits to those that match ref and alt.
        matched_hits = []
        for h in hits:
            if isinstance(h, (cyvcf2.Variant, pysam.VariantRecord)):
                start = h.start
            elif isinstance(h, basestring):
                start = int(h.split('\t', 2)[1]) - 1
            else:
                start = h.pos
            if start != coords[1]: continue

            anno_ref, anno_alt = _get_var_ref_and_alt(h)
            anno_alt = set(anno_alt)

            multiallele_warning(chrom, start, anno_alt, True)

            # Match via ref and set intersection of alternates.
            # the mappability uses "." as the alt for all rows. so
            if var_ref == anno_ref and (len(var_alt & anno_alt) >= 1 \
                    or anno_alt == set(".")):
                matched_hits.append(h)
        hits = matched_hits
    return hits


def bigwig_summary(var, anno, naming="ucsc"):
    coords = _get_var_coords(var, naming)
    if isinstance(anno, basestring):
        anno = annos[anno]
    return _get_bw_summary(coords, anno)


# ## Track-specific annotations
def get_cpg_island_info(var):
    """
    Returns a boolean indicating whether or not the
    variant overlaps a CpG island
    """
    for hit in annotations_in_region(var, "cpg_island", "bed"):
        return True
    return False

# def get_dbNSFP_info(var, impacts):
#     """
#     Returns Polyphen, SIFT, etc. from dbNSFP annotation file.
#     One prediction per transcript.

#     LIMITATION: only handles bi-allelic loci
#     """

#     # is this variant predicted to be nonsynonymous for any of the transcripts?
#     # if not, we can skip dnNSFP.
#     non_syn_impacts = [imp for imp in impacts \
#                                if imp.consequence == 'non_syn_coding']

#     if len(non_syn_impacts) > 0:
#         for hit in annotations_in_region(var, "dbnsfp", parser_type="tuple", naming="grch37"):

#             if var.POS == int(hit[1]) and \
#                var.REF == hit[2] and \
#                var.ALT[0] == hit[3]:

#                 transcripts = hit[7].split(';')
#                 aapos = hit[8].split(';')
#                 pp_scores = hit[11].split(';')

#                 if len(transcripts) != len(pp_scores):
#                     print var.POS, var.REF, var.ALT[0], [i.transcript for i in non_syn_impacts], \
#                           [i.polyphen_pred for i in non_syn_impacts], [i.polyphen_score for i in non_syn_impacts], \
#                           hit[7], hit[8], hit[11], hit[12]
#     else:
#         pass

def get_cyto_info(var):
    """
    Returns a comma-separated list of the chromosomal
    cytobands that a variant overlaps.
    """
    cyto_band = ''
    for hit in annotations_in_region(var, "cytoband", "bed"):
        if len(cyto_band) > 0:
            cyto_band += "," + hit.contig + hit.name
        else:
            cyto_band += hit.contig + hit.name
    return cyto_band if len(cyto_band) > 0 else None

def get_gerp_bp(var):
    """
    Returns a summary of the GERP scores for the variant.
    """
    if "gerp_bp" not in annos:
        raise IOError("Need to download BigWig file with GERP scores per base pair. "
                      "Run `gemini update --dataonly --extra gerp_bp")
    gerp = bigwig_summary(var, "gerp_bp")
    return gerp

def get_gerp_elements(var):
    """
    Returns the GERP element information.
    """
    p_vals = []
    for hit in annotations_in_region(var, "gerp_elements", "tuple"):
        p_vals.append(hit[3])
    if len(p_vals) == 1:
        return p_vals[0]
    elif len(p_vals) > 1:
        return min(float(p) for p in p_vals)
    else:
        return None

def get_vista_enhancers(var):
    """
    Returns the VISTA enhancer information.
    """
    vista_enhancers = []
    for hit in annotations_in_region(var, "vista_enhancers", "tuple"):
        vista_enhancers.append(hit[4])
    return ",".join(vista_enhancers) if len(vista_enhancers) > 0 else None

def get_fitcons(var):
    hmax = float('nan')
    for hit in annotations_in_region(var, "fitcons", None, "ucsc"):
        _, val = hit.rsplit("\t", 1)
        v = float(val)
        if not hmax > v:
            hmax = v
    return hmax

def get_cadd_scores(var):
    """
    Returns the C-raw scores & scaled scores (CADD) to predict deleterious
    variants. Implemented only for SNV's
    """
    if "cadd_score" not in annos:
        raise IOError("Need to download the CADD data file for deleteriousness."
                      "Run `gemini update --dataonly --extra cadd_score")

    cadd_raw = cadd_scaled = None
    labels = {"A":"CGT", "C":"AGT", "G":"ACT", "T":"ACG", "R":"ACGT", "M":"ACGT"}

    for hit in annotations_in_region(var, "cadd_score", "tuple", "grch37"):
        # we want exact position mapping here and not a range (end-start) as
        # returned in hit (e.g. indels) & we do not want to consider del & ins
        if str(hit[1]) == str(var.POS) and var.REF and var.ALT[0] and \
           len(var.REF) == 1 and len(var.ALT[0]) == 1:

            if str(hit[2]) == var.REF and str(var.ALT[0]) in labels[hit[2]]:
               (cadd_raw, cadd_scaled) = _get_cadd_scores(var, labels, hit)

            # consider ref cases with ambiguity codes R (G,A) and M (A,C)
            elif ((str(hit[2]) == 'R'  and var.REF in('G','A')) or \
                (str(hit[2]) == 'M'  and var.REF in('A','C'))) and \
                str(var.ALT[0]) in labels[hit[2]]:
                (cadd_raw, cadd_scaled) = _get_cadd_scores(var, labels, hit)

    return (cadd_raw, cadd_scaled)


def get_pfamA_domains(var):
    """
    Returns pfamA domains that a variant overlaps
    """
    pfam_domain = []
    for hit in annotations_in_region(var, "pfam_domain", "bed"):
        pfam_domain.append(hit.name)
    return ",".join(pfam_domain) if len(pfam_domain) > 0 else None


def get_cosmic_info(var):
    """
    Returns a list of COSMIC ids associated with given variant

    E.g. from COSMIC VCF
    #CHROM  POS     ID      REF     ALT     QUAL    FILTER  INFO
    chrM    1747    COSN408408      G       A       .       .       .
    chrM    2700    COSN408409      G       A       .       .       .
    chr1    42880262    COSM464635  G   C   .   .   AA=p.D224H;CDS=c.670G>C;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880269    COSM909628  G   A   .   .   AA=p.G226D;CDS=c.677G>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880284    COSM1502979 G   T   .   .   AA=p.C231F;CDS=c.692G>T;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880318    COSM681351  T   A   .   .   AA=p.F242L;CDS=c.726T>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880337    COSM464636  G   A   .   .   AA=p.D249N;CDS=c.745G>A;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880384    COSM909629  T   C   .   .   AA=p.N264N;CDS=c.792T>C;CNT=1;GENE=RIMKLA;STRAND=+
    chr1    42880415    COSM909630  G   C   .   .   AA=p.G275R;CDS=c.823G>C;CNT=1;GENE=RIMKLA;STRAND=+
    """
    # report the first overlapping ClinVar variant Most often, just one).
    cosmic_ids = []
    for hit in annotations_in_vcf(var, "cosmic", "vcf", "grch37"):
        cosmic_ids.append(hit.id)
    return ",".join(cosmic_ids) if len(cosmic_ids) > 0 else None


def get_clinvar_info(var):
    """
    Returns a suite of annotations from ClinVar

    ClinVarInfo named_tuple:
    --------------------------------------------------------------------------
    # clinvar_dbsource         = CLNSRC=OMIM Allelic Variant;
    # clinvar_dbsource_id      = CLNSRCID=103320.0001;
    # clinvar_origin           = CLNORIGIN=1
    # clinvar_sig              = CLNSIG=5
    # clinvar_dsdb             = CLNDSDB=GeneReviews:NCBI:OMIM:Orphanet;
    # clinvar_dsdbid           = CLNDSDBID=NBK1168:C1850792:254300:590;
    # clinvar_disease_name     = CLNDBN=Myasthenia\x2c limb-girdle\x2c familial;
    # clinvar_disease_acc      = CLNACC=RCV000019902.1
    # clinvar_in_omim          = OM
    # clinvar_causal_allele    = CLNALLE=1
    """

    clinvar = ClinVarInfo()

    # report the first overlapping ClinVar variant Most often, just one).
    for hit in annotations_in_vcf(var, "clinvar", "vcf", "grch37"):
        # load each VCF INFO key/value pair into a DICT
        info_map = {}
        vals = hit.info.split(';')

        for info in vals:
            if info.find("=") > 0:
                (key, value) = info.split("=")
                info_map[key] = value
            else:
                info_map[info] = True
        clinvar.clinvar_origin = clinvar.lookup_clinvar_origin(info_map.get('ORIGIN', '0'))
        clinvar.clinvar_sig = info_map.get('CLNSIG', info_map.get('CLNDISDBINCL', '')).lower()
        clinvar.clinvar_dsdb = info_map.get('CLNDISDB') or None
        # Remap all unicode characters into plain text string replacements
        raw_disease_name = info_map.get('CLNDN', '')
        try:
            clinvar.clinvar_disease_name = unidecode(raw_disease_name.decode('utf-8')).decode('string_escape')
        except:
            clinvar.clinvar_disease_name = unidecode(raw_disease_name.encode('utf-8').decode())
        # Clinvar represents commas as \x2c.  Make them commas.

        clinvar.clinvar_in_omim = 1 if 'OMIM' in hit.info else 0

        clinvar.clinvar_causal_allele = hit.alt
    return clinvar


def get_dbsnp_info(var):
    """
    Returns a suite of annotations from dbSNP
    """
    rs_ids = []
    for hit in annotations_in_vcf(var, "dbsnp", "vcf", "grch37"):
        rs_ids.append(hit.id)
    return ",".join(rs_ids) if len(rs_ids) > 0 else None


def get_esp_info(var):
    """
    Returns a suite of annotations from the ESP project

    ESP reports the minor allele frequency (MAF), not the
    alternate allele frequency (AAF). We must therefore figure
    out whther the reference or alternate allele is the minor allele.

    1       69496   rs150690004     G       A       .       PASS    DBSNP=dbSNP_134;EA_AC=2,6764;AA_AC=23,3785;TAC=25,10549;MAF=0.0296,0.604,0.2364;GTS=AA,AG,GG;EA_GTC=0,2,3381;AA_GTC=5,13,1886;GTC=5,15,5267;DP=91;GL=OR4F5;CP=0.5;CG=2.3;AA=G;CA=.;EXOME_CHIP=no;GWAS_PUBMED=.;GM=NM_001005484.1;FG=missense;AAC=SER/GLY;PP=136/306;CDP=406;GS=56;PH=benign
    1       69511   rs75062661      A       G       .       PASS    DBSNP=dbSNP_131;EA_AC=5337,677;AA_AC=1937,1623;TAC=7274,2300;MAF=11.2571,45.5899,24.0234;GTS=GG,GA,AA;EA_GTC=2430,477,100;AA_GTC=784,369,627;GTC=3214,846,727;DP=69;GL=OR4F5;CP=1.0;CG=1.1;AA=G;CA=.;EXOME_CHIP=no;GWAS_PUBMED=.;GM=NM_001005484.1;FG=missense;AAC=ALA/THR;PP=141/306;CDP=421;GS=58;PH=benign
    """
    fetched = []
    exome_chip = False
    found = False
    info_map = {}
    acs = {}
    for hit in annotations_in_vcf(var, "esp", "vcf", "grch37"):
        if hit.contig not in ['Y']:
            fetched.append(hit)
            # We need a single ESP entry for a variant
            if fetched != None and len(fetched) == 1 and \
                    hit.alt == var.ALT[0] and hit.ref == var.REF:
                found = True
                # loads each VCF INFO key/value pair into a DICT
                for info in hit.info.split(";"):
                    if info.find("=") > 0:
                    # splits on first occurence of '='
                    # useful to handle valuerror: too many values to unpack (e.g (a,b) = split(",", (a,b,c,d)) for cases like
                    # SA=http://www.ncbi.nlm.nih.gov/sites/varvu?gene=4524&amp%3Brs=1801131|http://omim.org/entry/607093#0004
                        (key, value) = info.split("=", 1)
                        info_map[key] = value

                # NOTE:, if we start ot use GTS, need to update preprocessing
                # script to handle weirdness on X, Y
                # get the allele counts so that we can compute alternate allele frequencies
                # example: EA_AC=2,6764;AA_AC=23,3785;TAC=25,10549
                for key in ('EA_AC', 'AA_AC', 'TAC'):
                    if info_map.get(key) is not None:
                        lines = info_map[key].split(",")
                        denom = float(lines[0]) + float(lines[1])
                        if denom == 0:
                            acs[key] = 0
                        else:
                            # alt allele is stored as 2nd.
                            acs[key] = float(lines[1]) / denom
                    else:
                        acs[key] = -1

                # Is the SNP on an human exome chip?
                if info_map.get('EXOME_CHIP') is not None and \
                        info_map['EXOME_CHIP'] == "no":
                    exome_chip = 0
                elif info_map.get('EXOME_CHIP') is not None and \
                        info_map['EXOME_CHIP'] == "yes":
                    exome_chip = 1
                break
    return ESPInfo(found, acs.get('EA_AC', -1), acs.get("AA_AC", -1),
            acs.get("TAC", -1), exome_chip)


EMPTY_1000G = ThousandGInfo(False, -1, -1, -1, -1, -1, -1)
def get_1000G_info(var, empty=EMPTY_1000G):
    """
    Returns a suite of annotations from the 1000 Genomes project
    """
    #fetched = []
    info_map = {}

    for hit in annotations_in_vcf(var, "1000g", "vcf", "grch37"):
        # We need to ensure we are dealing with the exact sample variant
        # based on position and the alleles present.
        # var.start is used since the chromosomal pos in pysam.asVCF is zero based (hit.pos)
        # and would be equivalent to (POS-1) i.e. var.start
        if var.start == hit.pos and \
           var.ALT and len(var.ALT) > 0 and \
           var.ALT[0] == hit.alt and \
           hit.ref == var.REF:
            for info in hit.info.split(";"):
                if "=" in info:
                    (key, value) = info.split("=", 1)
                    info_map[key] = value

            return ThousandGInfo(True, float(info_map.get('AF', -1)), float(info_map.get('AMR_AF', -1)),
                         float(info_map.get('EAS_AF', -1)), float(info_map.get('SAS_AF', -1)),
                         float(info_map.get('AFR_AF', -1)), float(info_map.get('EUR_AF', -1)))
    return empty

def get_geno2mp_ct(var):

    for hit in annotations_in_vcf(var, "geno2mp", "vcf", "grch37"):
        if not (var.start == hit.pos and var.REF == hit.ref):
            continue
        if not var.ALT[0] in hit.alt.split(","): continue

        ct = next(x for x in hit.info.split(";") if x.startswith("HPO_CT="))
        val = int(ct.split("=")[1])
        return val
    # missing is -1
    return -1

def get_gnomad_info(var, empty=GNOMAD_EMPTY):
    afs = {}
    for hit in annotations_in_vcf(var, "gnomad", "vcf", "grch37"):
        # Does not handle anything beyond var.ALT[0] in the VCF (in case of multi-allelic variants)
        # var.start is used since the chromosomal pos in pysam.asVCF is zero based (hit.pos)
        # and would be equivalent to (POS-1) i.e var.start
        # TODO: change to REF
        if var.start != hit.start or var.REF != hit.REF or var.ALT[0] != hit.ALT[0]:
            continue

        # This would look for var.ALT[0] matches to
        # any of the multiple alt alleles represented in the EXAC file
        # Population independent raw (non-adjusted) allele frequencies given by AF
        info_map = hit.INFO
        aaf_ALL = info_map.get("AF", -1.0)

        for grp in ('_afr', '_amr', '_asj', '_eas', '_fin', '_nfe', '_oth', '_sas'):
            ac = info_map.get('AC' + grp)
            if ac is None: continue

            an = info_map.get('AN' + grp)
            if an is None: continue

            if an == 0:
                afs[grp] = 0
                continue

            afs[grp] = float(ac) / float(an)
        if ac is None or an is None: continue

        num_homs = int(info_map.get("nhomalt", -1))
        num_hets = int(info_map.get("AC", -1))
        if num_hets != -1:
          num_hets -= 2 * num_homs

        called_chroms = int(info_map.get('AN', -1))

        return GnomadInfo(aaf_ALL, float(afs['_afr']),
                              float(afs['_amr']), float(afs['_asj']),
                              float(afs['_eas']), float(afs['_fin']),
                              float(afs['_nfe']), float(afs['_oth']),
                              float(afs['_sas']),
                              info_map.get('non_neuro_AF', -1.0),
                              info_map.get('controls_AF', -1.0),
                              info_map.get('non_cancer_AF', -1.0),
                              info_map.get("AF_popmax", -1.0),
                              num_hets, num_homs,
                              called_chroms)

    return empty


{'fraction_non_alphanumeric': 0.11584887440400748, 'fraction_numerical': 0.026434908564186132, 'mean_word_length': 1.729470389588996, 'pattern_counts': {'":': 28, '<': 3, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 154, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6848093', 'n_tokens_mistral': 22855, 'n_tokens_neox': 21017, 'n_words': 4897}