# -*- coding: utf-8 -*-

from scout.export.variant import export_mt_variants, export_verified_variants
from scout.constants.variants_export import MT_EXPORT_HEADER, VERIFIED_VARIANTS_HEADER

def test_export_mt_variants(case_obj, real_populated_database):
    """Test the function that creates lines with MT variants to be exported"""
    adapter = real_populated_database
    case_id = case_obj['_id']
    assert case_id

    # load MT variants for this case
    nr_loaded = adapter.load_variants(case_obj=case_obj,
                category='snv', chrom='MT', start=1, end=16500)
    assert nr_loaded > 0
    mt_variants = list(adapter.variant_collection.find({'chromosome':'MT'}))
    assert len(mt_variants) == nr_loaded # it's all MT variants, but double-checking it

    # Assert that there is at least one sample to create the excel file for
    samples = case_obj.get('individuals')
    assert samples

    # test function that exports variant lines
    for sample in samples:
        sample_lines = export_mt_variants(variants=mt_variants, sample_id=sample['individual_id'])

        # check that rows to write to excel corespond to number of variants
        assert len(sample_lines) == len(mt_variants)
        # check that cols to write to excel corespond to fields of excel header
        assert len(sample_lines[0]) == len(MT_EXPORT_HEADER)


def test_export_verified_variants(case_obj, real_populated_database, variant_objs):
    """Test the function that creates lines with verified variants to be exported"""

    adapter = real_populated_database
    case_id = case_obj['_id']

    # load variants to populated_database
    assert adapter.variants(case_id=case_id, nr_of_variants=-1).count() == 0
    nr_loaded = adapter.load_variants(case_obj=case_obj)
    assert nr_loaded > 0

    valid_status = ['False positive', 'True positive', 'False positive'] # for 3 vars

    test_vars = list(adapter.variant_collection.find().limit(3))
    assert len(test_vars) == 3

    # Make sure that no variant is set as validated:
    assert adapter.variant_collection.find({'validation':{'$exists':True}}).count()==0

    # Set test variant as validated
    for i in range(3):
        # Set validation status of a variant
        adapter.variant_collection.find_one_and_update(
            {'_id' : test_vars[i]['_id'] },
            {'$set' : {'validation' : valid_status[i] } }
         )

        # insert validation events
        adapter.event_collection.insert_one({
            'verb' : 'validate',
            'institute' : test_vars[i]['institute'],
            'variant_id' : test_vars[i]['variant_id'],
            'case' : case_id
         })

    # There should be 3 validated variants now
    assert adapter.variant_collection.find({'validation':{'$exists':True}}).count()==3

    # Call function to get aggregated data (variant + case data):
    cust = case_obj['owner']
    aggregated_vars = adapter.verified(cust)
    assert len(aggregated_vars) == 3 # same number of variants is returned

    # Call function that creates document lines
    document_lines = export_verified_variants(aggregated_vars)
    assert len(document_lines) == 3 # one line per variant should be returned

    for line in document_lines:
        # Make sure that document cells will be the same as in document header
        assert len(line) == len(VERIFIED_VARIANTS_HEADER)
