#!/usr/bin/env python
import sqlite3
import os
import numpy as np
import cPickle
import zlib
import collections
from collections import Counter

import gemini_utils as util
from gemini_constants import *
import GeminiQuery


def get_tstv(c, args):
    """
    Report the transition / transversion ratio.
    """
    ts_cmd = "SELECT count(1) \
           FROM  variants \
           WHERE type = \'snp\' \
           AND   sub_type = \'ts\'"
    tv_cmd = "SELECT count(1) \
          FROM  variants v \
          WHERE type = \'snp\' \
          AND   sub_type = \'tv\'"
    # get the number of transitions
    c.execute(ts_cmd)
    ts = c.fetchone()[0]
    # get the number of transversions
    c.execute(tv_cmd)
    tv = c.fetchone()[0]
    # report the transitions, transversions, and the ts/tv ratio
    print "ts" + '\t' + \
          "tv" + '\t' + "ts/tv"
    print str(ts) + '\t' + \
        str(tv) + '\t' + \
        str(tstv(ts,tv))


def get_tstv_coding(c, args):
    """
    Report the transition / transversion ratio in coding regions.
    """
    ts_cmd = "SELECT count(1) \
           FROM variants v \
           WHERE v.type = \'snp\' \
           AND v.sub_type = \'ts\' \
           AND v.is_coding = 1"
    tv_cmd = "SELECT count(1) \
          FROM variants v \
          WHERE v.type = \'snp\' \
          AND v.sub_type = \'tv\' \
          AND v.is_coding = 1"
    # get the number of transitions
    c.execute(ts_cmd)
    ts = c.fetchone()[0]

    # get the number of transversions
    c.execute(tv_cmd)
    tv = c.fetchone()[0]

    # report the transitions, transversions, and the ts/tv ratio
    print "ts" + '\t' + \
          "tv" + '\t' + "ts/tv"
    print str(ts) + '\t' + \
        str(tv) + '\t' + \
        str(tstv(ts,tv))


def get_tstv_noncoding(c, args):
    """
    Report the transition / transversion ratio in non-coding regions.
    """
    ts_cmd = "SELECT count(1) \
           FROM variants v \
           WHERE v.type = \'snp\' \
           AND v.sub_type = \'ts\' \
           AND v.is_coding = 0"
    tv_cmd = "SELECT count(1) \
          FROM variants v \
          WHERE v.type = \'snp\' \
          AND v.sub_type = \'tv\' \
          AND v.is_coding = 0"
    # get the number of transitions
    c.execute(ts_cmd)
    ts = c.fetchone()[0]

    # get the number of transversions
    c.execute(tv_cmd)
    tv = c.fetchone()[0]

    # report the transitions, transversions, and the ts/tv ratio
    print "ts" + '\t' + \
          "tv" + '\t' + "ts/tv"
    print str(ts) + '\t' + \
        str(tv) + '\t' + \
        str(tstv(ts,tv))


def tstv(ts, tv):
    """
    Calculate ts/tv, and avoid division by zero error
    """
    try:
        return round(float(ts) / float(tv), 4)
    except ZeroDivisionError:
        return 0

def get_snpcounts(c, args):
    """
    Report the count of each type of SNP.
    """
    query = "SELECT ref, alt, count(1) \
             FROM   variants \
             WHERE  type = \'snp\' \
             GROUP BY ref, alt"

    # get the ref and alt alleles for all snps.
    c.execute(query)
    print '\t'.join(['type', 'count'])
    for row in c:
        print '\t'.join([str(row['ref']) + "->" + str(row['alt']),
                         str(row['count(1)'])])


def get_sfs(c, args):
    """
    Report the site frequency spectrum
    """
    precision = 3
    query = "SELECT round(aaf," + str(precision) + "), count(1) \
             FROM variants \
             GROUP BY round(aaf," + str(precision) + ")"

    c.execute(query)
    print '\t'.join(['aaf', 'count'])
    for row in c:
        print '\t'.join([str(row[0]), str(row[1])])


def get_mds(c, args):
    """
    Compute the pairwise genetic distance between each sample.
    """
    idx_to_sample = {}
    c.execute("select sample_id, name from samples")
    for row in c:
        idx_to_sample[int(row['sample_id']) - 1] = row['name']

    query = "SELECT DISTINCT v.variant_id, v.gt_types\
    FROM variants v\
    WHERE v.type = 'snp'"
    c.execute(query)

    # keep a list of numeric genotype values
    # for each sample
    genotypes = collections.defaultdict(list)
    for row in c:

        gt_types = np.array(cPickle.loads(zlib.decompress(row['gt_types'])))

        # at this point, gt_types is a numpy array
        # idx:  0 1 2 3 4 5 6 .. #samples
        # type [0 1 2 1 2 0 0 ..         ]
        for idx, gt_type in enumerate(gt_types):
            sample = idx_to_sample[idx]
            genotypes[sample].append(gt_type)

    mds = collections.defaultdict(float)
    # convert the genotype list for each sample
    # to a numpy array for performance.
    # masks stores an array of T/F indicating which genotypes are
    # known (True, [0,1,2]) and unknown (False [-1]).
    masks = {}
    for s in genotypes:
        sample = str(s)
        x = np.array(genotypes[sample])
        genotypes[sample] = x
        masks[sample] = \
            np.ma.masked_where(genotypes[sample] != UNKNOWN,
                               genotypes[sample]).mask

    # compute the euclidean distance for each s1/s2 combination
    # using numpy's vectorized sum() and square() operations.
    # we use the mask arrays to identify the indices of known genotypes
    # for each sample.  by doing a bitwise AND of the mask arrays for the
    # two samples, we have a mask array of variants where __both__ samples
    # were called.
    for sample1 in genotypes:
        for sample2 in genotypes:
            pair = (sample1, sample2)
            # which variants have known genotypes for both samples?
            both_mask = masks[str(sample1)] & masks[str(sample2)]
            genotype1 = genotypes[sample1]
            genotype2 = genotypes[sample2]

            # distance between s1 and s2:
            eucl_dist = float(np.sum(np.square((genotype1 - genotype2)[both_mask]))) \
                / \
                float(np.sum(both_mask))

            mds[pair] = eucl_dist

    # report the pairwise MDS for each sample pair.
    print "sample1\tsample2\tdistance"
    for pair in mds:
        print "\t".join([str(pair[0]), str(pair[1]), str(round(mds[pair], 4))])


def get_variants_by_sample(c, args):
    """
    Report the number of variants observed for each sample
    where the sample had a non-ref genotype
    """
    idx_to_sample = util.map_indices_to_samples(c)

    # report.
    print '\t'.join(['sample', 'total'])

    query = "SELECT sample_id, \
             (num_het + num_hom_alt) as total \
             FROM sample_genotype_counts"
    c.execute(query)
    for row in c:
        sample = idx_to_sample[row['sample_id']]
        print "\t".join(str(s) for s in [sample,
                                         row['total']])


def get_gtcounts_by_sample(c, args):
    """
    Report the count of each genotype class
    observed for each sample.
    """
    idx_to_sample = util.map_indices_to_samples(c)

    # report.
    print '\t'.join(['sample', 'num_hom_ref', 'num_het',
                     'num_hom_alt', 'num_unknown', 'total'])

    query = "SELECT *, \
             (num_hom_ref + num_het + num_hom_alt + num_unknown) as total \
             FROM sample_genotype_counts"
    c.execute(query)
    # count the number of each genotype type obs. for each sample.
    for row in c:
        sample = idx_to_sample[row['sample_id']]
        print "\t".join(str(s) for s in [sample,
                                         row['num_hom_ref'],
                                         row['num_het'],
                                         row['num_hom_alt'],
                                         row['num_unknown'],
                                         row['total']])


def summarize_query_by_sample(args):
    gq = GeminiQuery.GeminiQuery(args.db)
    gq.run(args.query, show_variant_samples=True, gt_filter=args.gt_filter)
    total_counts = Counter()
    het_counts = Counter()
    hom_alt_counts = Counter()
    hom_ref_counts = Counter()
    print "\t".join(["sample", "total", "num_het", "num_hom_alt", "num_hom_ref"])
    for row in gq:
        total_counts.update(row["variant_samples"])
        het_counts.update(row["HET_samples"])
        hom_alt_counts.update(row["HOM_ALT_samples"])
        hom_ref_counts.update(row["HOM_REF_samples"])
    for key in total_counts.keys():
        count_row = [key, total_counts.get(key, 0), het_counts.get(key, 0),
                     hom_alt_counts.get(key, 0), hom_ref_counts.get(key, 0)]
        print "\t".join(map(str, count_row))


def stats(parser, args):

    if os.path.exists(args.db):
        conn = sqlite3.connect(args.db)
        conn.isolation_level = None
        conn.row_factory = sqlite3.Row
        c = conn.cursor()

        if args.tstv:
            get_tstv(c, args)
        elif args.tstv_coding:
            get_tstv_coding(c, args)
        elif args.tstv_noncoding:
            get_tstv_noncoding(c, args)
        elif args.snp_counts:
            get_snpcounts(c, args)
        elif args.sfs:
            get_sfs(c, args)
        elif args.variants_by_sample:
            get_variants_by_sample(c, args)
        elif args.genotypes_by_sample:
            get_gtcounts_by_sample(c, args)
        elif args.mds:
            get_mds(c, args)
        elif args.query:
            summarize_query_by_sample(args)
