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', ([], {'intervals': '[(1, 3), (4, 5)]'}), 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', (['self.testbed.dataset[:]', "['text']"], {}), False, 'from robustnessgym.cachedops.spacy import Spacy\n'), (21, 'numpy.array', 'np.array', (['[5, 5, 5, 5, 5, 5]'], {}), True, 'import numpy as np\n'), (30, 'numpy.array', 'np.array', (['([[0, 1]] * 6)'], {}), True, 'import numpy as np\n')]