from covigator.database.model import VariantCooccurrence, DataSource
from covigator.pipeline.cooccurrence_matrix import CooccurrenceMatrix
from covigator.tests.unit_tests.abstract_test import AbstractTest
from covigator.tests.unit_tests.mocked import get_mocked_ena_sample, get_mocked_variant, get_mocked_variant_observation


class CooccurrenceMatrixTests(AbstractTest):

    NUM_UNIQUE_VARIANTS = 10
    NUM_SAMPLES = 10
    NUM_VARIANT_OBSERVATIONS_PER_SAMPLE = 5

    def setUp(self) -> None:
        # mocks some unique variants
        self.mocked_variants = [get_mocked_variant(self.faker) for _ in range(self.NUM_UNIQUE_VARIANTS)]
        self.session.add_all(self.mocked_variants)
        self.session.commit()

        variant_observations = []
        self.samples = []
        for _ in range(self.NUM_SAMPLES):
            sample_ena = get_mocked_ena_sample(self.faker)
            self.session.add(sample_ena)
            self.session.commit()
            self.samples.append(sample_ena)
            # mocks observed variants
            for vo in self.faker.random_elements(
                    self.mocked_variants, length=self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE, unique=True):
                variant_observations.append(get_mocked_variant_observation(sample_ena, vo))
            self.session.add_all(variant_observations)
            self.session.commit()

    def test_missing_sample(self):
        self.assertRaises(
            AssertionError,
            CooccurrenceMatrix().compute,
            run_accession=None,
            source=None,
            session=self.session)

    def test_missing_session(self):
        self.assertRaises(
            AssertionError,
            CooccurrenceMatrix().compute,
            run_accession="12345",
            source=DataSource.ENA,
            session=None)

    def test_non_existing_sample_does_not_add_new_entries(self):
        count = self.session.query(VariantCooccurrence).count()
        self.assertEqual(count, 0)
        CooccurrenceMatrix().compute(run_accession="12345", source=DataSource.ENA.name, session=self.session)
        self.session.commit()
        count = self.session.query(VariantCooccurrence).count()
        self.assertEqual(count, 0)

    def test_one_existing_sample(self):
        count = self.session.query(VariantCooccurrence).count()
        self.assertEqual(count, 0)
        CooccurrenceMatrix().compute(
            run_accession=self.samples[0].run_accession, source=DataSource.ENA.name, session=self.session)
        self.session.commit()
        count = self.session.query(VariantCooccurrence).count()
        # size of matrix = n*(n-1)/2 + n (one half of matrix + diagonal)
        self.assertEqual(
            count, (self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE * (self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE - 1) / 2) +
                   self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE)

        self._assert_cooccurrence_matrix(assert_cooccurrence=False)

    def test_all_samples(self):
        count = self.session.query(VariantCooccurrence).count()
        self.assertEqual(count, 0)
        for s in self.samples:
            CooccurrenceMatrix().compute(run_accession=s.run_accession, source=DataSource.ENA.name, session=self.session)
            self.session.commit()
        count = self.session.query(VariantCooccurrence).count()
        self.assertLess(
            count,
            self.NUM_SAMPLES * self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE *
            (self.NUM_VARIANT_OBSERVATIONS_PER_SAMPLE - 1) / 2)
        self._assert_cooccurrence_matrix()

    def _assert_cooccurrence_matrix(self, assert_cooccurrence=True):
        variant_cooccurrences = self.session.query(VariantCooccurrence).all()
        found_cooccurrent_variant = False
        for vo in variant_cooccurrences:
            position_one = int(vo.variant_id_one.split(":")[0])
            position_two = int(vo.variant_id_two.split(":")[0])
            self.assertLessEqual(position_one, position_two)
            self.assertLessEqual(vo.count, self.NUM_SAMPLES)
            self.assertGreater(vo.count, 0)
            if vo.count > 1:
                found_cooccurrent_variant = True
        if assert_cooccurrence:
            self.assertTrue(found_cooccurrent_variant)

