class VCFReader(object):
    def __init__(self, fhand, compressed=None, filename=None,
                 min_calls_for_pop_stats=DEF_MIN_CALLS_FOR_POP_STATS):
        self.fhand = fhand
        self.pyvcf_reader = pyvcfReader(fsock=fhand, compressed=compressed,
                                        filename=filename)
        self.min_calls_for_pop_stats = min_calls_for_pop_stats
        self._snpcaller = None
        self._samples = self.pyvcf_reader.samples
        self._header_lines = self.pyvcf_reader._header_lines
        self._column_headers = self.pyvcf_reader._column_headers

    def parse_snvs(self):
        min_calls_for_pop_stats = self.min_calls_for_pop_stats
        last_snp = None
        try:
            counter =0
            for snp in self.pyvcf_reader:
                counter +=1
                snp = SNV(snp, reader=self,
                          min_calls_for_pop_stats=min_calls_for_pop_stats)
                last_snp = snp
                yield snp
        except Exception:
            from traceback import print_exception
            exc_type, exc_value, exc_traceback = sys.exc_info()

            print_exception(exc_type, exc_value, exc_traceback,
                            limit=20, file=sys.stderr)

            if last_snp is not None:
                chrom = str(last_snp.chrom)
                pos = last_snp.pos
                msg = 'Last parsed SNP was: {} {}\n'.format(chrom, pos + 1)
                sys.stderr.write(msg)
                raise

    def fetch_snvs(self, chrom, start, end=None):
        min_calls_for_pop_stats = self.min_calls_for_pop_stats
        try:
            snvs = self.pyvcf_reader.fetch(chrom, start + 1, end=end)
        except KeyError:
            snvs = []
        if snvs is None:
            snvs = []

        for snp in snvs:
            snp = SNV(snp, reader=self,
                      min_calls_for_pop_stats=min_calls_for_pop_stats)
            yield snp

    def sliding_windows(self, size, step=None, ref_fhand=None,
                        min_num_snps=DEF_MIN_NUM_SNPS_IN_WIN):
        random_snp_reader = VCFReader(open(self.fhand.name))
        sliding_window = _SNPSlidingWindow(snp_reader=random_snp_reader,
                                           win_size=size, win_step=step,
                                           min_num_snps=min_num_snps,
                                           ref_fhand=ref_fhand)
        for window in sliding_window.windows():
            yield window

    @property
    def snpcaller(self):
        if self._snpcaller is not None:
            return self._snpcaller

        metadata = self.pyvcf_reader.metadata
        if 'source' in metadata:
            if 'VarScan2' in metadata['source']:
                snpcaller = VARSCAN
            elif 'freebayes' in metadata['source'][0].lower():
                snpcaller = FREEBAYES
            else:
                snpcaller = GENERIC
        elif 'UnifiedGenotyper' in metadata:
            snpcaller = GATK
        else:
            snpcaller = GENERIC
        self._snpcaller = snpcaller
        return snpcaller

    @property
    def samples(self):
        return self._samples

    @property
    def filters(self):
        return self.pyvcf_reader.filters

    @property
    def infos(self):
        return self.pyvcf_reader.infos

    def _build_header(self, samples=None):
        if samples is None:
            samples = self.samples
        header = '\n'.join(self._header_lines)
        header += '\n#' + '\t'.join(self._column_headers)
        header += '\t' + '\t'.join(samples)
        return header

    @property
    def header(self):
        return self._build_header()

    def create_template_reader(self, samples=None):
        vcf = self._build_header(samples)
        fake_vcf = BytesIO(vcf)
        return VCFReader(fake_vcf) PASSED = 'passed' FILTERED_OUT = 'filtered_out' def group_in_filter_packets(items, items_per_packet):
    for packet in group_in_packets(items, items_per_packet):
        yield {PASSED: packet, FILTERED_OUT: []} class CallRateFilter(_BaseFilter):
    'Filter by the min. number of genotypes called'

    def __init__(self, min_calls=None, min_call_rate=None, reverse=False,
                 samples_to_consider=None):
        parent_kwargs = {'reverse': reverse,
                         'samples_to_consider': samples_to_consider}
        super(CallRateFilter, self).__init__(**parent_kwargs)
        if min_calls is not None and min_call_rate is not None:
            msg = 'Both min_calls and min_call rate cannot be given'
            msg += 'at the same time'
            raise ValueError(msg)
        elif min_calls is None and min_call_rate is None:
            msg = 'min_calls or call_rate should be given'
            raise ValueError(msg)
        self.min_calls = min_calls
        self.min_call_rate = min_call_rate
        self.call_rates = array('f')

    def _do_check(self, snv):
        value = snv.num_called if self.min_calls else snv.call_rate
        self.call_rates.append(value)
        if self.min_calls:
            if value >= self.min_calls:
                return True
            else:
                return False
        else:
            if value >= self.min_call_rate:
                return True
            else:
                return False

    def plot_hist(self, fhand):
        values = self.call_rates
        min_value = self.min_calls if self.min_calls else self.min_call_rate
        xlabel = 'Num. calls' if self.min_calls else 'Call rate'
        ylabel = 'num. SNPs'
        self._plot_hist(fhand, values, min_value=min_value, xlabel=xlabel,
                        ylabel=ylabel) class BiallelicFilter(_BaseFilter):
    'Filter the biallelic SNPs'

    def __init__(self, reverse=False, samples_to_consider=None):
        parent_kwargs = {'reverse': reverse,
                         'samples_to_consider': samples_to_consider}
        super(BiallelicFilter, self).__init__(**parent_kwargs)

    def _do_check(self, snv):
        if len(snv.alleles) == 2:
            return True
        else:
            return False class IsSNPFilter(_BaseFilter):
    def _do_check(self, snv):
        return snv.is_snp class SnvQualFilter(_BaseFilter):
    def __init__(self, min_qual, reverse=False):
        parent_kwargs = {'reverse': reverse}
        super(SnvQualFilter, self).__init__(**parent_kwargs)
        self.min_qual = min_qual
        self.values = array('f')

    def _do_check(self, snv):
        qual = snv.qual
        if qual is None:
            return False
        self.values.append(qual)
        return qual >= self.min_qual

    def plot_hist(self, fhand):
        values = self.values
        min_value = self.min_qual
        xlabel = 'SNV quality'
        ylabel = 'num. SNPs'
        self._plot_hist(fhand, values, min_value=min_value, xlabel=xlabel,
                        ylabel=ylabel) class ObsHetFilter(_BaseFilter):
    def __init__(self, min_het=None, max_het=None, remove_nd=True,
                 reverse=False, samples_to_consider=None):
        parent_kwargs = {'reverse': reverse,
                         'samples_to_consider': samples_to_consider}
        super(ObsHetFilter, self).__init__(**parent_kwargs)
        self.min_het = min_het
        self.max_het = max_het
        self.remove_nd = remove_nd
        self.hets = array('f')

    def _do_check(self, snv):
        min_het = self.min_het
        max_het = self.max_het
        het = snv.obs_het
        if het is None and self.remove_nd:
            return False
        self.hets.append(het)
        if min_het is not None and het < min_het:
            return False
        if max_het is not None and het > max_het:
            return False
        return True

    def plot_hist(self, fhand):
        values = self.hets
        min_value = self.min_het
        max_value = self.max_het
        xlabel = '% Obs. het.'
        ylabel = 'num. SNPs'
        self._plot_hist(fhand, values, min_value=min_value,
                        max_value=max_value, xlabel=xlabel, ylabel=ylabel) class MafFilter(_BaseFilter):
    def __init__(self, min_maf=None, max_maf=None, remove_nd=True,
                 reverse=False, samples_to_consider=None):
        parent_kwargs = {'reverse': reverse,
                         'samples_to_consider': samples_to_consider}
        super(MafFilter, self).__init__(**parent_kwargs)
        if min_maf is None and max_maf is None:
            msg = 'At least one value should be given for the min or max het'
            raise ValueError(msg)
        self.min_maf = min_maf
        self.max_maf = max_maf
        self.remove_nd = remove_nd
        self.values = array('f')

    def _do_check(self, snv):
        min_maf = self.min_maf
        max_maf = self.max_maf
        maf = snv.maf
        if maf is None and self.remove_nd:
            return False
        self.values.append(maf)
        if min_maf is not None and maf < min_maf:
            return False
        if max_maf is not None and maf > max_maf:
            return False
        return True

    def plot_hist(self, fhand):
        values = self.values
        min_value = self.min_maf
        max_value = self.max_maf
        xlabel = 'MAF'
        ylabel = 'num. SNPs'
        self._plot_hist(fhand, values, min_value=min_value,
                        max_value=max_value, xlabel=xlabel, ylabel=ylabel) def filter_snvs(in_fhand, out_fhand, filters, filtered_fhand=None,
                log_fhand=None, reader_kwargs=None):
    '''It filters an input vcf.

    The input fhand has to be uncompressed. The original file could be a
    gzipped file, but in that case it has to be opened with gzip.open before
    sending it to this function.
    '''
    if reader_kwargs is None:
        reader_kwargs = {}
    # The input fhand to this function cannot be compressed
    reader_kwargs.update({'compressed': False,
                         'filename': 'pyvcf_bug_workaround'})

    reader = VCFReader(in_fhand, **reader_kwargs)

    template_reader = VCFReader(StringIO(reader.header))
    writer = VCFWriter(out_fhand, template_reader=template_reader)
    if filtered_fhand:
        filtered_writer = VCFWriter(filtered_fhand,
                                    template_reader=template_reader)
    else:
        filtered_writer = None

    packets = group_in_filter_packets(reader.parse_snvs(),
                                      SNPS_PER_FILTER_PACKET)
    tot_snps = 00.01
    passed_snps = OrderedDict()
    broken_pipe = False
    for packet in packets:
        tot_snps += len(packet[PASSED]) + len(packet[FILTERED_OUT])
        for filter_ in filters:
            packet = filter_(packet)
            filter_name = filter_.__class__.__name__
            if filter_name not in passed_snps:
                passed_snps[filter_name] = 0
            passed_snps[filter_name] += len(packet[PASSED])

        for snv in packet[PASSED]:
            if not _safe_write_snv(writer, snv):
                broken_pipe = True
                break
        if filtered_writer:
            for snv in packet[FILTERED_OUT]:
                if not _safe_write_snv(filtered_writer, snv):
                    broken_pipe = True
                    break
        if broken_pipe:
            break

    if log_fhand:
        _write_log(log_fhand, tot_snps, passed_snps)

    writer.flush() class MonomorphicFilter(_BaseFilter):
    "Filters monomorphic snvs"
    def __init__(self, reverse=False, samples_to_consider=None,
                 freq_threshold=1):
        parent_kwargs = {'reverse': reverse,
                         'samples_to_consider': samples_to_consider}
        super(MonomorphicFilter, self).__init__(**parent_kwargs)
        self._freq_threslhold = freq_threshold

    def _do_check(self, snv):
        return snv.is_polymorphic(self._freq_threslhold) class WeirdSegregationFilter(object):
    def __init__(self, alpha=DEF_ALPHA, num_snvs_check=DEF_SNV_WIN,
                 max_failed_freq=DEF_MAX_FAILED_FREQ, samples=None,
                 win_width=DEF_MAX_DIST, win_mask_width=DEF_MIN_DIST,
                 min_num_snvs_check_in_win=DEF_MIN_NUM_CHECK_SNPS_IN_WIN,
                 debug_plot_dir=None):
        # We're assuming that most snps are ok and that a few have a weird
        # segregation
        self.alpha = alpha
        self.num_snvs_check = num_snvs_check
        self.max_failed_freq = max_failed_freq
        self.win_width = win_width
        if win_mask_width < 1:
            msg = 'You should mask at least a window of 3 bp to avoid the '
            msg += 'comparison with itself'
            raise ValueError(msg)
        self.win_mask_width = win_mask_width
        self.min_num_snvs_check_in_win = min_num_snvs_check_in_win
        self._failed_freqs = array('f')
        self.tot_snps = 0
        self.passed_snps = 0
        self.samples = samples
        if debug_plot_dir is not None and not exists(debug_plot_dir):
            mkdir(debug_plot_dir)
        self.plot_dir = debug_plot_dir

    def filter_vcf(self, vcf_fpath, min_samples=DEF_MIN_CALLS_FOR_POP_STATS):
        reader = VCFReader(open(vcf_fpath),
                           min_calls_for_pop_stats=min_samples)
        snvs = reader.parse_snvs()
        random_reader = VCFReader(open(vcf_fpath))

        for snv_1 in snvs:
            self.tot_snps += 1
            loc = snv_1.pos

            if self.plot_dir:
                chrom = str(snv_1.chrom)
                fname = chrom + '_' + str(loc) + '.png'
                chrom_dir = pjoin(self.plot_dir, chrom)
                if not exists(chrom_dir):
                    mkdir(chrom_dir)
                plot_fhand = open(pjoin(chrom_dir, fname), 'w')
                debug_plot_info = []
            else:
                plot_fhand = None

            win_1_start = loc - (self.win_width / 2)
            if win_1_start < 0:
                win_1_start = 0
            win_1_end = loc - (self.win_mask_width / 2)
            if win_1_end < 0:
                win_1_end = 0
            if win_1_end != 0:
                snvs_win_1 = random_reader.fetch_snvs(snv_1.chrom,
                                                      start=int(win_1_start),
                                                      end=int(win_1_end))
            else:
                snvs_win_1 = []

            win_2_start = loc + (self.win_mask_width / 2)
            win_2_end = loc + (self.win_width / 2)
            snvs_win_2 = random_reader.fetch_snvs(snv_1.chrom,
                                                  start=win_2_start,
                                                  end=win_2_end)
            snvs_in_win = list(chain(snvs_win_1, snvs_win_2))
            if len(snvs_in_win) > self.num_snvs_check:
                snvs_in_win = random.sample(snvs_in_win, self.num_snvs_check)
            if len(snvs_in_win) < self.min_num_snvs_check_in_win:
                # Not enough snps to check
                continue

            orig_snp = snv_1
            if self.samples is not None:
                snv_1 = snv_1.filter_calls_by_sample(self.samples)

            exp_cnts = snv_1.biallelic_genotype_counts

            if exp_cnts is None:
                continue

            test_values = []
            for snv_2 in snvs_in_win:
                if self.samples is not None:
                    snv_2 = snv_2.filter_calls_by_sample(self.samples)
                obs_cnts = snv_2.biallelic_genotype_counts
                if obs_cnts is None:
                    continue
                test_values.append(_fisher_extact_rxc(obs_cnts, exp_cnts))

                if plot_fhand:
                    debug_plot_info.append({'pos': snv_2.pos,
                                            'AA': obs_cnts[0],
                                            'Aa': obs_cnts[1],
                                            'aa': obs_cnts[2],
                                            'close_snp': True})
            alpha2 = self.alpha/len(test_values)
            results = []
            for idx, val in enumerate(test_values):
                result = False if val is None else val > alpha2
                results.append(result)

                if plot_fhand:
                    debug_plot_info[idx]['result'] = result

            if len(test_values) < self.min_num_snvs_check_in_win:
                # few snps can be tested for segregation
                continue

            tot_checked = len(test_values)
            if tot_checked > 0:
                failed_freq = results.count(False) / tot_checked
                passed = self.max_failed_freq > failed_freq
            else:
                failed_freq = None
                passed = False
            if failed_freq is not None:
                self._failed_freqs.append(failed_freq)

            if plot_fhand:
                debug_plot_info.append({'pos': snv_1.pos,
                                        'AA': exp_cnts[0],
                                        'Aa': exp_cnts[1],
                                        'aa': exp_cnts[2],
                                        'result': passed,
                                        'close_snp': False})
                self._plot_segregation_debug(debug_plot_info, plot_fhand)
            if passed:
                self.passed_snps += 1
                yield orig_snp

    @staticmethod
    def _plot_segregation_debug(plot_info, fhand):
        greens = ['#2d6e12', '#76b75b', '#164900']
        reds = ['#8f0007', '#fb1f2a', '#b90009']
        grays = ['0.5', '0.25', '0.75']
        fig = Figure(figsize=(10, 4))
        axes1 = fig.add_subplot(211)
        axes2 = fig.add_subplot(212)
        plot_info.sort(key=itemgetter('pos'))

        is_close_snp = [geno_info['close_snp'] for geno_info in plot_info]
        tested_snp_idx = is_close_snp.index(False)
        passed = [geno_info['result'] for geno_info in plot_info]
        total_cnts = [geno_info['AA'] + geno_info['Aa'] + geno_info['aa'] for geno_info in plot_info]

        num_snvs = len(passed)
        bottoms = [0] * num_snvs
        freq_bottoms = [0] * num_snvs
        lefts = range(num_snvs)
        for idx, geno in enumerate(('AA', 'Aa', 'aa')):
            cnts = [geno_info[geno] for geno_info in plot_info]
            backcolors = [greens[idx] if pss else reds[idx] for pss in passed]
            edgecolors = [grays[idx] for pss in passed]
            backcolors[tested_snp_idx], edgecolors[tested_snp_idx] = \
                edgecolors[tested_snp_idx], backcolors[tested_snp_idx]
            heights = cnts
            axes1.bar(lefts, heights, bottom=bottoms, color=backcolors,
                      edgecolor=edgecolors)

            freq_heights = [height/total_cnt for height, total_cnt in zip(heights, total_cnts)]
            axes2.bar(lefts, freq_heights, bottom=freq_bottoms,
                      color=backcolors, edgecolor=edgecolors)

            bottoms = [bot + heig for bot, heig in zip(bottoms, heights)]
            freq_bottoms = [bot + heig for bot, heig in zip(freq_bottoms, freq_heights)]

        axes1.tick_params(axis='y', which='both', left='off', right='off')
        axes2.tick_params(axis='y', which='both', left='off', right='off')

        canvas = FigureCanvas(fig)
        canvas.print_figure(fhand)
        fhand.flush()

    def plot_failed_freq_dist(self, fhand):
        fig = Figure()
        axes = fig.add_subplot(111)
        axes.hist(self._failed_freqs, fill=True, log=True, bins=20,
                  rwidth=1)
        axes.axvline(x=self.max_failed_freq)
        axes.set_xlabel('% of adjacent SNPs segregating differently')
        axes.set_ylabel('num. SNPs')
        _print_figure(axes, fig, fhand, plot_legend=False)

    def write_log(self, fhand):
        _write_log(fhand, self.tot_snps,
                   {self.__class__.__name__: self.passed_snps}) class WeirdRecombFilter(object):
    def __init__(self, pop_type, max_zero_dist_recomb=DEF_MAX_ZERO_DIST_RECOMB,
                 alpha_recomb_0=DEF_ALPHA,
                 snps_in_window=DEF_NUM_SNPS_IN_WIN_FOR_WEIRD_RECOMB,
                 min_num_snps=DEF_MIN_NUM_SNPS_WEIRD_RECOMB,
                 max_recomb_curve_fit=DEF_MAX_RECOMB_RATE_WEIRD_RECOMB,
                 debug_plot_dir=None, samples=None):
        self.pop_type = pop_type
        self.max_zero_dist_recomb = max_zero_dist_recomb
        self.alpha_recomb_0 = alpha_recomb_0
        self.snps_in_window = snps_in_window
        self.min_num_snps = min_num_snps
        self.max_recomb_curve_fit = max_recomb_curve_fit
        self.debug_plot_dir = debug_plot_dir
        self.samples = samples
        self.tot_snps = 0
        self.passed_snps = 0
        self.not_fitted_counter = Counter()
        self.recomb_rates = {'ok': array('f'),
                             'ok_conf_is_None': array('f'),
                             'not_ok': array('f')}

    def filter_snvs(self, snvs):

        snps = RandomAccessIterator(snvs, rnd_access_win=self.snps_in_window)
        rates = _calculate_segregation_rates(snps, self.pop_type,
                                             self.snps_in_window,
                                             samples=self.samples)
        max_zero_dist = self.max_zero_dist_recomb
        for snp, chrom, pos, rates in rates:
            self.tot_snps += 1
            dists, recombs = zip(*[(rate.pos - pos, rate.recomb_rate) for rate in rates])
            if len(dists) < self.min_num_snps:
                continue
            if self.debug_plot_dir is None:
                plot_fhand = None
            else:
                chrom_dir = pjoin(self.debug_plot_dir, str(chrom))
                if not exists(chrom_dir):
                    mkdir(chrom_dir)
                fname = str(chrom) + '_' + str(pos) + '.png'
                plot_fhand = open(pjoin(chrom_dir, fname), 'w')
            res = _calc_ajusted_recomb(dists, recombs,
                                       max_recomb=self.max_recomb_curve_fit,
                                       max_zero_dist_recomb=max_zero_dist,
                                       alpha_recomb_0=self.alpha_recomb_0,
                                       plot_fhand=plot_fhand)
            self._store_log_info(*res)
            if res[1]:
                self.passed_snps += 1
                yield snp

    def _store_log_info(self, recomb_at_0, snp_ok, debug_info):
        if 'reason_no_fit' in debug_info:
            self.not_fitted_counter[debug_info['reason_no_fit']] += 1
            return
        index = snp_ok, debug_info['conf_interval'] is None
        if snp_ok:
            if debug_info['conf_interval'] is None:
                index = 'ok_conf_is_None'
            else:
                index = 'ok'
        else:
            index = 'not_ok'
        self.recomb_rates[index].append(recomb_at_0)

    def plot_recomb_at_0_dist_hist(self, fhand):
        fig = Figure()
        axes = fig.add_subplot(111)
        data = [self.recomb_rates['ok'], self.recomb_rates['ok_conf_is_None'],
                self.recomb_rates['not_ok']]
        labels = ['OK', 'OK conf. is none', 'Not OK']
        colors = [(0.3, 1, 0.3), (0.3, 1, 0.6), (1, 0.3, 0.3)]
        some_data = [bool(dataset) for dataset in data]
        labels = [label for draw, label in zip(some_data, labels) if draw]
        colors = [color for draw, color in zip(some_data, colors) if draw]
        data = [dataset for draw, dataset in zip(some_data, data) if draw]
        axes.hist(data, stacked=True, fill=True, log=True, bins=20,
                  label=labels, rwidth=1, color=colors)
        _print_figure(axes, fig, fhand)

    def write_log(self, fhand):
        _write_log(fhand, self.tot_snps,
                   {self.__class__.__name__: self.passed_snps}) VCF_BIN_DIR = os.path.join(BIN_DIR, 'vcf') TEST_DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                             '..', 'test', 'test_data'))
from os.path import join
from tempfile import NamedTemporaryFile
from subprocess import check_output, Popen, PIPE, check_call
from StringIO import StringIO
from crumbs.vcf.snv import VCFReader
from crumbs.vcf.filters import (PASSED, FILTERED_OUT, group_in_filter_packets,
                                CallRateFilter, BiallelicFilter, IsSNPFilter,
                                SnvQualFilter, ObsHetFilter, MafFilter,
                                filter_snvs, MonomorphicFilter,
                                WeirdSegregationFilter,
                                WeirdRecombFilter)
from crumbs.utils.bin_utils import VCF_BIN_DIR
from crumbs.utils.test_utils import TEST_DATA_DIR
import unittest
import os
# Copyright 2012 Jose Blanca, Peio Ziarsolo, COMAV-Univ. Politecnica Valencia
# This file is part of ngs_crumbs.
# ngs_crumbs is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# ngs_crumbs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR  PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with ngs_crumbs. If not, see <http://www.gnu.org/licenses/>.




# Method could be a function
# pylint: disable=R0201
# Too many public methods
# pylint: disable=R0904
# Missing docstring
# pylint: disable=C0111

VCF_PATH = join(TEST_DATA_DIR, 'sample.vcf.gz')
VCF_INDEL_PATH = join(TEST_DATA_DIR, 'sample_indel.vcf.gz')
VARI_VCF_PATH = join(TEST_DATA_DIR, 'vari_filter.vcf')

VCF_HEADER = '''##fileformat=VCFv4.1
##fileDate=20090805
##source=myImputationProgramV3.1
##reference=file:///seq/references/1000GenomesPilot-NCBI36.fasta
##contig=<ID=20,length=62435964,assembly=B36,md5=f126cdf8a6e0c7f379d618ff66beb2da,species="Homo sapiens",taxonomy=x>
##phasing=partial
##INFO=<ID=NS,Number=1,Type=Integer,Description="Number of Samples With Data">
##INFO=<ID=DP,Number=1,Type=Integer,Description="Total Depth">
##INFO=<ID=AF,Number=A,Type=Float,Description="Allele Frequency">
##INFO=<ID=AA,Number=1,Type=String,Description="Ancestral Allele">
##INFO=<ID=DB,Number=0,Type=Flag,Description="dbSNP membership, build 129">
##INFO=<ID=H2,Number=0,Type=Flag,Description="HapMap2 membership">
##FILTER=<ID=q10,Description="Quality below 10">
##FILTER=<ID=s50,Description="Less than 50% of samples have data">
##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">
##FORMAT=<ID=DP,Number=1,Type=Integer,Description="Read Depth">
##FORMAT=<ID=HQ,Number=2,Type=Integer,Description="Haplotype Quality">
'''
VCF = '''#CHROM POS ID REF ALT QUAL FILTER INFO FORMAT NA00001 NA00002 NA00003
20\t14370\trs6054257\tG\tA\t29\tPASS\tNS=3;DP=14;AF=0.5;DB;H2\tGT:GQ:DP:HQ\t0|0:48:1:51,51\t1|0:48:8:51,51\t1/1:43:5:.,.
20\t17330\t.\tT\tA\t3\tq10\tNS=3;DP=11;AF=0.017\tGT:GQ:DP:HQ\t0|0:49:3:58,50\t0|1:3:5:65,3\t0/0:41:3
20\t1110696\trs6040355\tA\tG,T\t67\tPASS\tNS=2;DP=10;AF=0.333,0.667;AA=T;DB\tGT:GQ:DP:HQ\t1|2:21:6:23,27\t2|1:2:0:18,2\t2/2:35:4
20\t1230237\t.\tT\t.\t47\tPASS\tNS=3;DP=13;AA=T\tGT:GQ:DP:HQ\t0|0:54:7:56,60\t0|0:48:4:51,51\t0/0:61:2
20\t1234567\tmicrosat1\tGTC\tG,GTCT\t50\tPASS\tNS=3;DP=9;AA=G\tGT:GQ:DP\t0/1:35:4\t0/2:17:2\t1/1:40:3
20\t1234567\tmicrosat1\tGTC\tG,GTCT\t50\tPASS\tNS=3;DP=9;AA=G\tGT:GQ:DP\t./.:35:4\t0/2:17:2\t1/1:40:3
'''


def filter_vcf(vcf_fhand, filter_):
    snps = VCFReader(vcf_fhand, min_calls_for_pop_stats=1).parse_snvs()
    packet = list(group_in_filter_packets(snps, 10))[0]
    filtered_packet = filter_(packet)
    return filtered_packet


def eval_prop_in_packet(packet, prop):
    eval_prop = lambda snps: [getattr(snp, prop)for snp in snps]
    packet = {PASSED: eval_prop(packet[PASSED]),
              FILTERED_OUT: eval_prop(packet[FILTERED_OUT])}
    return packet


class FiltersTest(unittest.TestCase):

    def test_group_in_filter_packets(self):
        items = list(range(10))
        packets = list(group_in_filter_packets(items, 4))
        res = [{FILTERED_OUT: [], PASSED: (0, 1, 2, 3)},
               {FILTERED_OUT: [], PASSED: (4, 5, 6, 7)},
               {FILTERED_OUT: [], PASSED: (8, 9)}]
        assert res == packets

    def test_biallelic(self):
        packet = filter_vcf(open(VCF_PATH), filter_=BiallelicFilter())
        res = eval_prop_in_packet(packet, 'alleles')
        assert not(res[FILTERED_OUT])

        assert len(res[PASSED]) == 10