import os, unittest

from pytlex_core import TimeMLParser
from pytlex_tests.jTLEX_Data import jTLEX_Counts

problematic_bunch = {'ABC19980304.1830.1636.tml', 'APW19980213.1320.tml', 'APW19980227.0468.tml',
                     'APW19980227.0489.tml'}  # These files have extra events that is not used for instances


class TimeMLParserTestCase(unittest.TestCase):
    def test_instance_count(self):
        filepath = r"../pytlex_data/TimeBankCorpus"
        for subdir, dirs, files in os.walk(filepath):
            for filename in files:
                with open(filepath + "/" + filename) as timeMLFile:
                    instance_list = TimeMLParser.parse_instances(timeMLFile.read())
                    self.assertEqual(len(instance_list), jTLEX_Counts.jTLEX_instance_count[filename])

    def test_timex_count(self):
        filepath = r"../pytlex_data/TimeBankCorpus"
        for subdir, dirs, files in os.walk(filepath):
            for filename in files:
                with open(filepath + "/" + filename) as timeMLFile:
                    timex_list = TimeMLParser.parse_timex(timeMLFile.read())
                    self.assertEqual(len(timex_list), jTLEX_Counts.jTLEX_timex_count[filename])

    def test_link_count(self):
        filepath = r"../pytlex_data/TimeBankCorpus"
        for subdir, dirs, files in os.walk(filepath):
            for filename in files:
                with open(filepath + "/" + filename) as timeMLFile:
                    link_list = TimeMLParser.parse_links(timeMLFile.read())
                    self.assertEqual(len(link_list), jTLEX_Counts.jTLEX_link_count[filename])


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