from unittest import TestCase

import numpy as np

from robustnessgym.cachedops.spacy import Spacy
from robustnessgym.slicebuilders.subpopulations.length import LengthSubpopulation
from tests.testbeds import MockTestBedv0


class TestLengthSubpopulation(TestCase):
    def setUp(self):
        self.testbed = MockTestBedv0()
        self.testbed.dataset = Spacy()(self.testbed.dataset, columns=["text"])

    def test_score(self):
        # Create the length subpopulation
        length = LengthSubpopulation(intervals=[(1, 3), (4, 5)])

        # Compute scores
        scores = length.score(self.testbed.dataset[:], columns=["text"])
        self.assertTrue(np.allclose(scores, np.array([5, 5, 5, 5, 5, 5])))

        print(self.testbed.dataset.column_names)
        print(Spacy.retrieve(self.testbed.dataset[:], ["text"]))

        # Apply the subpopulation
        slices, slice_matrix = length(self.testbed.dataset, columns=["text"])

        # Check that the slice membership lines up
        self.assertTrue(np.allclose(slice_matrix, np.array([[0, 1]] * 6)))

[((12, 23, 12, 38), 'tests.testbeds.MockTestBedv0', 'MockTestBedv0', ({}, {}), '()', False, 'from tests.testbeds import MockTestBedv0\n'), ((17, 17, 17, 64), 'robustnessgym.slicebuilders.subpopulations.length.LengthSubpopulation', 'LengthSubpopulation', (), '', False, 'from robustnessgym.slicebuilders.subpopulations.length import LengthSubpopulation\n'), ((13, 31, 13, 38), 'robustnessgym.cachedops.spacy.Spacy', 'Spacy', ({}, {}), '()', False, 'from robustnessgym.cachedops.spacy import Spacy\n'), ((24, 14, 24, 63), 'robustnessgym.cachedops.spacy.Spacy.retrieve', 'Spacy.retrieve', ({(24, 29, 24, 52): 'self.testbed.dataset[:]', (24, 54, 24, 62): "['text']"}, {}), "(self.testbed.dataset[:], ['text'])", False, 'from robustnessgym.cachedops.spacy import Spacy\n'), ((21, 44, 21, 72), 'numpy.array', 'np.array', ({(21, 53, 21, 71): '[5, 5, 5, 5, 5, 5]'}, {}), '([5, 5, 5, 5, 5, 5])', True, 'import numpy as np\n'), ((30, 50, 30, 72), 'numpy.array', 'np.array', ({(30, 59, 30, 71): '([[0, 1]] * 6)'}, {}), '([[0, 1]] * 6)', True, 'import numpy as np\n')]