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, 'tests.testbeds.MockTestBedv0', 'MockTestBedv0', False, 'from tests.testbeds import MockTestBedv0\n'), (17, 'robustnessgym.slicebuilders.subpopulations.length.LengthSubpopulation', 'LengthSubpopulation', False, 'from robustnessgym.slicebuilders.subpopulations.length import LengthSubpopulation\n'), (13, 'robustnessgym.cachedops.spacy.Spacy', 'Spacy', False, 'from robustnessgym.cachedops.spacy import Spacy\n'), (24, 'robustnessgym.cachedops.spacy.Spacy.retrieve', 'Spacy.retrieve', False, 'from robustnessgym.cachedops.spacy import Spacy\n'), (21, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (30, 'numpy.array', 'np.array', True, 'import numpy as np\n')]