<commit_before>from collections import defaultdict, OrderedDict
import itertools
import json
import os
import sys
import random
import string
import copy
import sys
from datetime import date, datetime

import pysam
import pymongo
from xbrowse.utils import compressed_file, get_progressbar
from xbrowse.utils import slugify

from xbrowse import utils as xbrowse_utils
from xbrowse import vcf_stuff
from xbrowse.core.variant_filters import VariantFilter, passes_variant_filter
from xbrowse import Variant
import datastore
from pprint import pprint, pformat

import StringIO
import elasticsearch
import elasticsearch_dsl
from elasticsearch_dsl import Q

from pyliftover import LiftOver

liftover = LiftOver('hg38', 'hg19')


#ELASTICSEARCH_HOST="172.16.58.3"

#ELASTICSEARCH_HOST="192.168.3.11"

# make encoded values as human-readable as possible
ES_FIELD_NAME_ESCAPE_CHAR = '$'
ES_FIELD_NAME_BAD_LEADING_CHARS = set(['_', '-', '+', ES_FIELD_NAME_ESCAPE_CHAR])
ES_FIELD_NAME_SPECIAL_CHAR_MAP = {
    '.': '_$dot$_',
    ',': '_$comma$_',
    '#': '_$hash$_',
    '*': '_$star$_',
    '(': '_$lp$_',
    ')': '_$rp$_',
    '[': '_$lsb$_',
    ']': '_$rsb$_',
    '{': '_$lcb$_',
    '}': '_$rcb$_',
}

import StringIO
import elasticsearch
import elasticsearch_dsl
from elasticsearch_dsl import Q


# make encoded values as human-readable as possible
ES_FIELD_NAME_ESCAPE_CHAR = '$'
ES_FIELD_NAME_BAD_LEADING_CHARS = set(['_', '-', '+', ES_FIELD_NAME_ESCAPE_CHAR])
ES_FIELD_NAME_SPECIAL_CHAR_MAP = {
    '.': '_$dot$_',
    ',': '_$comma$_',
    '#': '_$hash$_',
    '*': '_$star$_',
    '(': '_$lp$_',
    ')': '_$rp$_',
    '[': '_$lsb$_',
    ']': '_$rsb$_',
    '{': '_$lcb$_',
    '}': '_$rcb$_',
}


def _encode_field_name(s):
    """Encodes arbitrary string into an elasticsearch field name

    See:
    https://discuss.elastic.co/t/special-characters-in-field-names/10658/2
    https://discuss.elastic.co/t/illegal-characters-in-elasticsearch-field-names/17196/2
    """
    field_name = StringIO.StringIO()
    for i, c in enumerate(s):
        if c == ES_FIELD_NAME_ESCAPE_CHAR:
            field_name.write(2*ES_FIELD_NAME_ESCAPE_CHAR)
        elif c in ES_FIELD_NAME_SPECIAL_CHAR_MAP:
            field_name.write(ES_FIELD_NAME_SPECIAL_CHAR_MAP[c])  # encode the char
        else:
            field_name.write(c)  # write out the char as is

    field_name = field_name.getvalue()

    # escape 1st char if necessary
    if any(field_name.startswith(c) for c in ES_FIELD_NAME_BAD_LEADING_CHARS):
        return ES_FIELD_NAME_ESCAPE_CHAR + field_name
    else:
        return field_name


def _decode_field_name(field_name):
    """Converts an elasticsearch field name back to the original unencoded string"""

    if field_name.startswith(ES_FIELD_NAME_ESCAPE_CHAR):
        field_name = field_name[1:]

    i = 0
    original_string = StringIO.StringIO()
    while i < len(field_name):
        current_string = field_name[i:]
        if current_string.startswith(2*ES_FIELD_NAME_ESCAPE_CHAR):
            original_string.write(ES_FIELD_NAME_ESCAPE_CHAR)
            i += 2
        else:
            for original_value, encoded_value in ES_FIELD_NAME_SPECIAL_CHAR_MAP.items():
                if current_string.startswith(encoded_value):
                    original_string.write(original_value)
                    i += len(encoded_value)
                    break
            else:
                original_string.write(field_name[i])
                i += 1

    return original_string.getvalue()

MONGO_QUERY_RESULTS_LIMIT = 6000

GENOTYPE_QUERY_MAP = {

    'ref_ref': 0,
    'ref_alt': 1,
    'alt_alt': 2,

    'has_alt': {'$gte': 1},
    'has_ref': {'$in': [0,1]},

    'not_missing': {'$gte': 0},
    'missing': -1,
}


CHROMOSOME_SIZES = { "1":249250621, "2":243199373, "3":198022430, "4":191154276, "5":180915260, "6":171115067, "7":159138663, "8":146364022, "9":141213431, "10":135534747, "11":135006516, "12":133851895, "13":115169878, "14":107349540, "15":102531392, "16":90354753, "17":81195210, "18":78077248, "19":59128983, "20":63025520, "21":48129895, "22":51304566, "X":155270560, "Y":59373566, "MT":16569, } # used for stats

def _add_genotype_filter_to_variant_query(db_query, genotype_filter):
    """
    Add conditions to db_query from the genotype filter
    Edits in place, returns True if successful
    """
    for indiv_id, genotype in genotype_filter.items():
        key = 'genotypes.%s.num_alt' % indiv_id
        db_query[key] = GENOTYPE_QUERY_MAP[genotype]
    return True


def _add_index_fields_to_variant(variant_dict, annotation=None):
    """
    Add fields to the vairant dictionary that you want to index on before load it
    """
    if annotation:
        variant_dict['db_freqs'] = annotation['freqs']
        variant_dict['db_tags'] = annotation['annotation_tags']
        variant_dict['db_gene_ids'] = annotation['gene_ids']


class MongoDatastore(datastore.Datastore):

    def __init__(self, db, annotator, custom_population_store=None, custom_populations_map=None):
        self._db = db
        self._annotator = annotator
        self._custom_population_store = custom_population_store
        self._custom_populations_map = custom_populations_map
        if self._custom_populations_map is None:
            self._custom_populations_map = {}

    def _make_db_query(self, genotype_filter=None, variant_filter=None):
        """
        Caller specifies filters to get_variants, but they are evaluated later.
        Here, we just inspect those filters and see what heuristics we can apply to avoid a full table scan,
        Query here must return a superset of the true get_variants results
        Note that the full annotation isn't stored, so use the fields added by _add_index_fields_to_variant
        """
        db_query = {}

        # genotype filter
        if genotype_filter is not None:
            _add_genotype_filter_to_variant_query(db_query, genotype_filter)

        if variant_filter:
            if variant_filter.locations:
                location_ranges = []
                for xstart, xend in variant_filter.locations:
                    location_ranges.append({'$and' : [ {'xpos' : {'$gte': xstart }}, {'xpos' : {'$lte': xend }}] })
                db_query['$or'] = location_ranges

            if variant_filter.so_annotations:
                db_query['db_tags'] = {'$in': variant_filter.so_annotations}
            if variant_filter.genes:
                db_query['db_gene_ids'] = {'$in': variant_filter.genes}
                db_query['db_exclude_genes'] = getattr(variant_filter, 'exclude_genes')
            if variant_filter.ref_freqs:
                for population, freq in variant_filter.ref_freqs:
                    if population in self._annotator.reference_population_slugs:
                        db_query['db_freqs.' + population] = {'$lte': freq}

        return db_query

    
    def get_elasticsearch_variants(self, query_json, project_id, family_id, variant_id_filter=None):
        from xbrowse_server.base.models import Individual

        ELASTICSEARCH_PORT="9200"
        if project_id == "ATGU_WGS-Jueppner":
            ELASTICSEARCH_HOST="10.4.3.24"
        else:
            ELASTICSEARCH_HOST="10.48.5.5"
        client = elasticsearch.Elasticsearch(ELASTICSEARCH_HOST, port=ELASTICSEARCH_PORT)

        if project_id == "Engle_WGS_900":
            indices = ["engle_wgs_900_samples__*coding_0", "engle_wgs_900_samples__*coding_1", "engle_wgs_900_samples__*coding_2", "engle_wgs_900_samples__*coding_4"]
        elif project_id == "rare_genomes_project":
            indices = ["rgp_callset_2"] #["rare_genomes_project"]
        elif project_id == "Engle_WGS_2_sample":
            indices = ["engle_wgs_2_sample"]
        elif project_id == "NIAID-gatk3dot4":
            indices = ["niaid-gatk3dot4"]
        elif project_id == "ATGU_WGS-Jueppner":
            indices = ["atgu_wgs_jueppner"]
        else:
            raise ValueError("Unexpected project_id: " + str(project_id))

        if family_id is not None: 
            family_individual_ids = [i.indiv_id for i in Individual.objects.filter(family__family_id=family_id)]
            # figure out which indices contain this family
            if len(indices) == 1:
                index_to_search = indices[0]
            else:
                index_to_search = None
                for index_name in indices:
                    mapping = elasticsearch_dsl.Index(index_name, using=client).get_mapping()
                    mapping_subset = {k: v for k, v in mapping[index_name.replace("*", "")]["mappings"]["variant"]["properties"].items() if k.startswith("S") or k.startswith("RGP") or k.startswith("APY-001")}

                    corrected_indiv_id = family_individual_ids[0].replace("S132-4_1_1", "S132-4.1.1").replace("S218-3_1", "S218-3.1").replace("S250-1_1", "S250-1.1").replace("S258-1_1", "S258-1.1")
                    genotype_field_name = "%s_num_alt" % _encode_field_name(corrected_indiv_id)
                    if genotype_field_name in mapping_subset:
                        index_to_search = index_name
                        break
                    else:
                        print("%s not found in %s" % (genotype_field_name, index_name))
                else:
                    raise ValueError("ERROR: family: %s not found in database" % family_id)

            # do the query 
            s = elasticsearch_dsl.Search(using=client, index=index_to_search)
        else:
            family_individual_ids = [i.indiv_id for i in Individual.objects.filter(project__project_id=project_id)]

            if project_id == "Engle_WGS_900":
                indices = ["engle_wgs_900_samples__*"]
            elif project_id == "rare_genomes_project":
                indices = ["rare_genomes_project__*coding"]
            elif project_id == "Engle_WGS_2_sample":
                indices = ["engle_wgs_2_sample"]
            elif project_id == "NIAID-gatk3dot4":
                indices = ["niaid-gatk3dot4"]
            elif project_id == "ATGU_WGS-Jueppner":
                indices = ["atgu_wgs_jueppner"]
            s = elasticsearch_dsl.Search(using=client, index=indices[0]) #",".join(indices))

        print("===> QUERY: ")
        pprint(query_json)

        if variant_id_filter is not None:
            s = s.filter('term', **{"variantId": variant_id_filter})

        # parse variant query
        #query_json = query_json.get('$and', [])
        for key, value in query_json.items():
            if key == 'db_tags':
                vep_consequences = query_json.get('db_tags', {}).get('$in', [])
                s = s.filter("terms",  transcriptConsequenceTerms=vep_consequences)
                print("==> transcriptConsequenceTerms: %s" % str(vep_consequences))

            if key.startswith("genotypes"):
                sample_id = ".".join(key.split(".")[1:-1])
                encoded_sample_id = _encode_field_name(sample_id)
                genotype_filter = value
                if type(genotype_filter) == int or type(genotype_filter) == basestring:
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": genotype_filter}))
                    s = s.filter('term', **{encoded_sample_id+"_num_alt": genotype_filter})

                elif '$gte' in genotype_filter:
                    genotype_filter = {k.replace("$", ""): v for k, v in genotype_filter.items()}
                    s = s.filter('range', **{encoded_sample_id+"_num_alt": genotype_filter})
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": genotype_filter}))
                elif "$in" in genotype_filter:
                    num_alt_values = genotype_filter['$in']
                    q = Q('term', **{encoded_sample_id+"_num_alt": num_alt_values[0]})
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": num_alt_values[0]}))
                    for num_alt_value in num_alt_values[1:]:
                        q = q | Q('term', **{encoded_sample_id+"_num_alt": num_alt_value})
                        print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": num_alt_value}))
                    s = s.filter(q)

            if key == "db_gene_ids":
                gene_ids = query_json.get('db_gene_ids', {}).get('$in', [])
                exclude_genes = query_json.get('db_exclude_genes')

                if exclude_genes:
                    s = s.exclude("terms", geneIds=gene_ids)
                else:
                    s =  s.filter("terms",  geneIds=gene_ids)
                print("==> %s %s" % ("exclude" if exclude_genes else "include", "geneIds: " + str(gene_ids)))

            if key == "$or" and type(value) == list:
                xpos_filters = value[0].get("$and", {})

                # for example: $or : [{'$and': [{'xpos': {'$gte': 12345}}, {'xpos': {'$lte': 54321}}]}]
                xpos_filters_dict = {}
                for xpos_filter in xpos_filters:
                    xpos_filter_setting = xpos_filter["xpos"]  # for example {'$gte': 12345} or {'$lte': 54321}
                    xpos_filters_dict.update(xpos_filter_setting)
                xpos_filter_setting = {k.replace("$", ""): v for k, v in xpos_filters_dict.items()}
                s = s.filter('range', **{"xpos": xpos_filter_setting})
                print("==> xpos range: " + str({"xpos": xpos_filter_setting}))


            af_key_map = {
                "db_freqs.1kg_wgs_phase3": "g1k_AF",
                "db_freqs.1kg_wgs_phase3_popmax": "g1k_POPMAX_AF",
                "db_freqs.exac_v3": "exac_AF",
                "db_freqs.exac_v3_popmax": "exac_AF_POPMAX",
                "db_freqs.topmed": "topmed_AF",
                "db_freqs.gnomad_exomes": "gnomad_exomes_AF",
                "db_freqs.gnomad_exomes_popmax": "gnomad_exomes_AF_POPMAX",
                "db_freqs.gnomad_genomes": "gnomad_genomes_AF",
                "db_freqs.gnomad_genomes_popmax": "gnomad_genomes_AF_POPMAX",
            }

            if key in af_key_map:
                filter_key = af_key_map[key]
                af_filter_setting = {k.replace("$", ""): v for k, v in value.items()}
                s = s.filter(Q('range', **{filter_key: af_filter_setting}) | ~Q('exists', field=filter_key))
                print("==> %s: %s" % (filter_key, af_filter_setting))

            s.sort("xpos")

        print("=====")
        print("FULL QUERY OBJ: " + pformat(s.__dict__))
        print("FILTERS: " + pformat(s.to_dict()))
        print("=====")
        print("Hits: ")
        # https://elasticsearch-py.readthedocs.io/en/master/helpers.html#elasticsearch.helpers.scan

        response = s.execute()
        print("TOTAL: " + str(response.hits.total))
        #print(pformat(response.to_dict()))
        for i, hit in enumerate(s.scan()):  # preserve_order=True
            #print("##### Raw result " + str(i) + "  "+ str(hit.meta))
            #pprint(hit.__dict__)
            filters = ",".join(hit["filters"]) if "filters" in hit else ""
            genotypes = {}
            all_num_alt = []
            for individual_id in family_individual_ids:
                encoded_individual_id = _encode_field_name(individual_id)
                num_alt =  int(hit["%s_num_alt" % encoded_individual_id]) if ("%s_num_alt" % encoded_individual_id) in hit else -1
                if num_alt is not None:
                    all_num_alt.append(num_alt)
                    
                alleles = []
                if num_alt == 0:
                    alleles = [hit["ref"], hit["ref"]]
                elif num_alt == 1:
                    alleles = [hit["ref"], hit["alt"]]
                elif num_alt == 2:
                    alleles = [hit["alt"], hit["alt"]]
                elif num_alt == -1 or num_alt == None:
                    alleles = []
                else:
                    raise ValueError("Invalid num_alt: " + str(num_alt))

                genotypes[individual_id] = {
                    'ab': hit["%s_ab" % encoded_individual_id] if ("%s_ab" % encoded_individual_id) in hit else '',
                    'alleles': map(str, alleles),
                    'extras': {
                        'ad': hit["%s_ab" % encoded_individual_id]  if ("%s_ad" % encoded_individual_id) in hit else '',
                        'dp': hit["%s_dp" % encoded_individual_id]  if ("%s_dp" % encoded_individual_id) in hit else '',
                        'pl': '',
                    },
                    'filter': filters or "pass",
                    'gq': hit["%s_gq" % encoded_individual_id] if ("%s_gq" % encoded_individual_id in hit and hit["%s_gq" % encoded_individual_id] is not None) else '',
                    'num_alt': num_alt,
                }

            if all([num_alt <= 0 for num_alt in all_num_alt]):
                #print("Filtered out due to genotype: " + str(genotypes))
                print("Filtered all_num_alt <= 0 - Result %s: GRCh38: %s:%s,  cadd: %s  %s - %s" % (i, hit["contig"], hit["start"], hit["cadd_PHRED"] if "cadd_PHRED" in hit else "", hit["transcriptConsequenceTerms"], all_num_alt))
                continue
            
            vep_annotation = json.loads(str(hit['sortedTranscriptConsequences']))
            
            if project_id not in ["NIAID-gatk3dot4"]:
                lifted_over_coord = liftover.convert_coordinate("chr%s" % hit["contig"].replace("chr", ""), int(hit["start"])), #
                if lifted_over_coord and lifted_over_coord[0] and lifted_over_coord[0][0]:
                    lifted_over_coord = "%s-%s-%s-%s "% (lifted_over_coord[0][0][0], lifted_over_coord[0][0][1], hit["ref"],  hit["alt"])
                else:
                    lifted_over_coord = None
            else:
                lifted_over_coord = hit["variantId"]
                                                     
            result = {
                #u'_id': ObjectId('596d2207ff66f729285ca588'),
                'alt': str(hit["alt"]) if "alt" in hit else None,
                'annotation': {
                    'fathmm': None,
                    'metasvm': None,
                    'muttaster': None,
                    'polyphen': None,
                    'sift': None,
                    'cadd_phred': hit["cadd_PHRED"] if "cadd_PHRED" in hit else None,
                    'annotation_tags': list(hit["transcriptConsequenceTerms"]),
                    'coding_gene_ids': list(hit['geneIds']),
                    'gene_ids': list(hit['geneIds']),
                    'vep_annotation': vep_annotation,
                    'vep_group': str(hit['mainTranscript_major_consequence']),
                    'vep_consequence': str(hit['mainTranscript_major_consequence']),
                    'worst_vep_annotation_index': 0,
                    'worst_vep_index_per_gene': {str(hit['mainTranscript_gene_id']): 0},
                },
                'chr': hit["contig"],
                'coding_gene_ids': None,
                'db_freqs': {
                    '1kg_wgs_AF': float(hit["g1k_AF"] or 0.0),
                    '1kg_wgs_popmax_AF': float(hit["g1k_POPMAX_AF"] or 0.0),
                    'exac_v3_AF': float(hit["exac_AF"] or 0.0) if "exac_AF" in hit else (hit["exac_AC_Adj"]/float(hit["exac_AN_Adj"]) if int(hit["exac_AN_Adj"] or 0) > 0 else 0.0),
                    'exac_v3_popmax_AF': float(hit["exac_AF_POPMAX"] or 0.0) if "exac_AF_POPMAX" in hit else 0.0,
                    'topmed_AF': float(hit["topmed_AF"] or 0.0) if "topmed_AF" in hit else 0.0,
                    'gnomad_exomes_AF': float(hit["gnomad_exomes_AF"] or 0.0),
                    'gnomad_exomes_popmax_AF': float(hit["gnomad_exomes_AF_POPMAX"] or 0.0),
                    'gnomad_genomes_AF': float(hit["gnomad_genomes_AF"] or 0.0),
                    'gnomad_genomes_popmax_AF': float(hit["gnomad_genomes_AF_POPMAX"] or 0.0),
                },
                'db_gene_ids': list(hit["geneIds"]),
                'db_tags': str(hit["transcriptConsequenceTerms"]),
                'extras': {
                    'grch37_coords': lifted_over_coord,
                    'grch38_coords': "%s-%s-%s-%s" % (hit["contig"], hit["start"], hit["ref"], hit["alt"]),
                    u'alt_allele_pos': 0,
                    u'orig_alt_alleles': map(str, [a.split("-")[-1] for a in hit["originalAltAlleles"]]) if "originalAltAlleles" in hit else []},
                'gene_ids': None,
                'genotypes': genotypes,
                'pos': long(hit['start']),<commit_msg>added revel and dann annotations<commit_after>from collections import defaultdict, OrderedDict
import itertools
import json
import os
import sys
import random
import string
import copy
import sys
from datetime import date, datetime

import pysam
import pymongo
from xbrowse.utils import compressed_file, get_progressbar
from xbrowse.utils import slugify

from xbrowse import utils as xbrowse_utils
from xbrowse import vcf_stuff
from xbrowse.core.variant_filters import VariantFilter, passes_variant_filter
from xbrowse import Variant
import datastore
from pprint import pprint, pformat

import StringIO
import elasticsearch
import elasticsearch_dsl
from elasticsearch_dsl import Q

from pyliftover import LiftOver

liftover = LiftOver('hg38', 'hg19')


#ELASTICSEARCH_HOST="172.16.58.3"

#ELASTICSEARCH_HOST="192.168.3.11"

# make encoded values as human-readable as possible
ES_FIELD_NAME_ESCAPE_CHAR = '$'
ES_FIELD_NAME_BAD_LEADING_CHARS = set(['_', '-', '+', ES_FIELD_NAME_ESCAPE_CHAR])
ES_FIELD_NAME_SPECIAL_CHAR_MAP = {
    '.': '_$dot$_',
    ',': '_$comma$_',
    '#': '_$hash$_',
    '*': '_$star$_',
    '(': '_$lp$_',
    ')': '_$rp$_',
    '[': '_$lsb$_',
    ']': '_$rsb$_',
    '{': '_$lcb$_',
    '}': '_$rcb$_',
}

import StringIO
import elasticsearch
import elasticsearch_dsl
from elasticsearch_dsl import Q


# make encoded values as human-readable as possible
ES_FIELD_NAME_ESCAPE_CHAR = '$'
ES_FIELD_NAME_BAD_LEADING_CHARS = set(['_', '-', '+', ES_FIELD_NAME_ESCAPE_CHAR])
ES_FIELD_NAME_SPECIAL_CHAR_MAP = {
    '.': '_$dot$_',
    ',': '_$comma$_',
    '#': '_$hash$_',
    '*': '_$star$_',
    '(': '_$lp$_',
    ')': '_$rp$_',
    '[': '_$lsb$_',
    ']': '_$rsb$_',
    '{': '_$lcb$_',
    '}': '_$rcb$_',
}


def _encode_field_name(s):
    """Encodes arbitrary string into an elasticsearch field name

    See:
    https://discuss.elastic.co/t/special-characters-in-field-names/10658/2
    https://discuss.elastic.co/t/illegal-characters-in-elasticsearch-field-names/17196/2
    """
    field_name = StringIO.StringIO()
    for i, c in enumerate(s):
        if c == ES_FIELD_NAME_ESCAPE_CHAR:
            field_name.write(2*ES_FIELD_NAME_ESCAPE_CHAR)
        elif c in ES_FIELD_NAME_SPECIAL_CHAR_MAP:
            field_name.write(ES_FIELD_NAME_SPECIAL_CHAR_MAP[c])  # encode the char
        else:
            field_name.write(c)  # write out the char as is

    field_name = field_name.getvalue()

    # escape 1st char if necessary
    if any(field_name.startswith(c) for c in ES_FIELD_NAME_BAD_LEADING_CHARS):
        return ES_FIELD_NAME_ESCAPE_CHAR + field_name
    else:
        return field_name


def _decode_field_name(field_name):
    """Converts an elasticsearch field name back to the original unencoded string"""

    if field_name.startswith(ES_FIELD_NAME_ESCAPE_CHAR):
        field_name = field_name[1:]

    i = 0
    original_string = StringIO.StringIO()
    while i < len(field_name):
        current_string = field_name[i:]
        if current_string.startswith(2*ES_FIELD_NAME_ESCAPE_CHAR):
            original_string.write(ES_FIELD_NAME_ESCAPE_CHAR)
            i += 2
        else:
            for original_value, encoded_value in ES_FIELD_NAME_SPECIAL_CHAR_MAP.items():
                if current_string.startswith(encoded_value):
                    original_string.write(original_value)
                    i += len(encoded_value)
                    break
            else:
                original_string.write(field_name[i])
                i += 1

    return original_string.getvalue()

MONGO_QUERY_RESULTS_LIMIT = 6000

GENOTYPE_QUERY_MAP = {

    'ref_ref': 0,
    'ref_alt': 1,
    'alt_alt': 2,

    'has_alt': {'$gte': 1},
    'has_ref': {'$in': [0,1]},

    'not_missing': {'$gte': 0},
    'missing': -1,
}


CHROMOSOME_SIZES = { "1":249250621, "2":243199373, "3":198022430, "4":191154276, "5":180915260, "6":171115067, "7":159138663, "8":146364022, "9":141213431, "10":135534747, "11":135006516, "12":133851895, "13":115169878, "14":107349540, "15":102531392, "16":90354753, "17":81195210, "18":78077248, "19":59128983, "20":63025520, "21":48129895, "22":51304566, "X":155270560, "Y":59373566, "MT":16569, } # used for stats

def _add_genotype_filter_to_variant_query(db_query, genotype_filter):
    """
    Add conditions to db_query from the genotype filter
    Edits in place, returns True if successful
    """
    for indiv_id, genotype in genotype_filter.items():
        key = 'genotypes.%s.num_alt' % indiv_id
        db_query[key] = GENOTYPE_QUERY_MAP[genotype]
    return True


def _add_index_fields_to_variant(variant_dict, annotation=None):
    """
    Add fields to the vairant dictionary that you want to index on before load it
    """
    if annotation:
        variant_dict['db_freqs'] = annotation['freqs']
        variant_dict['db_tags'] = annotation['annotation_tags']
        variant_dict['db_gene_ids'] = annotation['gene_ids']


class MongoDatastore(datastore.Datastore):

    def __init__(self, db, annotator, custom_population_store=None, custom_populations_map=None):
        self._db = db
        self._annotator = annotator
        self._custom_population_store = custom_population_store
        self._custom_populations_map = custom_populations_map
        if self._custom_populations_map is None:
            self._custom_populations_map = {}

    def _make_db_query(self, genotype_filter=None, variant_filter=None):
        """
        Caller specifies filters to get_variants, but they are evaluated later.
        Here, we just inspect those filters and see what heuristics we can apply to avoid a full table scan,
        Query here must return a superset of the true get_variants results
        Note that the full annotation isn't stored, so use the fields added by _add_index_fields_to_variant
        """
        db_query = {}

        # genotype filter
        if genotype_filter is not None:
            _add_genotype_filter_to_variant_query(db_query, genotype_filter)

        if variant_filter:
            if variant_filter.locations:
                location_ranges = []
                for xstart, xend in variant_filter.locations:
                    location_ranges.append({'$and' : [ {'xpos' : {'$gte': xstart }}, {'xpos' : {'$lte': xend }}] })
                db_query['$or'] = location_ranges

            if variant_filter.so_annotations:
                db_query['db_tags'] = {'$in': variant_filter.so_annotations}
            if variant_filter.genes:
                db_query['db_gene_ids'] = {'$in': variant_filter.genes}
                db_query['db_exclude_genes'] = getattr(variant_filter, 'exclude_genes')
            if variant_filter.ref_freqs:
                for population, freq in variant_filter.ref_freqs:
                    if population in self._annotator.reference_population_slugs:
                        db_query['db_freqs.' + population] = {'$lte': freq}

        return db_query

    
    def get_elasticsearch_variants(self, query_json, project_id, family_id, variant_id_filter=None):
        from xbrowse_server.base.models import Individual

        ELASTICSEARCH_PORT="9200"
        if project_id == "ATGU_WGS-Jueppner":
            ELASTICSEARCH_HOST="10.4.3.24"
        elif project_id == "rare_genomes_project":
            ELASTICSEARCH_HOST="192.168.3.11"
        else:
            ELASTICSEARCH_HOST="10.48.5.14"
        client = elasticsearch.Elasticsearch(ELASTICSEARCH_HOST, port=ELASTICSEARCH_PORT)

        if project_id == "Engle_WGS_900":
            indices = ["engle_wgs_900_samples__*coding_0", "engle_wgs_900_samples__*coding_1", "engle_wgs_900_samples__*coding_2", "engle_wgs_900_samples__*coding_4"]
        elif project_id == "rare_genomes_project":
            indices = ["rgp_callset_2"] #["rare_genomes_project"]
        elif project_id == "Engle_WGS_2_sample":
            indices = ["engle_wgs_2_sample"]
        elif project_id == "NIAID-gatk3dot4":
            indices = ["niaid-gatk3dot4"]
        elif project_id == "ATGU_WGS-Jueppner":
            indices = ["atgu_wgs_jueppner"]
        else:
            raise ValueError("Unexpected project_id: " + str(project_id))

        if family_id is not None: 
            family_individual_ids = [i.indiv_id for i in Individual.objects.filter(family__family_id=family_id)]
            # figure out which indices contain this family
            if len(indices) == 1:
                index_to_search = indices[0]
            else:
                index_to_search = None
                for index_name in indices:
                    mapping = elasticsearch_dsl.Index(index_name, using=client).get_mapping()
                    mapping_subset = {k: v for k, v in mapping[index_name.replace("*", "")]["mappings"]["variant"]["properties"].items() if k.startswith("S") or k.startswith("RGP") or k.startswith("APY-001")}

                    corrected_indiv_id = family_individual_ids[0].replace("S132-4_1_1", "S132-4.1.1").replace("S218-3_1", "S218-3.1").replace("S250-1_1", "S250-1.1").replace("S258-1_1", "S258-1.1")
                    genotype_field_name = "%s_num_alt" % _encode_field_name(corrected_indiv_id)
                    if genotype_field_name in mapping_subset:
                        index_to_search = index_name
                        break
                    else:
                        print("%s not found in %s" % (genotype_field_name, index_name))
                else:
                    raise ValueError("ERROR: family: %s not found in database" % family_id)

            # do the query 
            s = elasticsearch_dsl.Search(using=client, index=index_to_search)
        else:
            family_individual_ids = [i.indiv_id for i in Individual.objects.filter(project__project_id=project_id)]

            if project_id == "Engle_WGS_900":
                indices = ["engle_wgs_900_samples__*"]
            elif project_id == "rare_genomes_project":
                indices = ["rare_genomes_project__*coding"]
            elif project_id == "Engle_WGS_2_sample":
                indices = ["engle_wgs_2_sample"]
            elif project_id == "NIAID-gatk3dot4":
                indices = ["niaid-gatk3dot4"]
            elif project_id == "ATGU_WGS-Jueppner":
                indices = ["atgu_wgs_jueppner"]
            s = elasticsearch_dsl.Search(using=client, index=indices[0]) #",".join(indices))

        print("===> QUERY: ")
        pprint(query_json)

        if variant_id_filter is not None:
            s = s.filter('term', **{"variantId": variant_id_filter})

        # parse variant query
        #query_json = query_json.get('$and', [])
        for key, value in query_json.items():
            if key == 'db_tags':
                vep_consequences = query_json.get('db_tags', {}).get('$in', [])
                s = s.filter("terms",  transcriptConsequenceTerms=vep_consequences)
                print("==> transcriptConsequenceTerms: %s" % str(vep_consequences))

            if key.startswith("genotypes"):
                sample_id = ".".join(key.split(".")[1:-1])
                encoded_sample_id = _encode_field_name(sample_id)
                genotype_filter = value
                if type(genotype_filter) == int or type(genotype_filter) == basestring:
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": genotype_filter}))
                    s = s.filter('term', **{encoded_sample_id+"_num_alt": genotype_filter})

                elif '$gte' in genotype_filter:
                    genotype_filter = {k.replace("$", ""): v for k, v in genotype_filter.items()}
                    s = s.filter('range', **{encoded_sample_id+"_num_alt": genotype_filter})
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": genotype_filter}))
                elif "$in" in genotype_filter:
                    num_alt_values = genotype_filter['$in']
                    q = Q('term', **{encoded_sample_id+"_num_alt": num_alt_values[0]})
                    print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": num_alt_values[0]}))
                    for num_alt_value in num_alt_values[1:]:
                        q = q | Q('term', **{encoded_sample_id+"_num_alt": num_alt_value})
                        print("==> genotypes: %s" % str({encoded_sample_id+"_num_alt": num_alt_value}))
                    s = s.filter(q)

            if key == "db_gene_ids":
                gene_ids = query_json.get('db_gene_ids', {}).get('$in', [])
                exclude_genes = query_json.get('db_exclude_genes')

                if exclude_genes:
                    s = s.exclude("terms", geneIds=gene_ids)
                else:
                    s =  s.filter("terms",  geneIds=gene_ids)
                print("==> %s %s" % ("exclude" if exclude_genes else "include", "geneIds: " + str(gene_ids)))

            if key == "$or" and type(value) == list:
                xpos_filters = value[0].get("$and", {})

                # for example: $or : [{'$and': [{'xpos': {'$gte': 12345}}, {'xpos': {'$lte': 54321}}]}]
                xpos_filters_dict = {}
                for xpos_filter in xpos_filters:
                    xpos_filter_setting = xpos_filter["xpos"]  # for example {'$gte': 12345} or {'$lte': 54321}
                    xpos_filters_dict.update(xpos_filter_setting)
                xpos_filter_setting = {k.replace("$", ""): v for k, v in xpos_filters_dict.items()}
                s = s.filter('range', **{"xpos": xpos_filter_setting})
                print("==> xpos range: " + str({"xpos": xpos_filter_setting}))


            af_key_map = {
                "db_freqs.1kg_wgs_phase3": "g1k_AF",
                "db_freqs.1kg_wgs_phase3_popmax": "g1k_POPMAX_AF",
                "db_freqs.exac_v3": "exac_AF",
                "db_freqs.exac_v3_popmax": "exac_AF_POPMAX",
                "db_freqs.topmed": "topmed_AF",
                "db_freqs.gnomad_exomes": "gnomad_exomes_AF",
                "db_freqs.gnomad_exomes_popmax": "gnomad_exomes_AF_POPMAX",
                "db_freqs.gnomad_genomes": "gnomad_genomes_AF",
                "db_freqs.gnomad_genomes_popmax": "gnomad_genomes_AF_POPMAX",
            }

            if key in af_key_map:
                filter_key = af_key_map[key]
                af_filter_setting = {k.replace("$", ""): v for k, v in value.items()}
                s = s.filter(Q('range', **{filter_key: af_filter_setting}) | ~Q('exists', field=filter_key))
                print("==> %s: %s" % (filter_key, af_filter_setting))

            s.sort("xpos")

        print("=====")
        print("FULL QUERY OBJ: " + pformat(s.__dict__))
        print("FILTERS: " + pformat(s.to_dict()))
        print("=====")
        print("Hits: ")
        # https://elasticsearch-py.readthedocs.io/en/master/helpers.html#elasticsearch.helpers.scan

        response = s.execute()
        print("TOTAL: " + str(response.hits.total))
        #print(pformat(response.to_dict()))
        for i, hit in enumerate(s.scan()):  # preserve_order=True
            if i == 0:
                print("Hit columns: " + str(hit.__dict__))
            #print("##### Raw result " + str(i) + "  "+ str(hit.meta))
            #pprint(hit.__dict__)
            filters = ",".join(hit["filters"]) if "filters" in hit else ""
            genotypes = {}
            all_num_alt = []
            for individual_id in family_individual_ids:
                encoded_individual_id = _encode_field_name(individual_id)
                num_alt =  int(hit["%s_num_alt" % encoded_individual_id]) if ("%s_num_alt" % encoded_individual_id) in hit else -1
                if num_alt is not None:
                    all_num_alt.append(num_alt)
                    
                alleles = []
                if num_alt == 0:
                    alleles = [hit["ref"], hit["ref"]]
                elif num_alt == 1:
                    alleles = [hit["ref"], hit["alt"]]
                elif num_alt == 2:
                    alleles = [hit["alt"], hit["alt"]]
                elif num_alt == -1 or num_alt == None:
                    alleles = []
                else:
                    raise ValueError("Invalid num_alt: " + str(num_alt))

                genotypes[individual_id] = {
                    'ab': hit["%s_ab" % encoded_individual_id] if ("%s_ab" % encoded_individual_id) in hit else '',
                    'alleles': map(str, alleles),
                    'extras': {
                        'ad': hit["%s_ab" % encoded_individual_id]  if ("%s_ad" % encoded_individual_id) in hit else '',
                        'dp': hit["%s_dp" % encoded_individual_id]  if ("%s_dp" % encoded_individual_id) in hit else '',
                        'pl': '',
                    },
                    'filter': filters or "pass",
                    'gq': hit["%s_gq" % encoded_individual_id] if ("%s_gq" % encoded_individual_id in hit and hit["%s_gq" % encoded_individual_id] is not None) else '',
                    'num_alt': num_alt,
                }

            if all([num_alt <= 0 for num_alt in all_num_alt]):
                #print("Filtered out due to genotype: " + str(genotypes))
                print("Filtered all_num_alt <= 0 - Result %s: GRCh38: %s:%s,  cadd: %s  %s - %s" % (i, hit["contig"], hit["start"], hit["cadd_PHRED"] if "cadd_PHRED" in hit else "", hit["transcriptConsequenceTerms"], all_num_alt))
                continue
            
            vep_annotation = json.loads(str(hit['sortedTranscriptConsequences']))
            
            if project_id not in ["NIAID-gatk3dot4"]:
                lifted_over_coord = liftover.convert_coordinate("chr%s" % hit["contig"].replace("chr", ""), int(hit["start"])), #
                if lifted_over_coord and lifted_over_coord[0] and lifted_over_coord[0][0]:
                    lifted_over_coord = "%s-%s-%s-%s "% (lifted_over_coord[0][0][0], lifted_over_coord[0][0][1], hit["ref"],  hit["alt"])
                else:
                    lifted_over_coord = None
            else:
                lifted_over_coord = hit["variantId"]
                                                     
            result = {
                #u'_id': ObjectId('596d2207ff66f729285ca588'),
                'alt': str(hit["alt"]) if "alt" in hit else None,
                'annotation': {
                    'fathmm': None,
                    'metasvm': None,
                    'muttaster': None,
                    'polyphen': None,
                    'sift': None,

                    'cadd_phred': hit["cadd_PHRED"] if "cadd_PHRED" in hit else None,
                    'dann_score': hit["dbnsfp_DANN_score"] if "dbnsfp_DANN_score" in hit else None,
                    'revel_score': hit["dbnsfp_REVEL_score"] if "dbnsfp_REVEL_score" in hit else None,

                    'annotation_tags': list(hit["transcriptConsequenceTerms"]) if "transcriptConsequenceTerms" in hit else None,
                    'coding_gene_ids': list(hit['geneIds']),
                    'gene_ids': list(hit['geneIds']),
                    'vep_annotation': vep_annotation,
                    'vep_group': str(hit['mainTranscript_major_consequence']),
                    'vep_consequence': str(hit['mainTranscript_major_consequence']),
                    'worst_vep_annotation_index': 0,
                    'worst_vep_index_per_gene': {str(hit['mainTranscript_gene_id']): 0},
                },
                'chr': hit["contig"],
                'coding_gene_ids': None,
                'db_freqs': {
                    '1kg_wgs_AF': float(hit["g1k_AF"] or 0.0),
                    '1kg_wgs_popmax_AF': float(hit["g1k_POPMAX_AF"] or 0.0),
                    'exac_v3_AC': float(hit["exac_AC_Adj"] or 0.0) if "exac_AC_Adj" in hit else 0.0,
                    'exac_v3_AF': float(hit["exac_AF"] or 0.0) if "exac_AF" in hit else (hit["exac_AC_Adj"]/float(hit["exac_AN_Adj"]) if int(hit["exac_AN_Adj"] or 0) > 0 else 0.0),                    
                    'exac_v3_popmax_AF': float(hit["exac_AF_POPMAX"] or 0.0) if "exac_AF_POPMAX" in hit else 0.0,
                    'topmed_AF': float(hit["topmed_AF"] or 0.0) if "topmed_AF" in hit else 0.0,
                    'gnomad_exomes_AC': float(hit["gnomad_exomes_AC"] or 0.0) if "gnomad_exomes_AC" in hit else 0.0,
                    'gnomad_exomes_Hom': float(hit["gnomad_exomes_HOM"] or 0.0) if "gnomad_exomes_HOM" in hit else 0.0,
                    'gnomad_exomes_AF': float(hit["gnomad_exomes_AF"] or 0.0) if "gnomad_exomes_AF" in hit else 0.0,
                    'gnomad_exomes_popmax_AF': float(hit["gnomad_exomes_AF_POPMAX"] or 0.0) if "gnomad_exomes_AF_POPMAX" in hit else 0.0,
                    'gnomad_genomes_AC': float(hit["gnomad_genomes_AC"] or 0.0) if "gnomad_genomes_AC" in hit else 0.0,
                    'gnomad_genomes_Hom': float(hit["gnomad_genomes_HOM"] or 0.0) if "gnomad_genomes_HOM" in hit else 0.0,
                    'gnomad_genomes_AF': float(hit["gnomad_genomes_AF"] or 0.0) if "gnomad_genomes_AF" in hit else 0.0,
                    'gnomad_genomes_popmax_AF': float(hit["gnomad_genomes_AF_POPMAX"] or 0.0) if "gnomad_genomes_AF_POPMAX" in hit else 0.0,
                },
                'db_gene_ids': list(hit["geneIds"]),
                'db_tags': str(hit["transcriptConsequenceTerms"]) if "transcriptConsequenceTerms" in hit else None,
                'extras': {
                    'grch37_coords': lifted_over_coord,
                    'grch38_coords': "%s-%s-%s-%s" % (hit["contig"], hit["start"], hit["ref"], hit["alt"]),
                    u'alt_allele_pos': 0,
                    u'orig_alt_alleles': map(str, [a.split("-")[-1] for a in hit["originalAltAlleles"]]) if "originalAltAlleles" in hit else []},
                'gene_ids': None,
                'genotypes': genotypes,
                'pos': long(hit['start']),
{'fraction_non_alphanumeric': 0.12898882773932965, 'fraction_numerical': 0.02527340151640409, 'mean_word_length': 2.305332188305098, 'pattern_counts': {'":': 115, '<': 9, '<?xml version=': 0, '>': 23, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 9, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5379048', 'n_tokens_mistral': 14847, 'n_tokens_neox': 13557, 'n_words': 2621}