from __future__ import print_function, absolute_import
import sys
import os
try:
    from compiler import compile
except ImportError:
    basestring = str
    pass
import operator
import itertools as it
from argparse import ArgumentParser
from .GeminiQuery import GeminiQuery

def add_args(a=None):
    if a is None:
        a = ArgumentParser()
    a.add_argument("--min-filters", type=int, default=1)
    a.add_argument("--where", default="is_exonic = 1 AND impact_severity != 'LOW'",
            help="where clause to subset variants. [default \"%(default)s\"]")
    a.add_argument("--gt-filter", required=False, default=[], action='append')
    a.add_argument("--gt-filter-required", required=False, default=[],
            action='append', help="specify filter(s) that must be met."
            "a variant passing this does filter is required and does not "
            " contribute to '--min-filters")
    a.add_argument("--filter")
    a.add_argument("--columns", default="chrom,start,end,gene,impact,impact_severity,max_aaf_all")
    a.add_argument("db")
    return a


def add_cols(cols, gt_filter):
    assert isinstance(cols, list)
    all_cols = ["gts", "gt_types", "gt_phases", "gt_depths",
                "gt_ref_depths", "gt_alt_depths", "gt_quals",
                "gt_phred_ll_homref", "gt_phred_ll_het",
                "gt_phred_ll_homalt"]
    return [x for x in all_cols if x in gt_filter and not x in cols]


def gen_results(rows, gt_filters, gt_req_filters, min_filters, min_variants, columns,
                user_dict=None):
    """
    gen_results is called on rows from the same group (e.g. gene).
    it does filtering and aggregation.
    """
    # we track the index of the passed filter in passed_filters.
    gene_passed_filters = {}
    if user_dict is None:
        user_dict = {}
    subset = []
    for row in rows:
        cols = {c: row[c] for c in columns}
        cols.update(user_dict)

        row_passed_filters = []
        # check required filters first.
        for i, gt_req in enumerate(gt_req_filters, start=1):
            if not eval(gt_req, cols): break
            row_passed_filters.append("required[%d]" % i)

        if len(row_passed_filters) < len(gt_req_filters): continue

        # only test the other filters if we met the req'd filter.

        # have to test all filters since 1 row can meet multiple fitlers.
        for i, gt_filter in enumerate(gt_filters, start=1):
            if eval(gt_filter, cols):
                # track that this filter passed.
                gene_passed_filters[i] = True
                row_passed_filters.append(i)
        # make sure that some non-required filters passed in order to display
        # the row.
        if row_passed_filters and (len(gt_filters) == 0 or sum(isinstance(f, int) for f in row_passed_filters) > 0):
            row.print_fields['variant_filters'] = ",".join(map(str, row_passed_filters))
            subset.append(row)
    if len(gene_passed_filters) < min_filters or len(subset) < min_variants:
        raise StopIteration

    # e.g. 1,2 indicating which filters passed
    passed_filters = ",".join(str(x) for x in sorted(gene_passed_filters))
    for row in subset:
        row.print_fields['n_gene_variants'] = len(subset)
        row.print_fields['gene_filters'] = passed_filters
        yield row


def genewise(db, gt_filters, gt_req_filters, filter=None, columns=None, min_filters=None,
             min_variants=1,
             where=None,
             grouper="gene"):
    assert os.path.exists(db)

    orig_columns = [x.strip() for x in (columns or "").split(",")]
    added_cols = add_cols(orig_columns, "||".join(gt_filters))
    req_cols = add_cols(added_cols, "||".join(gt_req_filters))
    added_cols = added_cols + req_cols
    if grouper not in orig_columns:
        added_cols.append(grouper)
    columns = orig_columns + added_cols
    assert not any(';' in c for c in columns)

    # NOTE: we could make the WHERE part customizable.
    query = "SELECT {columns} FROM variants WHERE (%s)" % where
    if filter:
        query += " AND  " + filter
    query += (" ORDER BY CHROM, %s" % grouper)

    gq = GeminiQuery(db, include_gt_cols=True)

    # use the GeminiQuery machinery to transform to something that's eval'able
    # then compile once for speed.
    cleaned_filters = []
    for gt_filter in gt_filters:
        assert gq._is_gt_filter_safe(gt_filter)
        gt_filter = gq._correct_genotype_filter(gt_filter)
        cleaned_filters.append(compile(gt_filter, gt_filter, 'eval'))

    cleaned_reqs = []
    for gt_filter in gt_req_filters:
        assert gq._is_gt_filter_safe(gt_filter)
        gt_filter = gq._correct_genotype_filter(gt_filter)
        cleaned_reqs.append(compile(gt_filter, gt_filter, 'eval'))


    if not "gt_types" in columns:
        columns.append("gt_types")
        added_cols.append("gt_types")

    gq.run(query.format(columns=", ".join(columns)), needs_genotypes=True)
    columns = [c for c in columns if not c in gq.gt_name_to_idx_map or (gq.gt_name_to_idx_map.get(c) == c)]

    if isinstance(grouper, basestring):
        grouper = operator.itemgetter(grouper)

    user_dict = dict(sample_info=gq.sample_info, HOM_REF=0, HET=1, UNKNOWN=2,
            HOM_ALT=3, MISSING=None, UNAFFECTED=1, AFFECTED=2)
    header_printed = False
    for groupkey, grp in it.groupby(gq, grouper):
        grp = list(grp)
        for x in gen_results(grp, cleaned_filters, cleaned_reqs, min_filters or 0,
                             min_variants, columns, user_dict=user_dict):
            for c in added_cols:
                if c != grouper:
                    del x.print_fields[c]
            if not header_printed:
                print("\t".join(x.print_fields.keys()))
                header_printed = True
            print(x)


def run(args):
    if args.min_filters > len(args.gt_filter):
        if args.min_filters == 1 and len(args.gt_filter) == 0:
            # only do this if they didn't specify args.gt_filter
            args.min_filters = 0
        else:
            raise RuntimeError("gene-wise: specified --min-filter > the number of --gt-filters\n")
    if len(args.gt_filter) == 0 and len(args.gt_filter_required) == 0:
        raise RuntimeError("ERROR gene-wise: specified neither --min-filter or --min-filter-required\n")

    genewise(args.db, args.gt_filter, args.gt_filter_required, args.filter, args.columns,
             args.min_filters, where=args.where)
