import unittest
from scripts.ner_utils import convert_bsf, parse_bsf, BsfInfo


class TestBsf2Iob(unittest.TestCase):

    def test_1line_follow_markup_iob(self):
        data = 'тележурналіст Василь .'
        bsf_markup = 'T1	PERS 14 20	Василь'
        expected = '''тележурналіст O
Василь B-PERS
. O'''
        self.assertEqual(expected, convert_bsf(data, bsf_markup, converter='iob'))

    def test_1line_2tok_markup_iob(self):
        data = 'тележурналіст Василь Нагірний .'
        bsf_markup = 'T1	PERS 14 29	Василь Нагірний'
        expected = '''тележурналіст O
Василь B-PERS
Нагірний I-PERS
. O'''
        self.assertEqual(expected, convert_bsf(data, bsf_markup, converter='iob'))

    def test_1line_Long_tok_markup_iob(self):
        data = 'А в музеї Гуцульщини і Покуття можна '
        bsf_markup = 'T12	ORG 4 30	музеї Гуцульщини і Покуття'
        expected = '''Рђ O
РІ O
музеї B-ORG
Гуцульщини I-ORG
С– I-ORG
Покуття I-ORG
можна O'''
        self.assertEqual(expected, convert_bsf(data, bsf_markup, converter='iob'))

    def test_2line_2tok_markup_iob(self):
        data = '''тележурналіст Василь Нагірний .
В івано-франківському видавництві «Лілея НВ» вийшла друком'''
        bsf_markup = '''T1	PERS 14 29	Василь Нагірний
T2	ORG 67 75	Лілея НВ'''
        expected = '''тележурналіст O
Василь B-PERS
Нагірний I-PERS
. O


Р’ O
івано-франківському O
видавництві O
В« O
Лілея B-ORG
РќР’ I-ORG
В» O
вийшла O
РґСЂСѓРєРѕРј O'''
        self.assertEqual(expected, convert_bsf(data, bsf_markup, converter='iob'))

    def test_all_multiline_iob(self):
        data = '''його книжечка «А .
Kubler .
Світло і тіні маестро» .
Причому'''
        bsf_markup = '''T4	MISC 15 49	Рђ .
Kubler .
Світло і тіні маестро
'''
        expected = '''Р№РѕРіРѕ O
книжечка O
В« O
Рђ B-MISC
. I-MISC
Kubler I-MISC
. I-MISC
Світло I-MISC
С– I-MISC
тіні I-MISC
маестро I-MISC
В» O
. O


Причому O'''
        self.assertEqual(expected, convert_bsf(data, bsf_markup, converter='iob'))


if __name__ == '__main__':
    unittest.main()