<commit_before>"""
Tests for variant_filter.py.
"""

import os
import pickle

from django.contrib.auth.models import User
from django.test import TestCase
from sympy.core.function import sympify

from main.models import Dataset
from main.models import ExperimentSample
from main.models import Project
from main.models import ReferenceGenome
from main.models import Variant
from main.models import VariantAlternate
from main.models import VariantCallerCommonData
from main.models import VariantEvidence
from main.models import VariantSet
from main.models import VariantToVariantSet
from scripts.dynamic_snp_filter_key_map import initialize_filter_key_map
from scripts.dynamic_snp_filter_key_map import update_filter_key_map
from settings import PWD as GD_ROOT
from variants.variant_filter import EXPRESSION_REGEX
from variants.variant_filter import SAMPLE_SCOPE_REGEX
from variants.variant_filter import SAMPLE_SCOPE_REGEX_NAMED
from variants.variant_filter import SET_REGEX
from variants.variant_filter import SET_REGEX_NAMED
from variants.variant_filter import get_variants_that_pass_filter
from variants.variant_filter import SymbolGenerator
from variants.variant_filter import ParseError
from variants.variant_filter import VariantFilterEvaluator

TEST_DIR = os.path.join(GD_ROOT, 'test_data', 'genbank_aligned')

TEST_ANNOTATED_VCF = os.path.join(TEST_DIR, 'bwa_align_annotated.vcf')


class BaseTestVariantFilterTestCase(TestCase):
    """Base class for the test classes here.
    """

    def setUp(self):
        user = User.objects.create_user('testuser', password='password',
                email='<EMAIL>')
        self.project = Project.objects.create(owner=user.get_profile(),
                title='Test Project')
        self.ref_genome = ReferenceGenome.objects.create(project=self.project,
                label='refgenome', num_chromosomes=1, num_bases=1000)

        # Make sure the reference genome has the required vcf keys.
        initialize_filter_key_map(self.ref_genome)
        update_filter_key_map(self.ref_genome, TEST_ANNOTATED_VCF)

        self.vcf_dataset = Dataset.objects.create(
                label='test_data_set',
                type=Dataset.TYPE.VCF_FREEBAYES,
                filesystem_location=TEST_ANNOTATED_VCF)


class TestVariantFilter(BaseTestVariantFilterTestCase):

    def test_filter__by_position(self):
        """Test filtering by position.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        # Test querying Variants with position > 5.
        result = get_variants_that_pass_filter('position > 5', self.ref_genome)
        variants_above_5 = result.variant_set
        self.assertEqual(4, len(variants_above_5))
        for var in variants_above_5:
            self.assertTrue(var.position > 5)

        # Test querying Variants with position >= 5.
        result = get_variants_that_pass_filter('position >= 5', self.ref_genome)
        variants_above_5 = result.variant_set
        self.assertEqual(5, len(variants_above_5))
        for var in variants_above_5:
            self.assertTrue(var.position >= 5)


    def test_filter__by_chromosome(self):
        """Test filtering by position.
        """
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        CHROM_2 = 'chrom2'
        for pos in range(9):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_2,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        result = get_variants_that_pass_filter('chromosome = chrom',
                self.ref_genome)
        self.assertEqual(6, len(result.variant_set))

        result = get_variants_that_pass_filter('chromosome = chrom2',
                self.ref_genome)
        self.assertEqual(9, len(result.variant_set))


    def test_filter__by_position_and_chromosome(self):
        """Test filtering by position.
        """
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        CHROM_2 = 'chrom2'
        for pos in range(9):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_2,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        QUERY_STRING = 'position > 4 & chromosome = chrom'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        self.assertEqual(1, len(result.variant_set))

        QUERY_STRING = 'position >= 5 & chromosome = chrom2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        self.assertEqual(4, len(result.variant_set))


    def test_filter__invalid_key(self):
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))            

        QUERY_STRING = 'dinosaur > 4 & chromosome = chrom'
        with self.assertRaises(ParseError):
            get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)


    def test_filter__by_position_complex(self):
        """Test filtering by position.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        # Test AND case.
        QUERY_STRING = 'position < 1 & position > 7'
        result = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome)
        self.assertEqual(0, len(result.variant_set))

        # Test OR case.
        QUERY_STRING = 'position < 1 | position > 7'
        result = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome)
        variant_list = result.variant_set
        self.assertEqual(3, len(variant_list))
        for pos in [0, 8, 9]:
            found = False
            for var in variant_list:
                if var.position == pos:
                    found = True
                    break
            self.assertTrue(found, "Expected variant at pos %d not found" % pos)


    def test_filter__with_non_SQL_key(self):
        """Tests handling of non-SQL key.

        NOTE: This is not actually testing filtering, just showing that we
        can handle the key.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        QUERY_STRING = 'position < 1 & INFO_XRM > 0'
        variants = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome).variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position < 1 | position >= 7 & INFO_XRM > 0'
        variants = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))


    def test_filter__common_data(self):
        """Test with VariantCallerCommonData filtering.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        raw_data_dict = {
                'INFO_XRM': pickle.dumps(0.12),
        }
        VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_data_dict
        )

        QUERY_STRING = 'position < 5 & INFO_XRM > 0'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & INFO_XRM > 1'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))


    def test_filter__variant_evidence(self):
        """Test with VariantEvidence objects.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        raw_common_data_dict = {
                'INFO_XRM': pickle.dumps(0.12),
        }
        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_common_data_dict
        )

        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(0),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(None),
                'gt_bases': pickle.dumps(None)

        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        QUERY_STRING = 'position < 5 & gt_type = 2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & gt_type = 0'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & gt_type = 1'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position > 5 & gt_type = 2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position < 5 & (gt_type = 0 | gt_type = 2)'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        # Check that the number of passing samples is correct.
        passing_sample_ids = result.variant_id_to_metadata_dict[variant.id][
                'passing_sample_ids']
        self.assertEqual(2, len(passing_sample_ids))


    def test_filter__equality(self):
        """Test filtering with equality operators.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        variants = get_variants_that_pass_filter('position == 5',
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))
        self.assertEqual(5, variants.pop().position)

        variants = get_variants_that_pass_filter('position = 5',
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))
        self.assertEqual(5, variants.pop().position)

        variants = get_variants_that_pass_filter('position != 5',
                self.ref_genome).variant_set
        self.assertEqual(9, len(variants))
        for var in variants:
            self.assertTrue(var.position != 5)


    def test_filter__scoped(self):
        """Test with VariantEvidence objects.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        # Common data shared by all.
        raw_common_data_dict = {
                'INFO_XRM': pickle.dumps(0.12)
        }
        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_common_data_dict
        )

        # Sample 1
        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        # Sample 2
        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(0),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        # Sample 3
        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(1),
                'gt_bases': pickle.dumps(None)
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        # Sample 4
        sample_obj_4 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 4',
                group='Plate 1',
                well='A04',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_4,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        QUERY_STRING = '(position < 5 & gt_type = 2) in ANY(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ALL(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ALL(%s)' % (
                sample_obj.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ONLY(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        # This should fail because condition passes for sample_obj and
        # sample_obj_4.
        QUERY_STRING = '(position < 5 & gt_type = 2) in ONLY(%s)' % (
                sample_obj.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))


    def test_filter__common_data_per_alt(self):
        """Test filtering for common data of type '-1', 
        with different values for each alternate allele. 
        """

        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        for alt_value, alt_eff in zip(['G','T'], ['NON_SYNONYMOUS_CODING',
                        'SYNONYMOUS_CODING']):

            alt = VariantAlternate.objects.create(
                        variant=variant,
                        alt_value=alt_value)
            alt.data = {'INFO_EFF_EFFECT': alt_eff}

            variant.variantalternate_set.add(alt)

        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=dict()
        )
        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'GT': pickle.dumps('1/1'),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'GT': pickle.dumps('1/2'),
                'gt_bases': pickle.dumps('G/T')
        }
        VariantEvidence.objects.create(
                uid=1000,
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(1),
                'GT': pickle.dumps('0/0'),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        # per_alt_dict, per_alt_types = get_per_alt_dict(
        #         'INFO_EFF_EFFECT', variant,
        #         VariantEvidence.objects.get(uid=1000),
        #         self.ref_genome.variant_key_map['snp_caller_common_data'])

        # assert('INFO_EFF_EFFECT' in per_alt_dict.keys())
        # assert('NON_SYNONYMOUS_CODING' in per_alt_dict['INFO_EFF_EFFECT'])

        QUERY_STRING = (
                '(position < 5 & INFO_EFF_EFFECT = NON_SYNONYMOUS_CODING)')
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set

        self.assertEqual(1, len(variants))  

        passing_sample_ids = result.variant_id_to_metadata_dict[variant.id][
                'passing_sample_ids']

        self.assertEqual(2, len(passing_sample_ids))

    def test_filter__sets(self):
        """Test filtering relative to sets.
        """
        # Some variant sets for testing.
        set_1 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set1')
        set_2 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set2')
        set_3 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set3')

        # Add some variants to each of the above sets.
        def _create_variants_in_set(pos_range, var_set_list):
            for pos in pos_range:
                var = Variant.objects.create(
                    type=Variant.TYPE.TRANSITION,
                    reference_genome=self.ref_genome,
                    chromosome='chrom',
                    position=pos,
                    ref_value='A')

                var.variantalternate_set.add(
                        VariantAlternate.objects.create(
                                variant=var,
                                alt_value='G'))

                for var_set in var_set_list:
                    if var_set is not None:
                        VariantToVariantSet.objects.create(variant=var,
                                variant_set=var_set)
        _create_variants_in_set(range(3), [set_1]) # 3
        _create_variants_in_set(range(3, 6), [set_2]) # 3
        _create_variants_in_set(range(6, 9), [set_3]) # 3
        _create_variants_in_set(range(9, 13), [None]) # 4
        _create_variants_in_set(range(13, 19), [set_1, set_2]) # 6

        # We have 19 total variants.
        self.assertEqual(19, len(Variant.objects.all()))

        QUERY_STRING = 'IN_SET(%s)' % (set_1.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(9, len(variants))

        QUERY_STRING = 'IN_SET(%s) | IN_SET(%s)' % (set_1.uid, set_2.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(12, len(variants))

        QUERY_STRING = 'IN_SET(%s) | IN_SET(%s) | IN_SET(%s)' % (
                set_1.uid, set_2.uid, set_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(15, len(variants))
<commit_msg>Add some clarifying comments to test case to help discuss with DBG.<commit_after>"""
Tests for variant_filter.py.
"""

import os
import pickle

from django.contrib.auth.models import User
from django.test import TestCase
from sympy.core.function import sympify

from main.models import Dataset
from main.models import ExperimentSample
from main.models import Project
from main.models import ReferenceGenome
from main.models import Variant
from main.models import VariantAlternate
from main.models import VariantCallerCommonData
from main.models import VariantEvidence
from main.models import VariantSet
from main.models import VariantToVariantSet
from scripts.dynamic_snp_filter_key_map import initialize_filter_key_map
from scripts.dynamic_snp_filter_key_map import update_filter_key_map
from settings import PWD as GD_ROOT
from variants.variant_filter import EXPRESSION_REGEX
from variants.variant_filter import SAMPLE_SCOPE_REGEX
from variants.variant_filter import SAMPLE_SCOPE_REGEX_NAMED
from variants.variant_filter import SET_REGEX
from variants.variant_filter import SET_REGEX_NAMED
from variants.variant_filter import get_variants_that_pass_filter
from variants.variant_filter import SymbolGenerator
from variants.variant_filter import ParseError
from variants.variant_filter import VariantFilterEvaluator

TEST_DIR = os.path.join(GD_ROOT, 'test_data', 'genbank_aligned')

TEST_ANNOTATED_VCF = os.path.join(TEST_DIR, 'bwa_align_annotated.vcf')


class BaseTestVariantFilterTestCase(TestCase):
    """Base class for the test classes here.
    """

    def setUp(self):
        user = User.objects.create_user('testuser', password='password',
                email='<EMAIL>')
        self.project = Project.objects.create(owner=user.get_profile(),
                title='Test Project')
        self.ref_genome = ReferenceGenome.objects.create(project=self.project,
                label='refgenome', num_chromosomes=1, num_bases=1000)

        # Make sure the reference genome has the required vcf keys.
        initialize_filter_key_map(self.ref_genome)
        update_filter_key_map(self.ref_genome, TEST_ANNOTATED_VCF)

        self.vcf_dataset = Dataset.objects.create(
                label='test_data_set',
                type=Dataset.TYPE.VCF_FREEBAYES,
                filesystem_location=TEST_ANNOTATED_VCF)


class TestVariantFilter(BaseTestVariantFilterTestCase):

    def test_filter__by_position(self):
        """Test filtering by position.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        # Test querying Variants with position > 5.
        result = get_variants_that_pass_filter('position > 5', self.ref_genome)
        variants_above_5 = result.variant_set
        self.assertEqual(4, len(variants_above_5))
        for var in variants_above_5:
            self.assertTrue(var.position > 5)

        # Test querying Variants with position >= 5.
        result = get_variants_that_pass_filter('position >= 5', self.ref_genome)
        variants_above_5 = result.variant_set
        self.assertEqual(5, len(variants_above_5))
        for var in variants_above_5:
            self.assertTrue(var.position >= 5)


    def test_filter__by_chromosome(self):
        """Test filtering by position.
        """
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        CHROM_2 = 'chrom2'
        for pos in range(9):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_2,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        result = get_variants_that_pass_filter('chromosome = chrom',
                self.ref_genome)
        self.assertEqual(6, len(result.variant_set))

        result = get_variants_that_pass_filter('chromosome = chrom2',
                self.ref_genome)
        self.assertEqual(9, len(result.variant_set))


    def test_filter__by_position_and_chromosome(self):
        """Test filtering by position.
        """
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        CHROM_2 = 'chrom2'
        for pos in range(9):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_2,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        QUERY_STRING = 'position > 4 & chromosome = chrom'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        self.assertEqual(1, len(result.variant_set))

        QUERY_STRING = 'position >= 5 & chromosome = chrom2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        self.assertEqual(4, len(result.variant_set))


    def test_filter__invalid_key(self):
        CHROM_1 = 'chrom'
        for pos in range(6):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome=CHROM_1,
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))            

        QUERY_STRING = 'dinosaur > 4 & chromosome = chrom'
        with self.assertRaises(ParseError):
            get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)


    def test_filter__by_position_complex(self):
        """Test filtering by position.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        # Test AND case.
        QUERY_STRING = 'position < 1 & position > 7'
        result = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome)
        self.assertEqual(0, len(result.variant_set))

        # Test OR case.
        QUERY_STRING = 'position < 1 | position > 7'
        result = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome)
        variant_list = result.variant_set
        self.assertEqual(3, len(variant_list))
        for pos in [0, 8, 9]:
            found = False
            for var in variant_list:
                if var.position == pos:
                    found = True
                    break
            self.assertTrue(found, "Expected variant at pos %d not found" % pos)


    def test_filter__with_non_SQL_key(self):
        """Tests handling of non-SQL key.

        NOTE: This is not actually testing filtering, just showing that we
        can handle the key.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))

        QUERY_STRING = 'position < 1 & INFO_XRM > 0'
        variants = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome).variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position < 1 | position >= 7 & INFO_XRM > 0'
        variants = get_variants_that_pass_filter(QUERY_STRING,
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))


    def test_filter__common_data(self):
        """Test with VariantCallerCommonData filtering.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        raw_data_dict = {
                'INFO_XRM': pickle.dumps(0.12),
        }
        VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_data_dict
        )

        QUERY_STRING = 'position < 5 & INFO_XRM > 0'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & INFO_XRM > 1'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))


    def test_filter__variant_evidence(self):
        """Test with VariantEvidence objects.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        raw_common_data_dict = {
                'INFO_XRM': pickle.dumps(0.12),
        }
        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_common_data_dict
        )

        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(0),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(None),
                'gt_bases': pickle.dumps(None)

        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        QUERY_STRING = 'position < 5 & gt_type = 2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & gt_type = 0'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = 'position < 5 & gt_type = 1'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position > 5 & gt_type = 2'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = 'position < 5 & (gt_type = 0 | gt_type = 2)'
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        # Check that the number of passing samples is correct.
        passing_sample_ids = result.variant_id_to_metadata_dict[variant.id][
                'passing_sample_ids']
        self.assertEqual(2, len(passing_sample_ids))


    def test_filter__equality(self):
        """Test filtering with equality operators.
        """
        # Create several Variants with positions:
        # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        for pos in range(10):
            var = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=pos,
                ref_value='A')

            var.variantalternate_set.add(
                    VariantAlternate.objects.create(
                            variant=var,
                            alt_value='G'))


        variants = get_variants_that_pass_filter('position == 5',
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))
        self.assertEqual(5, variants.pop().position)

        variants = get_variants_that_pass_filter('position = 5',
                self.ref_genome).variant_set
        self.assertEqual(1, len(variants))
        self.assertEqual(5, variants.pop().position)

        variants = get_variants_that_pass_filter('position != 5',
                self.ref_genome).variant_set
        self.assertEqual(9, len(variants))
        for var in variants:
            self.assertTrue(var.position != 5)


    def test_filter__scoped(self):
        """Test with VariantEvidence objects.
        """
        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        variant.variantalternate_set.add(
                VariantAlternate.objects.create(
                        variant=variant,
                        alt_value='G'))

        # Common data shared by all.
        raw_common_data_dict = {
                'INFO_XRM': pickle.dumps(0.12)
        }
        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=raw_common_data_dict
        )

        # Sample 1
        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        # Sample 2
        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(0),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        # Sample 3
        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(1),
                'gt_bases': pickle.dumps(None)
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        # Sample 4
        sample_obj_4 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 4',
                group='Plate 1',
                well='A04',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_4,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        QUERY_STRING = '(position < 5 & gt_type = 2) in ANY(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ALL(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ALL(%s)' % (
                sample_obj.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(1, len(variants))

        QUERY_STRING = '(position < 5 & gt_type = 2) in ONLY(%s, %s, %s)' % (
                sample_obj.uid, sample_obj_2.uid, sample_obj_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))

        # This should fail because condition passes for sample_obj and
        # sample_obj_4.
        QUERY_STRING = '(position < 5 & gt_type = 2) in ONLY(%s)' % (
                sample_obj.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(0, len(variants))


    def test_filter__common_data_per_alt(self):
        """Test filtering for common data of type '-1', 
        with different values for each alternate allele. 
        """

        variant = Variant.objects.create(
                type=Variant.TYPE.TRANSITION,
                reference_genome=self.ref_genome,
                chromosome='chrom',
                position=2,
                ref_value='A')

        for alt_value, alt_eff in zip(['G','T'], ['NON_SYNONYMOUS_CODING',
                        'SYNONYMOUS_CODING']):

            alt = VariantAlternate.objects.create(
                        variant=variant,
                        alt_value=alt_value)
            alt.data = {'INFO_EFF_EFFECT': alt_eff}

            variant.variantalternate_set.add(alt)

        common_data_obj = VariantCallerCommonData.objects.create(
                variant=variant,
                source_dataset=self.vcf_dataset,
                data=dict()
        )
        sample_obj = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample',
                group='Plate 1',
                well='A01',
                num_reads=100,
        )
        raw_sample_data_dict = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'GT': pickle.dumps('1/1'),
                'gt_bases': pickle.dumps('G/G')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)

        sample_obj_2 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 2',
                group='Plate 1',
                well='A02',
                num_reads=100,
        )
        raw_sample_data_dict_2 = {
                'called': True,
                'gt_type': pickle.dumps(2),
                'GT': pickle.dumps('1/2'),
                'gt_bases': pickle.dumps('G/T')
        }
        VariantEvidence.objects.create(
                uid=1000,
                experiment_sample=sample_obj_2,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict_2)

        sample_obj_3 = ExperimentSample.objects.create(
                project=self.project,
                label='fake sample 3',
                group='Plate 1',
                well='A03',
                num_reads=100
        )
        raw_sample_data_dict = {
                'called': False,
                'gt_type': pickle.dumps(1),
                'GT': pickle.dumps('0/0'),
                'gt_bases': pickle.dumps('A/A')
        }
        VariantEvidence.objects.create(
                experiment_sample=sample_obj_3,
                variant_caller_common_data=common_data_obj,
                data=raw_sample_data_dict)


        # per_alt_dict, per_alt_types = get_per_alt_dict(
        #         'INFO_EFF_EFFECT', variant,
        #         VariantEvidence.objects.get(uid=1000),
        #         self.ref_genome.variant_key_map['snp_caller_common_data'])

        # assert('INFO_EFF_EFFECT' in per_alt_dict.keys())
        # assert('NON_SYNONYMOUS_CODING' in per_alt_dict['INFO_EFF_EFFECT'])

        QUERY_STRING = (
                '(position < 5 & INFO_EFF_EFFECT = NON_SYNONYMOUS_CODING)')
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set

        self.assertEqual(1, len(variants))  

        passing_sample_ids = result.variant_id_to_metadata_dict[variant.id][
                'passing_sample_ids']

        self.assertEqual(2, len(passing_sample_ids))

    def test_filter__sets(self):
        """Test filtering relative to sets.
        """
        # We create 3 VariantSets.
        set_1 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set1')
        set_2 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set2')
        set_3 = VariantSet.objects.create(
                reference_genome=self.ref_genome,
                label='set3')

        # Add some variants to each of the above sets.
        def _create_variants_in_set(pos_range, var_set_list):
            for pos in pos_range:
                var = Variant.objects.create(
                    type=Variant.TYPE.TRANSITION,
                    reference_genome=self.ref_genome,
                    chromosome='chrom',
                    position=pos,
                    ref_value='A')

                var.variantalternate_set.add(
                        VariantAlternate.objects.create(
                                variant=var,
                                alt_value='G'))

                for var_set in var_set_list:
                    if var_set is not None:
                        VariantToVariantSet.objects.create(variant=var,
                                variant_set=var_set)
        _create_variants_in_set(range(3), [set_1]) # 3
        _create_variants_in_set(range(3, 6), [set_2]) # 3
        _create_variants_in_set(range(6, 9), [set_3]) # 3
        _create_variants_in_set(range(9, 13), [None]) # 4
        _create_variants_in_set(range(13, 19), [set_1, set_2]) # 6

        # We have 19 total variants.
        self.assertEqual(19, len(Variant.objects.all()))

        # 3 + 6 = 9 in set_1
        QUERY_STRING = 'IN_SET(%s)' % (set_1.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(9, len(variants))

        # 3 + 6 + 3 = 12 between set_1 and set_2
        QUERY_STRING = 'IN_SET(%s) | IN_SET(%s)' % (set_1.uid, set_2.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(12, len(variants))

        # All but 4, so 15, in any set.
        QUERY_STRING = 'IN_SET(%s) | IN_SET(%s) | IN_SET(%s)' % (
                set_1.uid, set_2.uid, set_3.uid)
        result = get_variants_that_pass_filter(QUERY_STRING, self.ref_genome)
        variants = result.variant_set
        self.assertEqual(15, len(variants))

        # All but 9, or 10 not in set_1.
{'fraction_non_alphanumeric': 0.08867677196610232, 'fraction_numerical': 0.012794712257445119, 'mean_word_length': 1.7867874048158057, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1457087', 'n_tokens_mistral': 16671, 'n_tokens_neox': 15118, 'n_words': 2647}