#!/usr/bin/env python
import pysam
import sqlite3
import os
import sys
import collections
import re
from unidecode import unidecode
from bx.bbi.bigwig_file import BigWigFile
from gemini.config import read_gemini_config

# 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.b141.20140813.hg19.tidy.vcf.gz'),
     'clinvar': os.path.join(anno_dirname, 'clinvar_20150305.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.autosomes.phase3_shapeit2_mvncall_integrated_v5.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')
    }
    # 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, 'whole_genome_SNVs.tsv.compressed.gz')):
        annos['cadd_score'] = os.path.join(anno_dirname, 'whole_genome_SNVs.tsv.compressed.gz')
    return annos

class ClinVarInfo(object):
    def __init__(self):
        self.clinvar_dbsource = None
        self.clinvar_dbsource_id = None
        self.clinvar_origin = None
        self.clinvar_sig = None
        self.clinvar_dsdb = None
        self.clinvar_dsdbid = None
        self.clinvar_disease_name = None
        self.clinvar_disease_acc = None
        self.clinvar_in_omim = None
        self.clinvar_in_locus_spec_db = None
        self.clinvar_on_diag_assay = None
        self.clinvar_causal_allele = 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([self.clinvar_dbsource,
                          self.clinvar_dbsource_id,
                          self.clinvar_origin,
                          self.clinvar_sig,
                          self.clinvar_dsdb,
                          self.clinvar_dsdbid,
                          self.clinvar_disease_name,
                          self.clinvar_disease_acc,
                          str(self.clinvar_in_omim),
                          str(self.clinvar_in_locus_spec_db),
                          str(self.clinvar_on_diag_assay),
                          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")

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"):
                annos[anno] = pysam.Tabixfile(anno_files[anno])
            # .bw denotes BigWig files.
            elif anno_files[anno].endswith(".bw"):
                annos[anno] = BigWigFile(open(anno_files[anno]))

        except IOError:
            sys.exit("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
    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 = []
    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, dict) or isinstance(var, sqlite3.Row):
        chrom = var["chrom"]
        start = int(var["start"])
        end = int(var["end"])
    else:
        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, basestring):
        # Assume var is a line from a VCF.
        ref, alt = var.split('\t')[3:5]
    elif isinstance(var, dict) or isinstance(var, sqlite3.Row):
        ref = var["ref"]
        alt = var["alt"]
    else:
        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(",")
    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'